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

@@ -46,6 +46,11 @@ type workflowEngine struct {
 	// backendClient makes calls to backend MCP servers.
 	backendClient vmcp.BackendClient
 
+	// tools is the resolved tool list for the session. When non-nil it is
+	// used by getToolInputSchema instead of reading DiscoveredCapabilities
+	// from the request context, making the engine independent of middleware.
+	tools []vmcp.Tool
+
 	// templateExpander handles template expansion.
 	templateExpander TemplateExpander
 
@@ -74,14 +79,17 @@ type workflowEngine struct {
 // will not be available. Use NewInMemoryStateStore() for basic state tracking.
 //
 // The auditor parameter is optional. If nil, workflow execution will not be audited.
+//
+// Optional behaviour can be configured via EngineOption values (see WithTools).
 func NewWorkflowEngine(
 	rtr router.Router,
 	backendClient vmcp.BackendClient,
 	elicitationHandler ElicitationProtocolHandler,
 	stateStore WorkflowStateStore,
 	auditor *audit.WorkflowAuditor,
+	opts ...EngineOption,
 ) Composer {
-	return &workflowEngine{
+	e := &workflowEngine{
 		router:             rtr,
 		backendClient:      backendClient,
 		templateExpander:   NewTemplateExpander(),
@@ -91,6 +99,21 @@ func NewWorkflowEngine(
 		stateStore:         stateStore,
 		auditor:            auditor,
 	}
+	for _, o := range opts {
+		o(e)
+	}
+	return e
+}
+
+// EngineOption configures a workflowEngine.
+type EngineOption func(*workflowEngine)
+
+// WithTools binds a resolved tool list to the engine. When set, getToolInputSchema
+// uses this list instead of reading DiscoveredCapabilities from the request context,
+// making the engine independent of the discovery middleware. Use this when creating
+// per-session engines via router.NewSessionRouter.
+func WithTools(tools []vmcp.Tool) EngineOption {
+	return func(e *workflowEngine) { e.tools = tools }
 }
 
 // ExecuteWorkflow executes a composite tool workflow.
@@ -1223,20 +1246,31 @@ 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 {
+// getToolInputSchema looks up a tool's InputSchema.
+// When the engine was created with WithTools, it resolves from that list directly.
+// Otherwise it falls back to reading DiscoveredCapabilities from the request context
+// (legacy path used in tests).
+// Returns nil if the tool is not found.
+func (e *workflowEngine) getToolInputSchema(ctx context.Context, toolName string) map[string]any {
+	// Prefer session-bound tools list when available (no context dependency).
+	if e.tools != nil {
+		for i := range e.tools {
+			if e.tools[i].Name == toolName {
+				return e.tools[i].InputSchema
+			}
+		}
+		return nil
+	}
+
+	// Fallback: read from context (engines created via NewWorkflowEngine).
 	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
 		}
 	}
-
 	return nil
 }

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
+}