Skip to content

Remove legacy optimizer decorator path #6119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kocaemre wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Remove legacy optimizer decorator path #6119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
011003b
Remove legacy optimizer decorator path
kocaemre Jul 28, 2026
b908e1e
Update session manager comments
kocaemre Aug 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
147 changes: 8 additions & 139 deletions pkg/vmcp/server/sessionmanager/factory.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@ package sessionmanager

import (
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"time"
Expand All @@ -18,15 +16,10 @@ import (

"github.com/stacklok/toolhive-core/mcpcompat/mcp"
mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server"
"github.com/stacklok/toolhive/pkg/auth"
"github.com/stacklok/toolhive/pkg/telemetry"
"github.com/stacklok/toolhive/pkg/vmcp"
"github.com/stacklok/toolhive/pkg/vmcp/conversion"
"github.com/stacklok/toolhive/pkg/vmcp/health"
"github.com/stacklok/toolhive/pkg/vmcp/optimizer"
vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session"
"github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec"
sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types"
)

const instrumentationName = "github.com/stacklok/toolhive/pkg/vmcp"
Expand All @@ -36,10 +29,9 @@ const instrumentationName = "github.com/stacklok/toolhive/pkg/vmcp"
// CacheCapacity from a config does not silently enable unbounded growth.
const defaultCacheCapacity = 1000

// FactoryConfig holds the session factory construction parameters that the
// session manager needs to build its decorating factory. It is separate from
// server.Config to avoid a circular import between the server and sessionmanager
// packages.
// FactoryConfig holds the session manager construction parameters. It is
// separate from server.Config to avoid a circular import between the server and
// sessionmanager packages.
type FactoryConfig struct {
// Base is the underlying session factory. Required.
Base vmcpsession.MultiSessionFactory
Expand Down Expand Up @@ -75,14 +67,14 @@ type FactoryConfig struct {
// tools — the single writer of the shared FTS5 store (the AC6
// no-double-index guarantee).
// - false: fine without an optimizer; with one, New rejects the config at
// construction, because the Serve layer discards the decorator's
// per-session tools (the optimizer would index the store yet serve
// nobody) and the Modern capability gate would fail open.
// construction, because the Serve layer is the only supported
// place to build an optimizer over the core-advertised tool set;
// building one here would index a store that no request path uses.
// New resolves the optimizer factory and owns its store/cleanup. server.New
// sets this unconditionally (server.go), so every in-tree composition
// advertises from the core; the flag exists for direct-Serve embedders. The
// decorator branch the false case used to select is now unreachable — its
// deletion is tracked in #6103.
// legacy decorator branch the false case used to select has been removed;
// optimizers are now built only by the Serve layer.
AdvertiseFromCore bool

// BackendHealth gates which backends a NEW session attempts to connect to.
Expand Down Expand Up @@ -152,129 +144,6 @@ func resolveOptimizer(cfg *FactoryConfig) (
}
}

// buildDecoratingFactory builds the decorating session factory from cfg.
// terminateSession is the session manager's own Terminate method, captured
// here to avoid the forward-reference dance previously needed in server.New().
func buildDecoratingFactory(
cfg *FactoryConfig,
optimizerFactory func(context.Context, []mcpserver.ServerTool) (optimizer.Optimizer, error),
terminateSession func(string) (bool, error),
) vmcpsession.MultiSessionFactory {
var decorators []vmcpsession.Decorator

// On the Serve path (AdvertiseFromCore) the optimizer is built by the Serve layer
// over the core's advertised set, so the factory's optimizer decorator is skipped
// to avoid double-indexing the shared store (see FactoryConfig.AdvertiseFromCore).
// Composite tools and their telemetry are owned by the core, not the factory.
// This branch is unreachable: New rejects an optimizer without AdvertiseFromCore,
// so optimizerFactory is nil whenever the flag is false. Deleting the decorator
// path is tracked in #6103.
if optimizerFactory != nil && !cfg.AdvertiseFromCore {
decorators = append(decorators, optimizerDecoratorFn(optimizerFactory, terminateSession))
}

return vmcpsession.NewDecoratingFactory(cfg.Base, decorators...)
}

// optimizerDecoratorFn returns a Decorator that indexes all session tools into
// the optimizer and replaces the tool list with find_tool + call_tool.
func optimizerDecoratorFn(
optimizerFactory func(context.Context, []mcpserver.ServerTool) (optimizer.Optimizer, error),
terminateSession func(string) (bool, error),
) vmcpsession.Decorator {
return func(ctx context.Context, sess vmcpsession.MultiSession) (vmcpsession.MultiSession, error) {
sdkTools, err := adaptToolsForFactory(sess, terminateSession)
if err != nil {
return nil, fmt.Errorf("failed to adapt tools for optimizer: %w", err)
}

opt, err := optimizerFactory(ctx, sdkTools)
if err != nil {
return nil, fmt.Errorf("failed to create optimizer: %w", err)
}

slog.Info("session capabilities decorated (optimizer mode)", "indexed_tool_count", len(sdkTools))
return optimizerdec.NewDecorator(sess, opt), nil
}
}

// adaptToolsForFactory converts domain tools from sess to SDK-format ServerTools.
// Unlike GetAdaptedTools in session_manager.go, this version accepts an explicit
// terminateSession callback so that auth failures still terminate the session,
// preserving hijack-prevention parity with the non-optimizer tool path.
func adaptToolsForFactory(
sess sessiontypes.MultiSession,
terminateSession func(string) (bool, error),
) ([]mcpserver.ServerTool, error) {
domainTools := sess.Tools()
sdkTools := make([]mcpserver.ServerTool, 0, len(domainTools))

for _, domainTool := range domainTools {
schemaJSON, err := json.Marshal(domainTool.InputSchema)
if err != nil {
return nil, fmt.Errorf("failed to marshal schema for tool %s: %w", domainTool.Name, err)
}

tool := mcp.Tool{
Name: domainTool.Name,
Description: domainTool.Description,
RawInputSchema: schemaJSON,
Annotations: conversion.ToMCPToolAnnotations(domainTool.Annotations),
}
if domainTool.OutputSchema != nil {
outputSchemaJSON, marshalErr := json.Marshal(domainTool.OutputSchema)
if marshalErr != nil {
slog.Warn("failed to marshal tool output schema", "tool", domainTool.Name, "error", marshalErr)
} else {
tool.RawOutputSchema = outputSchemaJSON
}
}

capturedSess := sess
capturedSessionID := sess.ID()
capturedToolName := domainTool.Name
handler := func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
args, ok := req.Params.Arguments.(map[string]any)
if !ok {
wrappedErr := fmt.Errorf("%w: arguments must be object, got %T", vmcp.ErrInvalidInput, req.Params.Arguments)
slog.Warn("invalid arguments for tool", "tool", capturedToolName, "error", wrappedErr)
return mcp.NewToolResultError(wrappedErr.Error()), nil
}

meta := conversion.FromMCPMeta(req.Params.Meta)
caller, _ := auth.IdentityFromContext(ctx)

result, callErr := capturedSess.CallTool(ctx, caller, capturedToolName, args, meta)
if callErr != nil {
if errors.Is(callErr, sessiontypes.ErrUnauthorizedCaller) || errors.Is(callErr, sessiontypes.ErrNilCaller) {
slog.Warn("caller authorization failed, terminating session",
"session_id", capturedSessionID, "tool", capturedToolName, "error", callErr)
if _, termErr := terminateSession(capturedSessionID); termErr != nil {
slog.Error("failed to terminate session after auth failure",
"session_id", capturedSessionID, "error", termErr)
}
return mcp.NewToolResultError(fmt.Sprintf("Unauthorized: %v", callErr)), nil
}
return mcp.NewToolResultError(callErr.Error()), nil
}

return &mcp.CallToolResult{
Result: mcp.Result{Meta: conversion.ToMCPMeta(result.Meta)},
Content: conversion.ToMCPContents(result.Content),
StructuredContent: result.StructuredContent,
IsError: result.IsError,
}, nil
}

sdkTools = append(sdkTools, mcpserver.ServerTool{
Tool: tool,
Handler: handler,
})
}

return sdkTools, nil
}

// monitorOptimizer wraps an optimizer factory so that every Optimizer instance
// produced by it is decorated with telemetry (metrics + traces).
func monitorOptimizer(
Expand Down
10 changes: 5 additions & 5 deletions pkg/vmcp/server/sessionmanager/session_manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,8 @@ type Manager struct {
}

// New creates a Manager backed by the given SessionDataStorage and backend
// registry. It builds the decorating session factory from cfg, wiring the
// optimizer and composite tool layers internally.
// registry. It wraps cfg.Base as the session factory and exposes any resolved
// optimizer factory for the Serve layer to compose with the core tools.
//
// An optimizer (FactoryConfig.OptimizerFactory or OptimizerConfig) requires
// FactoryConfig.AdvertiseFromCore; New rejects the combination otherwise. The
Expand Down Expand Up @@ -135,8 +135,8 @@ func New(
return nil, nil, err
}

// Build the Manager first so we can reference sm.Terminate and sm.sessions
// directly in closures, eliminating the forward-reference variable pattern.
// Build the Manager first so cache callbacks can reference sm.loadSession,
// sm.checkSession, and sm.sessions without a forward-reference variable.
sm := &Manager{
storage: storage,
backendReg: backendRegistry,
Expand Down Expand Up @@ -168,7 +168,7 @@ func New(
},
)

sm.factory = buildDecoratingFactory(cfg, optimizerFactory, sm.Terminate)
sm.factory = vmcpsession.NewDecoratingFactory(cfg.Base)

@ChrisJBurns ChrisJBurns Aug 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New()'s doc comment at lines 95-96 still says: "It builds the decorating session factory from cfg, wiring the optimizer and composite tool layers internally." That's no longer accurate — the factory is now just cfg.Base with no decorators, and the optimizer is surfaced solely via sm.optimizerFactory/OptimizerFactory() for the Serve layer to consume. Worth rewording.

@ChrisJBurns ChrisJBurns Aug 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment at lines 132-133 above sm := &Manager{...} says "Build the Manager first so we can reference sm.Terminate and sm.sessions directly in closures..." — with buildDecoratingFactory (and its sm.Terminate param) gone, nothing built in New() references sm.Terminate anymore, only sm.loadSession/sm.checkSession/sm.sessions. The comment should drop the sm.Terminate mention.


cleanup := func(ctx context.Context) error {
return optimizerCleanup(ctx)
Expand Down
Loading