Route server.New through core.New and Serve

pkg/vmcp/aggregator/caching_aggregator.go

@@ -0,0 +1,140 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package aggregator
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
+	"sort"
+	"time"
+
+	lru "github.com/hashicorp/golang-lru/v2"
+
+	"github.com/stacklok/toolhive/pkg/auth"
+	"github.com/stacklok/toolhive/pkg/vmcp"
+	"github.com/stacklok/toolhive/pkg/vmcp/headerforward"
+)
+
+// capabilityCacheMaxEntries bounds the per-identity capability cache so it cannot grow
+// without limit (one entry per distinct identity + forwarded-credential + backend-set key).
+// Beyond it, the LRU evicts the least-recently-used entry. 1024 distinct active keys per
+// node is generous for a vMCP instance; tune if real workloads exceed it.
+const capabilityCacheMaxEntries = 1024
+
+// cachingAggregator wraps an Aggregator and memoizes AggregateCapabilities for a bounded
+// TTL, so the Serve path does not re-sweep every backend's tools/list on every tool call.
+//
+// On the New/Serve path, core.CallTool/ReadResource re-derive the advertised view on every
+// call (the core holds no per-session cache); without this, a session with M calls performs
+// ~M+1 full backend sweeps instead of the legacy ~1-per-session. This decorator restores
+// once-per-(identity, TTL) freshness without coupling the core or Serve to a cache — it sits
+// transparently below the core, wrapping the Aggregator the core already calls.
+//
+// It delegates the cache mechanics (LRU eviction, size bounding, thread-safety) to
+// hashicorp/golang-lru and adds only a lazy TTL check on read. The base (non-expirable) LRU
+// is used deliberately: the expirable variant runs a perpetual background cleanup goroutine
+// with no Stop, which a per-server cache would leak (the repo's tests run under goleak).
+//
+// Security: backend enumeration is identity-dependent — what each backend returns depends on
+// the credential presented to it — so the cache key MUST include the identity and the
+// forwarded credentials, not just the backend set. Keying on (subject, forwarded headers,
+// backend IDs) ensures one caller's capability view is never served to another, while a
+// single caller's view is shared across their sessions. The key is a SHA-256 digest, so raw
+// credential values are not retained as map keys. The cache is node-local and never persisted
+// (it would be a credentialed view in shared state otherwise).
+type cachingAggregator struct {
+	// Aggregator is the wrapped aggregator; embedding delegates every method except the
+	// AggregateCapabilities override below.
+	Aggregator
+
+	ttl   time.Duration
+	cache *lru.Cache[string, cacheEntry]
+}
+
+type cacheEntry struct {
+	caps *AggregatedCapabilities
+	at   time.Time
+}
+
+// NewCachingAggregator wraps next so AggregateCapabilities results are memoized per identity
+// for ttl, backed by a size-bounded LRU. A ttl <= 0 disables caching (next is returned
+// unwrapped) so a misconfiguration cannot silently serve permanently-stale capabilities. A
+// nil next is returned as-is so the downstream nil-aggregator validation (core.New) still
+// fires rather than being masked by a non-nil wrapper.
+func NewCachingAggregator(next Aggregator, ttl time.Duration) Aggregator {
+	if next == nil || ttl <= 0 {
+		return next
+	}
+	cache, err := lru.New[string, cacheEntry](capabilityCacheMaxEntries)
+	if err != nil {
+		// lru.New only errors on a non-positive size, which is a positive constant here, so
+		// this is unreachable; degrade to the uncached aggregator rather than panicking.
+		return next
+	}
+	return &cachingAggregator{Aggregator: next, ttl: ttl, cache: cache}
+}
+
+// AggregateCapabilities returns a cached view when a fresh entry exists for the caller's
+// identity + forwarded credentials + backend set, and otherwise sweeps the backends (via the
+// wrapped aggregator) and caches the result. Errors are never cached. The returned value is
+// treated as immutable by the core (it derives fresh per-call routers from it), so the cached
+// pointer is shared rather than deep-copied.
+func (c *cachingAggregator) AggregateCapabilities(
+	ctx context.Context, backends []vmcp.Backend,
+) (*AggregatedCapabilities, error) {
+	key := cacheKey(ctx, backends)
+	if e, ok := c.cache.Get(key); ok && time.Since(e.at) < c.ttl {
+		return e.caps, nil
+	}
+
+	// Miss/expiry: sweep with the lock released (Get/Add are individually locked) so callers
+	// with different keys are not serialized behind one backend sweep. Concurrent misses for
+	// the same key may each sweep once (last writer wins) — acceptable for a cold/expired key.
+	caps, err := c.Aggregator.AggregateCapabilities(ctx, backends)
+	if err != nil {
+		return nil, err
+	}
+	c.cache.Add(key, cacheEntry{caps: caps, at: time.Now()})
+	return caps, nil
+}
+
+// cacheKey derives a collision-resistant key from the inputs that drive backend enumeration:
+// the caller's subject, the forwarded headers (passthrough credentials/scopes), and the
+// backend ID set. Hashing keeps raw credential values out of the cache keys.
+func cacheKey(ctx context.Context, backends []vmcp.Backend) string {
+	h := sha256.New()
+
+	if id, ok := auth.IdentityFromContext(ctx); ok && id != nil {
+		_, _ = io.WriteString(h, id.Subject)
+	}
+	_, _ = h.Write([]byte{0})
+
+	fwd := headerforward.ForwardedHeadersFromContext(ctx)
+	fwdKeys := make([]string, 0, len(fwd))
+	for k := range fwd {
+		fwdKeys = append(fwdKeys, k)
+	}
+	sort.Strings(fwdKeys)
+	for _, k := range fwdKeys {
+		_, _ = io.WriteString(h, k)
+		_, _ = h.Write([]byte{'='})
+		_, _ = io.WriteString(h, fwd[k])
+		_, _ = h.Write([]byte{0})
+	}
+	_, _ = h.Write([]byte{0})
+
+	ids := make([]string, 0, len(backends))
+	for _, b := range backends {
+		ids = append(ids, b.ID)
+	}
+	sort.Strings(ids)
+	for _, id := range ids {
+		_, _ = io.WriteString(h, id)
+		_, _ = h.Write([]byte{0})
+	}
+
+	return hex.EncodeToString(h.Sum(nil))
+}

pkg/vmcp/aggregator/caching_aggregator_test.go

@@ -0,0 +1,126 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package aggregator_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+
+	"github.com/stacklok/toolhive/pkg/auth"
+	"github.com/stacklok/toolhive/pkg/vmcp"
+	"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
+	"github.com/stacklok/toolhive/pkg/vmcp/aggregator/mocks"
+	"github.com/stacklok/toolhive/pkg/vmcp/headerforward"
+)
+
+func ctxWithSubject(subject string) context.Context {
+	return auth.WithIdentity(context.Background(),
+		&auth.Identity{PrincipalInfo: auth.PrincipalInfo{Subject: subject}})
+}
+
+var testBackends = []vmcp.Backend{{ID: "b1"}, {ID: "b2"}}
+
+// TestCachingAggregator_CacheHitWithinTTL: two calls with the same identity and backend set
+// within the TTL hit the cache, so the wrapped aggregator sweeps only once.
+func TestCachingAggregator_CacheHitWithinTTL(t *testing.T) {
+	t.Parallel()
+	ctrl := gomock.NewController(t)
+	mock := mocks.NewMockAggregator(ctrl)
+	caps := &aggregator.AggregatedCapabilities{}
+	mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(caps, nil).Times(1)
+
+	c := aggregator.NewCachingAggregator(mock, time.Hour)
+	ctx := ctxWithSubject("alice")
+
+	first, err := c.AggregateCapabilities(ctx, testBackends)
+	require.NoError(t, err)
+	second, err := c.AggregateCapabilities(ctx, testBackends)
+	require.NoError(t, err)
+	assert.Same(t, first, second, "the cached view must be returned on a hit")
+	assert.Same(t, caps, second)
+}
+
+// TestCachingAggregator_KeyedByIdentityAndHeaders: distinct subjects, and distinct forwarded
+// credentials for the same subject, are distinct cache keys — each sweeps the backends. This
+// is the security-critical property: one caller's view is never served to another.
+func TestCachingAggregator_KeyedByIdentityAndHeaders(t *testing.T) {
+	t.Parallel()
+	ctrl := gomock.NewController(t)
+	mock := mocks.NewMockAggregator(ctrl)
+	// alice, bob, and alice-with-a-forwarded-token are three distinct keys → three sweeps.
+	mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).
+		Return(&aggregator.AggregatedCapabilities{}, nil).Times(3)
+
+	c := aggregator.NewCachingAggregator(mock, time.Hour)
+
+	_, err := c.AggregateCapabilities(ctxWithSubject("alice"), testBackends)
+	require.NoError(t, err)
+	_, err = c.AggregateCapabilities(ctxWithSubject("bob"), testBackends)
+	require.NoError(t, err)
+	aliceWithToken := headerforward.WithForwardedHeaders(ctxWithSubject("alice"),
+		map[string]string{"Authorization": "Bearer xyz"})
+	_, err = c.AggregateCapabilities(aliceWithToken, testBackends)
+	require.NoError(t, err)
+}
+
+// TestCachingAggregator_TTLExpiry: once the TTL elapses, the next call re-sweeps.
+func TestCachingAggregator_TTLExpiry(t *testing.T) {
+	t.Parallel()
+	ctrl := gomock.NewController(t)
+	mock := mocks.NewMockAggregator(ctrl)
+	mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).
+		Return(&aggregator.AggregatedCapabilities{}, nil).Times(2)
+
+	c := aggregator.NewCachingAggregator(mock, 20*time.Millisecond)
+	ctx := ctxWithSubject("alice")
+
+	_, err := c.AggregateCapabilities(ctx, testBackends)
+	require.NoError(t, err)
+	time.Sleep(40 * time.Millisecond)
+	_, err = c.AggregateCapabilities(ctx, testBackends)
+	require.NoError(t, err)
+}
+
+// TestCachingAggregator_DisabledWhenTTLNonPositive: a ttl <= 0 returns the wrapped aggregator
+// unwrapped, so callers cannot silently get a permanently-stale cache.
+func TestCachingAggregator_DisabledWhenTTLNonPositive(t *testing.T) {
+	t.Parallel()
+	ctrl := gomock.NewController(t)
+	mock := mocks.NewMockAggregator(ctrl)
+
+	assert.Same(t, mock, aggregator.NewCachingAggregator(mock, 0))
+	assert.Same(t, mock, aggregator.NewCachingAggregator(mock, -time.Second))
+
+	// A nil aggregator is returned as-is so the downstream nil check (core.New) still fires
+	// rather than being masked by a non-nil caching wrapper.
+	var nilAgg aggregator.Aggregator
+	assert.Nil(t, aggregator.NewCachingAggregator(nilAgg, time.Hour))
+}
+
+// TestCachingAggregator_ErrorNotCached: a failed sweep is not cached, so the next call retries
+// the wrapped aggregator.
+func TestCachingAggregator_ErrorNotCached(t *testing.T) {
+	t.Parallel()
+	ctrl := gomock.NewController(t)
+	mock := mocks.NewMockAggregator(ctrl)
+	gomock.InOrder(
+		mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).Return(nil, errors.New("boom")),
+		mock.EXPECT().AggregateCapabilities(gomock.Any(), gomock.Any()).
+			Return(&aggregator.AggregatedCapabilities{}, nil),
+	)
+
+	c := aggregator.NewCachingAggregator(mock, time.Hour)
+	ctx := ctxWithSubject("alice")
+
+	_, err := c.AggregateCapabilities(ctx, testBackends)
+	require.Error(t, err)
+	_, err = c.AggregateCapabilities(ctx, testBackends)
+	require.NoError(t, err, "the error must not have been cached")
+}

pkg/vmcp/auth/factory/incoming.go

@@ -123,6 +123,40 @@ func newCedarAuthzMiddleware(
 
 	slog.Info("creating Cedar authorization middleware", "policies", len(authzCfg.Policies))
 
+	authzConfig, err := buildCedarAuthzConfig(authzCfg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create authz config: %w", err)
+	}
+
+	// Create the middleware using the existing factory
+	middlewareFn, err := authz.CreateMiddlewareFromConfig(authzConfig, serverName, passThroughTools)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create Cedar middleware: %w", err)
+	}
+
+	return middlewareFn, nil
+}
+
+// BuildAuthzConfig builds the authorizer-agnostic *authz.Config that the vMCP core
+// admission seam (core.Config.Authz) consumes from the incoming-auth config, or
+// (nil, nil) when no Cedar policies are configured.
+//
+// It is the SAME config newCedarAuthzMiddleware builds the HTTP authz middleware from,
+// surfaced so the composition root can feed it to core.New via server.Config.Authz once
+// server.New routes through Serve. The nil return mirrors that middleware's nil result
+// for the no-policies case, preserving allow-all parity (a nil core Authz is allow-all).
+func BuildAuthzConfig(authzCfg *config.AuthzConfig) (*authz.Config, error) {
+	if authzCfg == nil || authzCfg.Type != "cedar" || len(authzCfg.Policies) == 0 {
+		return nil, nil
+	}
+	return buildCedarAuthzConfig(authzCfg)
+}
+
+// buildCedarAuthzConfig converts the vMCP Cedar authz config into the authorizer-agnostic
+// authorizers.Config (aliased as authz.Config) consumed by both the HTTP authz middleware
+// (newCedarAuthzMiddleware) and the core admission seam (via BuildAuthzConfig). Callers
+// guarantee authzCfg is non-nil with at least one policy.
+func buildCedarAuthzConfig(authzCfg *config.AuthzConfig) (*authz.Config, error) {
 	// Default EntitiesJSON to "[]" when the operator/CLI did not set it. Cedar
 	// requires a valid JSON array; an empty string would fail to parse.
 	entitiesJSON := authzCfg.EntitiesJSON
@@ -149,19 +183,7 @@ func newCedarAuthzMiddleware(
 		},
 	}
 
-	// Create the authz Config using the factory method
-	authzConfig, err := authorizers.NewConfig(cedarConfig)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create authz config: %w", err)
-	}
-
-	// Create the middleware using the existing factory
-	middlewareFn, err := authz.CreateMiddlewareFromConfig(authzConfig, serverName, passThroughTools)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create Cedar middleware: %w", err)
-	}
-
-	return middlewareFn, nil
+	return authorizers.NewConfig(cedarConfig)
 }
 
 // newOIDCAuthMiddleware creates OIDC authentication middleware.

pkg/vmcp/cli/serve.go

@@ -378,6 +378,16 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 		return fmt.Errorf("failed to create authentication middleware: %w", err)
 	}
 
+	// Build the authorizer-agnostic authz config the core admission seam consumes. It is
+	// the same config the HTTP authz middleware above is built from; server.New routes
+	// through Serve, which ignores the (vestigial) AuthzMiddleware and enforces authz in
+	// the core from this config instead. Nil when no Cedar policies are configured
+	// (allow-all parity).
+	authzConfig, err := authfactory.BuildAuthzConfig(vmcpCfg.IncomingAuth.Authz)
+	if err != nil {
+		return fmt.Errorf("failed to build authorization config: %w", err)
+	}
+
 	slog.Info(fmt.Sprintf("Incoming authentication configured: %s", vmcpCfg.IncomingAuth.Type))
 
 	namespace := vmcpNamespace()
@@ -425,6 +435,11 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 		OptimizerConfig:         optCfg,
 		SessionFactory:          sessionFactory,
 		SessionStorage:          vmcpCfg.SessionStorage,
+		// Core collaborators: server.New routes through core.New + Serve, so the core
+		// is the single aggregator and authorizer. The aggregator is the same instance
+		// that backs discovery; Authz feeds the core admission seam (nil = allow-all).
+		Aggregator: agg,
+		Authz:      authzConfig,
 	})
 
 	// Assign Watcher only when backendWatcher is non-nil. A typed nil