Skip to content

Commit 96c4e7b

Browse files
qstearnsqstearns
committed
add caching for registry client
1 parent 26cf077 commit 96c4e7b

File tree

7 files changed

+147
-15
lines changed

7 files changed

+147
-15
lines changed

‎server/cmd/gram/deps.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import (
3838
"github.com/speakeasy-api/gram/server/internal/assets"
3939
"github.com/speakeasy-api/gram/server/internal/attr"
4040
"github.com/speakeasy-api/gram/server/internal/billing"
41+
"github.com/speakeasy-api/gram/server/internal/cache"
4142
"github.com/speakeasy-api/gram/server/internal/conv"
4243
"github.com/speakeasy-api/gram/server/internal/encryption"
4344
"github.com/speakeasy-api/gram/server/internal/externalmcp"
@@ -545,6 +546,7 @@ func newFunctionOrchestrator(
545546
type mcpRegistryClientOptions struct {
546547
pulseTenantID string
547548
pulseAPIKey conv.Secret
549+
cacheImpl cache.Cache
548550
}
549551

550552
func newMCPRegistryClient(logger *slog.Logger, tracerProvider trace.TracerProvider, opts mcpRegistryClientOptions) (*externalmcp.RegistryClient, error) {
@@ -555,5 +557,5 @@ func newMCPRegistryClient(logger *slog.Logger, tracerProvider trace.TracerProvid
555557

556558
backend := externalmcp.NewPulseBackend(pulseURL, opts.pulseTenantID, opts.pulseAPIKey)
557559

558-
return externalmcp.NewRegistryClient(logger, tracerProvider, backend), nil
560+
return externalmcp.NewRegistryClient(logger, tracerProvider, backend, opts.cacheImpl), nil
559561
}

‎server/cmd/gram/start.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -587,6 +587,7 @@ func newStartCommand() *cli.Command {
587587
mcpRegistryClient, err := newMCPRegistryClient(logger, tracerProvider, mcpRegistryClientOptions{
588588
pulseTenantID: c.String("pulse-registry-tenant"),
589589
pulseAPIKey: conv.NewSecret([]byte(c.String("pulse-registry-api-key"))),
590+
cacheImpl: cache.NewRedisCacheAdapter(redisClient),
590591
})
591592
if err != nil {
592593
return fmt.Errorf("failed to create mcp registry client: %w", err)

‎server/cmd/gram/worker.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,7 @@ func newWorkerCommand() *cli.Command {
402402
mcpRegistryClient, err := newMCPRegistryClient(logger, tracerProvider, mcpRegistryClientOptions{
403403
pulseTenantID: c.String("pulse-registry-tenant"),
404404
pulseAPIKey: conv.NewSecret([]byte(c.String("pulse-registry-api-key"))),
405+
cacheImpl: cache.NewRedisCacheAdapter(redisClient),
405406
})
406407
if err != nil {
407408
return fmt.Errorf("failed to create mcp registry client: %w", err)

‎server/internal/externalmcp/registry_cache.go‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package externalmcp
2+
3+
import (
4+
"crypto/sha256"
5+
"fmt"
6+
"net/http"
7+
"sort"
8+
"strings"
9+
"time"
10+
11+
"github.com/speakeasy-api/gram/server/gen/types"
12+
"github.com/speakeasy-api/gram/server/internal/cache"
13+
)
14+
15+
const registryCacheTTL = 24 * time.Hour
16+
17+
// CachedListServersResponse wraps a list of external MCP servers for caching.
18+
type CachedListServersResponse struct {
19+
Key string
20+
Servers []*types.ExternalMCPServer
21+
}
22+
23+
var _ cache.CacheableObject[CachedListServersResponse] = (*CachedListServersResponse)(nil)
24+
25+
func (c CachedListServersResponse) CacheKey() string { return c.Key }
26+
func (c CachedListServersResponse) AdditionalCacheKeys() []string { return []string{} }
27+
func (c CachedListServersResponse) TTL() time.Duration { return registryCacheTTL }
28+
29+
// CachedServerDetailsResponse wraps server details for caching.
30+
type CachedServerDetailsResponse struct {
31+
Key string
32+
Details *ServerDetails
33+
}
34+
35+
var _ cache.CacheableObject[CachedServerDetailsResponse] = (*CachedServerDetailsResponse)(nil)
36+
37+
func (c CachedServerDetailsResponse) CacheKey() string { return c.Key }
38+
func (c CachedServerDetailsResponse) AdditionalCacheKeys() []string { return []string{} }
39+
func (c CachedServerDetailsResponse) TTL() time.Duration { return registryCacheTTL }
40+
41+
// registryCacheKey builds a cache key from a prefix and the request's URL + headers.
42+
// Headers are sorted and hashed with SHA-256 to capture tenant/auth identity.
43+
func registryCacheKey(prefix string, req *http.Request) string {
44+
// Sort header keys for deterministic hashing
45+
keys := make([]string, 0, len(req.Header))
46+
for k := range req.Header {
47+
keys = append(keys, k)
48+
}
49+
sort.Strings(keys)
50+
51+
h := sha256.New()
52+
for _, k := range keys {
53+
vals := req.Header[k]
54+
sort.Strings(vals)
55+
_, _ = fmt.Fprintf(h, "%s=%s\n", k, strings.Join(vals, ","))
56+
}
57+
58+
return fmt.Sprintf("registry:%s:%s:%x", prefix, req.URL.String(), h.Sum(nil))
59+
}

‎server/internal/externalmcp/registryclient.go‎

Lines changed: 78 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616

1717
"github.com/speakeasy-api/gram/server/gen/types"
1818
"github.com/speakeasy-api/gram/server/internal/attr"
19+
"github.com/speakeasy-api/gram/server/internal/cache"
1920
externalmcptypes "github.com/speakeasy-api/gram/server/internal/externalmcp/repo/types"
2021
"github.com/speakeasy-api/gram/server/internal/o11y"
2122
"github.com/speakeasy-api/gram/server/internal/oops"
@@ -28,23 +29,46 @@ type RegistryBackend interface {
2829

2930
// RegistryClient handles communication with external MCP registries.
3031
type RegistryClient struct {
31-
httpClient *http.Client
32-
logger *slog.Logger
33-
backend RegistryBackend
32+
httpClient *http.Client
33+
logger *slog.Logger
34+
backend RegistryBackend
35+
listCache *cache.TypedCacheObject[CachedListServersResponse]
36+
detailsCache *cache.TypedCacheObject[CachedServerDetailsResponse]
3437
}
3538

36-
// NewRegistryClient creates a new registry client.
37-
func NewRegistryClient(logger *slog.Logger, tracerProvider trace.TracerProvider, backend RegistryBackend) *RegistryClient {
38-
return &RegistryClient{
39+
// NewRegistryClient creates a new registry client. The cacheImpl parameter is
40+
// optional — pass nil to disable caching.
41+
func NewRegistryClient(logger *slog.Logger, tracerProvider trace.TracerProvider, backend RegistryBackend, cacheImpl cache.Cache) *RegistryClient {
42+
rc := &RegistryClient{
3943
httpClient: &http.Client{
4044
Transport: otelhttp.NewTransport(
4145
retryablehttp.NewClient().StandardClient().Transport,
4246
otelhttp.WithTracerProvider(tracerProvider),
4347
),
4448
},
45-
logger: logger.With(attr.SlogComponent("mcp-registry-client")),
46-
backend: backend,
49+
logger: logger.With(attr.SlogComponent("mcp-registry-client")),
50+
backend: backend,
51+
listCache: nil,
52+
detailsCache: nil,
4753
}
54+
55+
if cacheImpl != nil {
56+
listCache := cache.NewTypedObjectCache[CachedListServersResponse](
57+
logger.With(attr.SlogCacheNamespace("registry-list")),
58+
cacheImpl,
59+
cache.SuffixNone,
60+
)
61+
rc.listCache = &listCache
62+
63+
detailsCache := cache.NewTypedObjectCache[CachedServerDetailsResponse](
64+
logger.With(attr.SlogCacheNamespace("registry-details")),
65+
cacheImpl,
66+
cache.SuffixNone,
67+
)
68+
rc.detailsCache = &detailsCache
69+
}
70+
71+
return rc
4872
}
4973

5074
// Registry represents an MCP registry endpoint.
@@ -173,6 +197,16 @@ func (c *RegistryClient) ListServers(ctx context.Context, registry Registry, par
173197
}
174198
}
175199

200+
// Check cache after authorization so headers are populated.
201+
if c.listCache != nil {
202+
cacheKey := registryCacheKey("list", req)
203+
cached, err := c.listCache.Get(ctx, cacheKey)
204+
if err == nil {
205+
c.logger.DebugContext(ctx, "registry list cache hit", attr.SlogCacheKey(cacheKey))
206+
return cached.Servers, nil
207+
}
208+
}
209+
176210
resp, err := c.httpClient.Do(req)
177211
if err != nil {
178212
return nil, fmt.Errorf("failed to fetch from registry: %w", err)
@@ -231,6 +265,17 @@ func (c *RegistryClient) ListServers(ctx context.Context, registry Registry, par
231265
servers = append(servers, server)
232266
}
233267

268+
// Store in cache on success.
269+
if c.listCache != nil {
270+
cacheKey := registryCacheKey("list", req)
271+
if storeErr := c.listCache.Store(ctx, CachedListServersResponse{
272+
Key: cacheKey,
273+
Servers: servers,
274+
}); storeErr != nil {
275+
c.logger.WarnContext(ctx, "failed to store registry list in cache", attr.SlogError(storeErr))
276+
}
277+
}
278+
234279
return servers, nil
235280
}
236281

@@ -266,6 +311,16 @@ func (c *RegistryClient) GetServerDetails(ctx context.Context, registry Registry
266311
}
267312
}
268313

314+
// Check cache after authorization so headers are populated.
315+
if c.detailsCache != nil {
316+
cacheKey := registryCacheKey("details", req)
317+
cached, err := c.detailsCache.Get(ctx, cacheKey)
318+
if err == nil {
319+
c.logger.DebugContext(ctx, "registry details cache hit", attr.SlogCacheKey(cacheKey))
320+
return cached.Details, nil
321+
}
322+
}
323+
269324
resp, err := c.httpClient.Do(req)
270325
if err != nil {
271326
return nil, fmt.Errorf("send external mcp server details request: %w", err)
@@ -326,13 +381,26 @@ func (c *RegistryClient) GetServerDetails(ctx context.Context, registry Registry
326381
tools = serverResp.Meta.Version.FifthRemote.Tools
327382
}
328383

329-
return &ServerDetails{
384+
details := &ServerDetails{
330385
Name: serverResp.Server.Name,
331386
Description: serverResp.Server.Description,
332387
Version: serverResp.Server.Version,
333388
RemoteURL: remoteURL,
334389
TransportType: transportType,
335390
Tools: tools,
336391
Headers: headers,
337-
}, nil
392+
}
393+
394+
// Store in cache on success.
395+
if c.detailsCache != nil {
396+
cacheKey := registryCacheKey("details", req)
397+
if storeErr := c.detailsCache.Store(ctx, CachedServerDetailsResponse{
398+
Key: cacheKey,
399+
Details: details,
400+
}); storeErr != nil {
401+
c.logger.WarnContext(ctx, "failed to store registry details in cache", attr.SlogError(storeErr))
402+
}
403+
}
404+
405+
return details, nil
338406
}

‎server/internal/externalmcp/registryclient_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ func TestListServers_FiltersDeletedServers(t *testing.T) {
8282
}))
8383
defer server.Close()
8484

85-
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{})
85+
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{}, nil)
8686
client.httpClient = server.Client()
8787
registry := Registry{
8888
ID: uuid.New(),
@@ -124,7 +124,7 @@ func TestGetServerDetails_OnlyStreamableHTTP(t *testing.T) {
124124
}))
125125
defer server.Close()
126126

127-
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{})
127+
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{}, nil)
128128
client.httpClient = server.Client()
129129
registry := Registry{
130130
ID: uuid.New(),
@@ -169,7 +169,7 @@ func TestGetServerDetails_OnlySSE(t *testing.T) {
169169
}))
170170
defer server.Close()
171171

172-
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{})
172+
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{}, nil)
173173
client.httpClient = server.Client()
174174
registry := Registry{
175175
ID: uuid.New(),
@@ -215,7 +215,7 @@ func TestGetServerDetails_PrefersStreamableHTTPOverSSE(t *testing.T) {
215215
}))
216216
defer server.Close()
217217

218-
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{})
218+
client := NewRegistryClient(logger, tracernoop.NewTracerProvider(), &PassthroughBackend{}, nil)
219219
client.httpClient = server.Client()
220220
registry := Registry{
221221
ID: uuid.New(),

‎server/internal/testenv/testing.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ func NewMCPRegistryClient(t *testing.T, logger *slog.Logger, tracerProvider trac
8282
NewLogger(t),
8383
tracerProvider,
8484
externalmcp.NewPulseBackend(pulseURL, "test-tenant-id", conv.NewSecret([]byte("test-api-key"))),
85+
nil,
8586
)
8687
require.NoError(t, err, "expected mcp registry client to initialize without error")
8788

0 commit comments

Comments
 (0)