Bring composite tools into session abstraction

Composite tool workflow engines were previously relying on the
discovery middleware to inject DiscoveredCapabilities into the
request context so that the shared stateless router could route
backend tool calls within workflows. This created an implicit
coupling between the middleware and composite tool execution that
made unit-testing harder and was a source of integration bugs.

Affected components: pkg/vmcp/router, pkg/vmcp/composer,
pkg/vmcp/server, pkg/vmcp/discovery

Related-to: #3872

Author: taskbot

pkg/vmcp/composer/workflow_engine.go

@@ -17,7 +17,6 @@ import (
 	"github.com/stacklok/toolhive/pkg/audit"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"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/schema"
 )
@@ -46,6 +45,10 @@ type workflowEngine struct {
 	// backendClient makes calls to backend MCP servers.
 	backendClient vmcp.BackendClient
 
+	// tools is the resolved tool list for the session, used by getToolInputSchema
+	// for argument type coercion. Set via NewSessionWorkflowEngine.
+	tools []vmcp.Tool
+
 	// templateExpander handles template expansion.
 	templateExpander TemplateExpander
 
@@ -93,6 +96,30 @@ func NewWorkflowEngine(
 	}
 }
 
+// NewSessionWorkflowEngine creates a per-session workflow engine bound to a resolved tool list.
+// tools is required: it enables argument type coercion against the session's tool schemas.
+// Use this when creating per-session engines via router.NewSessionRouter.
+func NewSessionWorkflowEngine(
+	rtr router.Router,
+	backendClient vmcp.BackendClient,
+	elicitationHandler ElicitationProtocolHandler,
+	stateStore WorkflowStateStore,
+	auditor *audit.WorkflowAuditor,
+	tools []vmcp.Tool,
+) Composer {
+	return &workflowEngine{
+		router:             rtr,
+		backendClient:      backendClient,
+		templateExpander:   NewTemplateExpander(),
+		contextManager:     newWorkflowContextManager(),
+		elicitationHandler: elicitationHandler,
+		dagExecutor:        newDAGExecutor(defaultMaxParallelSteps),
+		stateStore:         stateStore,
+		auditor:            auditor,
+		tools:              tools,
+	}
+}
+
 // ExecuteWorkflow executes a composite tool workflow.
 //
 // TODO(rate-limiting): Add rate limiting per user/session to prevent workflow execution DoS.
@@ -407,7 +434,7 @@ func (e *workflowEngine) executeToolStep(
 	// Coerce expanded arguments to expected types based on backend tool schema.
 	// Template expansion returns strings, but backend tools expect typed values
 	// (integer, boolean, number) as defined in their InputSchema.
-	rawSchema := e.getToolInputSchema(ctx, step.Tool)
+	rawSchema := e.getToolInputSchema(step.Tool)
 	s := schema.MakeSchema(rawSchema)
 	if coerced, ok := s.TryCoerce(expandedArgs).(map[string]any); ok {
 		expandedArgs = coerced
@@ -1223,20 +1250,13 @@ func (e *workflowEngine) auditStepSkipped(
 	}
 }
 
-// getToolInputSchema looks up a tool's InputSchema from discovered capabilities.
-// Returns nil if the tool is not found or capabilities are not in context.
-func (*workflowEngine) getToolInputSchema(ctx context.Context, toolName string) map[string]any {
-	caps, ok := discovery.DiscoveredCapabilitiesFromContext(ctx)
-	if !ok || caps == nil {
-		return nil
-	}
-
-	// Search in backend tools
-	for i := range caps.Tools {
-		if caps.Tools[i].Name == toolName {
-			return caps.Tools[i].InputSchema
+// getToolInputSchema looks up a tool's InputSchema from the session-bound tools list.
+// Returns nil if the engine has no tools list or the tool is not found.
+func (e *workflowEngine) getToolInputSchema(toolName string) map[string]any {
+	for i := range e.tools {
+		if e.tools[i].Name == toolName {
+			return e.tools[i].InputSchema
 		}
 	}
-
 	return nil
 }

pkg/vmcp/composer/workflow_engine_test.go

@@ -741,3 +741,98 @@ func TestWorkflowEngine_WorkflowMetadataAvailableInTemplates(t *testing.T) {
 	assert.Equal(t, WorkflowStatusCompleted, result.Status)
 	assert.Len(t, result.Steps, 2)
 }
+
+func TestWorkflowEngine_SessionEngine_CoercesTemplateStringToTypedArg(t *testing.T) {
+	t.Parallel()
+
+	// Template expansion always produces strings. When the engine is created
+	// with NewSessionWorkflowEngine, getToolInputSchema resolves the target tool's InputSchema
+	// and the schema coercion layer converts "42" → 42 before calling the backend.
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+
+	mockRouter := routermocks.NewMockRouter(ctrl)
+	mockBackend := mocks.NewMockBackendClient(ctrl)
+
+	tools := []vmcp.Tool{
+		{
+			Name: "count_items",
+			InputSchema: map[string]any{
+				"type": "object",
+				"properties": map[string]any{
+					"limit": map[string]any{"type": "integer"},
+				},
+			},
+		},
+	}
+
+	engine := NewSessionWorkflowEngine(mockRouter, mockBackend, nil, nil, nil, tools)
+
+	target := &vmcp.BackendTarget{WorkloadID: "backend1", BaseURL: "http://backend1:8080"}
+	mockRouter.EXPECT().RouteTool(gomock.Any(), "count_items").Return(target, nil)
+
+	// Expect the backend to receive the coerced integer, not the string "42".
+	coercedArgs := map[string]any{"limit": int64(42)}
+	mockBackend.EXPECT().
+		CallTool(gomock.Any(), target, "count_items", coercedArgs, gomock.Any()).
+		Return(&vmcp.ToolCallResult{StructuredContent: map[string]any{"items": []any{}}, Content: []vmcp.Content{}}, nil)
+
+	workflow := &WorkflowDefinition{
+		Name: "coerce_test",
+		Parameters: map[string]any{
+			"type": "object",
+			"properties": map[string]any{
+				"n": map[string]any{"type": "string"},
+			},
+		},
+		Steps: []WorkflowStep{
+			{
+				ID:   "step1",
+				Type: StepTypeTool,
+				Tool: "count_items",
+				// Template expansion produces a string; coercion must convert it to int.
+				Arguments: map[string]any{"limit": "{{.params.n}}"},
+			},
+		},
+	}
+
+	result, err := engine.ExecuteWorkflow(context.Background(), workflow, map[string]any{"n": "42"})
+	require.NoError(t, err)
+	assert.Equal(t, WorkflowStatusCompleted, result.Status)
+}
+
+func TestWorkflowEngine_SessionEngine_ToolNotInList_ReturnsNilSchema(t *testing.T) {
+	t.Parallel()
+
+	// When NewSessionWorkflowEngine is used but the requested tool is not in the list,
+	// getToolInputSchema returns nil and coercion is a no-op.
+	ctrl := gomock.NewController(t)
+	t.Cleanup(ctrl.Finish)
+
+	mockRouter := routermocks.NewMockRouter(ctrl)
+	mockBackend := mocks.NewMockBackendClient(ctrl)
+
+	// Tools list does not include "other_tool".
+	tools := []vmcp.Tool{{Name: "known_tool", InputSchema: map[string]any{"type": "object"}}}
+	engine := NewSessionWorkflowEngine(mockRouter, mockBackend, nil, nil, nil, tools)
+
+	target := &vmcp.BackendTarget{WorkloadID: "backend1", BaseURL: "http://backend1:8080"}
+	mockRouter.EXPECT().RouteTool(gomock.Any(), "other_tool").Return(target, nil)
+
+	// Args pass through unmodified (string stays a string).
+	rawArgs := map[string]any{"value": "hello"}
+	mockBackend.EXPECT().
+		CallTool(gomock.Any(), target, "other_tool", rawArgs, gomock.Any()).
+		Return(&vmcp.ToolCallResult{StructuredContent: map[string]any{"ok": true}, Content: []vmcp.Content{}}, nil)
+
+	workflow := &WorkflowDefinition{
+		Name: "no_schema_test",
+		Steps: []WorkflowStep{
+			{ID: "s1", Type: StepTypeTool, Tool: "other_tool", Arguments: rawArgs},
+		},
+	}
+
+	result, err := engine.ExecuteWorkflow(context.Background(), workflow, nil)
+	require.NoError(t, err)
+	assert.Equal(t, WorkflowStatusCompleted, result.Status)
+}

pkg/vmcp/discovery/middleware.go

@@ -281,10 +281,12 @@ func handleSubsequentRequest(
 		return ctx, fmt.Errorf("session not found: %s", sessionID)
 	}
 
-	// Backend tool calls are routed by session-scoped handlers registered with the SDK.
-	// However, composite tool workflow steps go through the shared router which requires
-	// DiscoveredCapabilities in the context. Inject capabilities built from the session's
-	// routing table so composite workflows can route backend tool calls correctly.
+	// Backend tool handlers (created by DefaultHandlerFactory) resolve their backend
+	// target by calling router.RouteTool(ctx, name), which reads DiscoveredCapabilities
+	// from the request context. Inject capabilities built from the session's routing
+	// table so these handlers can route correctly on subsequent requests.
+	// Note: composite tool workflow engines are created per-session and route via
+	// SessionRouter directly, so they no longer depend on this context value.
 	multiSess, isMulti := rawSess.(vmcpsession.MultiSession)
 	if !isMulti {
 		// The session is still a StreamableSession placeholder — Phase 2

pkg/vmcp/router/session_router.go

@@ -0,0 +1,66 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package router
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/stacklok/toolhive/pkg/vmcp"
+)
+
+// sessionRouter is a Router implementation backed directly by a RoutingTable,
+// requiring no request context to resolve capabilities. It is used by
+// per-session workflow engines so that composite tool execution does not depend
+// on the discovery middleware injecting DiscoveredCapabilities into the context.
+type sessionRouter struct {
+	routingTable *vmcp.RoutingTable
+}
+
+// NewSessionRouter creates a Router that routes from the provided RoutingTable
+// without reading the request context. This is the preferred router for
+// composite tool workflow engines because it couples routing to the session
+// rather than to middleware-managed context values.
+func NewSessionRouter(rt *vmcp.RoutingTable) Router {
+	return &sessionRouter{routingTable: rt}
+}
+
+// RouteTool resolves a tool name to its backend target using the session's
+// routing table directly.
+func (r *sessionRouter) RouteTool(_ context.Context, toolName string) (*vmcp.BackendTarget, error) {
+	if r.routingTable == nil || r.routingTable.Tools == nil {
+		return nil, fmt.Errorf("%w: %s", ErrToolNotFound, toolName)
+	}
+	target, exists := r.routingTable.Tools[toolName]
+	if !exists {
+		return nil, fmt.Errorf("%w: %s", ErrToolNotFound, toolName)
+	}
+	return target, nil
+}
+
+// RouteResource resolves a resource URI to its backend target using the
+// session's routing table directly.
+func (r *sessionRouter) RouteResource(_ context.Context, uri string) (*vmcp.BackendTarget, error) {
+	if r.routingTable == nil || r.routingTable.Resources == nil {
+		return nil, fmt.Errorf("%w: %s", ErrResourceNotFound, uri)
+	}
+	target, exists := r.routingTable.Resources[uri]
+	if !exists {
+		return nil, fmt.Errorf("%w: %s", ErrResourceNotFound, uri)
+	}
+	return target, nil
+}
+
+// RoutePrompt resolves a prompt name to its backend target using the session's
+// routing table directly.
+func (r *sessionRouter) RoutePrompt(_ context.Context, name string) (*vmcp.BackendTarget, error) {
+	if r.routingTable == nil || r.routingTable.Prompts == nil {
+		return nil, fmt.Errorf("%w: %s", ErrPromptNotFound, name)
+	}
+	target, exists := r.routingTable.Prompts[name]
+	if !exists {
+		return nil, fmt.Errorf("%w: %s", ErrPromptNotFound, name)
+	}
+	return target, nil
+}