fix(devshard): bound the gateway chat response cache

devshard/cmd/devshardctl/response_cache.go

@@ -6,17 +6,29 @@ import (
 	"encoding/hex"
 	"io"
 	"net/http"
+	"sort"
 	"strings"
 	"sync"
 	"time"
 )
 
 const chatResponseCacheTTL = time.Hour
 
+// defaultChatResponseCacheMaxEntries bounds how many cached chat responses the
+// gateway keeps in memory at once. The cache is a single process-wide dedup
+// layer keyed by the full request body, so under a normal chat workload
+// (mostly unique bodies) it would otherwise grow by one full response body per
+// distinct request and never shrink: an expired entry is only reclaimed when
+// the *same* key is looked up again, which for unique bodies never happens.
+// Capping the entry count keeps the footprint bounded regardless of traffic or
+// uptime, matching the bounded-map idiom used by rateLimiter and gossip dedup.
+const defaultChatResponseCacheMaxEntries = 2048
+
 type chatResponseCache struct {
-	mu      sync.Mutex
-	ttl     time.Duration
-	entries map[string]cachedChatResponse
+	mu         sync.Mutex
+	ttl        time.Duration
+	maxEntries int
+	entries    map[string]cachedChatResponse
 }
 
 type cachedChatResponse struct {
@@ -34,8 +46,9 @@ func newChatResponseCache(ttl time.Duration) *chatResponseCache {
 		ttl = chatResponseCacheTTL
 	}
 	return &chatResponseCache{
-		ttl:     ttl,
-		entries: make(map[string]cachedChatResponse),
+		ttl:        ttl,
+		maxEntries: defaultChatResponseCacheMaxEntries,
+		entries:    make(map[string]cachedChatResponse),
 	}
 }
 
@@ -83,9 +96,43 @@ func (c *chatResponseCache) Set(key string, entry cachedChatResponse, now time.T
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	if _, exists := c.entries[key]; !exists && c.maxEntries > 0 && len(c.entries) >= c.maxEntries {
+		c.evictLocked(now)
+	}
 	c.entries[key] = entry
 }
 
+// evictLocked bounds the cache. It first reclaims every expired entry (the
+// common case for a unique-body workload) and, only if the cache is still at
+// capacity with live entries, drops the oldest half in a single pass. Batching
+// amortizes the O(n) sweep across ~maxEntries/2 inserts instead of running it
+// on every Set once the cache is hot. Callers must hold c.mu. On return the
+// cache holds fewer than maxEntries entries, leaving room for the insert.
+func (c *chatResponseCache) evictLocked(now time.Time) {
+	for key, entry := range c.entries {
+		if !entry.ExpiresAt.After(now) {
+			delete(c.entries, key)
+		}
+	}
+	if len(c.entries) < c.maxEntries {
+		return
+	}
+	// ExpiresAt is insertion time plus a constant TTL, so ordering by ExpiresAt
+	// is insertion order. Keep the newest maxEntries/2 entries, drop the rest.
+	type keyExpiry struct {
+		key     string
+		expires time.Time
+	}
+	live := make([]keyExpiry, 0, len(c.entries))
+	for key, entry := range c.entries {
+		live = append(live, keyExpiry{key: key, expires: entry.ExpiresAt})
+	}
+	sort.Slice(live, func(i, j int) bool { return live[i].expires.Before(live[j].expires) })
+	for i := 0; i < len(live)-c.maxEntries/2; i++ {
+		delete(c.entries, live[i].key)
+	}
+}
+
 func serveCachedChatResponse(w http.ResponseWriter, r *http.Request, entry cachedChatResponse) {
 	if rid, ok := requestLogFromContext(r.Context()); ok {
 		w.Header().Set("X-Request-Id", rid)

devshard/cmd/devshardctl/response_cache_test.go

@@ -0,0 +1,135 @@
+package main
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func cacheTestEntry(escrow, body string, expires time.Time) cachedChatResponse {
+	return cachedChatResponse{
+		EscrowID:  escrow,
+		Body:      []byte(body),
+		ExpiresAt: expires,
+	}
+}
+
+// cacheLen reports the entry count without depending on any exported accessor,
+// so this bug-fix test stays self-contained.
+func cacheLen(c *chatResponseCache) int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return len(c.entries)
+}
+
+func TestChatResponseCache_SetGetRoundTrip(t *testing.T) {
+	c := newChatResponseCache(time.Hour)
+	now := time.Unix(1000, 0)
+	key := chatCacheKey("model-a", []byte("body-a"))
+
+	c.Set(key, cacheTestEntry("escrow-1", "hello", time.Time{}), now)
+
+	got, ok := c.Get(key, now)
+	require.True(t, ok)
+	require.Equal(t, "escrow-1", got.EscrowID)
+	require.Equal(t, "hello", string(got.Body))
+
+	// The returned body must be a copy: mutating it cannot corrupt the cache.
+	got.Body[0] = 'H'
+	again, ok := c.Get(key, now)
+	require.True(t, ok)
+	require.Equal(t, "hello", string(again.Body))
+}
+
+func TestChatResponseCache_ExpiredEntryNotServed(t *testing.T) {
+	c := newChatResponseCache(time.Hour)
+	now := time.Unix(1000, 0)
+	key := chatCacheKey("m", []byte("b"))
+
+	c.Set(key, cacheTestEntry("e", "x", now.Add(time.Minute)), now)
+
+	_, ok := c.Get(key, now.Add(2*time.Minute))
+	require.False(t, ok)
+	require.Equal(t, 0, cacheLen(c), "expired entry should be dropped on access")
+}
+
+// TestChatResponseCache_BoundsEntryCount is the regression test for the
+// unbounded-growth bug: streaming unique request bodies (the normal chat
+// workload) must never let the cache exceed its configured maximum.
+func TestChatResponseCache_BoundsEntryCount(t *testing.T) {
+	c := newChatResponseCache(time.Hour)
+	c.maxEntries = 8
+	base := time.Unix(1000, 0)
+
+	const inserted = 500
+	for i := 0; i < inserted; i++ {
+		now := base.Add(time.Duration(i) * time.Millisecond)
+		key := chatCacheKey("m", []byte(fmt.Sprintf("body-%d", i)))
+		c.Set(key, cacheTestEntry("e", fmt.Sprintf("resp-%d", i), time.Time{}), now)
+		require.LessOrEqualf(t, cacheLen(c), c.maxEntries,
+			"cache exceeded max entries (%d) after %d inserts", c.maxEntries, i+1)
+	}
+
+	// The most recent insert must survive: eviction drops the oldest, not the
+	// newest, so a just-cached response is still available for a fast retry.
+	lastNow := base.Add(time.Duration(inserted-1) * time.Millisecond)
+	lastKey := chatCacheKey("m", []byte(fmt.Sprintf("body-%d", inserted-1)))
+	got, ok := c.Get(lastKey, lastNow)
+	require.True(t, ok)
+	require.Equal(t, fmt.Sprintf("resp-%d", inserted-1), string(got.Body))
+}
+
+// TestChatResponseCache_EvictionReclaimsExpiredBeforeLive verifies that
+// eviction frees expired entries before touching live ones, so a burst of
+// stale entries does not evict a fresh, still-useful response.
+func TestChatResponseCache_EvictionReclaimsExpiredBeforeLive(t *testing.T) {
+	c := newChatResponseCache(time.Hour)
+	c.maxEntries = 4
+	base := time.Unix(1000, 0)
+
+	for i := 0; i < 4; i++ {
+		key := chatCacheKey("m", []byte(fmt.Sprintf("old-%d", i)))
+		c.Set(key, cacheTestEntry("e", "old", base.Add(10*time.Second)), base)
+	}
+	require.Equal(t, 4, cacheLen(c))
+
+	// Insert a fresh entry after the old ones have expired. Eviction should
+	// reclaim the expired entries instead of dropping the newcomer.
+	later := base.Add(20 * time.Second)
+	freshKey := chatCacheKey("m", []byte("fresh"))
+	c.Set(freshKey, cacheTestEntry("e", "fresh", time.Time{}), later)
+
+	require.LessOrEqual(t, cacheLen(c), c.maxEntries)
+	got, ok := c.Get(freshKey, later)
+	require.True(t, ok)
+	require.Equal(t, "fresh", string(got.Body))
+}
+
+// TestChatResponseCache_OverwriteDoesNotEvict confirms that re-Setting an
+// existing key updates in place without triggering eviction (the entry count
+// does not grow, so a full cache is not needlessly churned).
+func TestChatResponseCache_OverwriteDoesNotEvict(t *testing.T) {
+	c := newChatResponseCache(time.Hour)
+	c.maxEntries = 2
+	now := time.Unix(1000, 0)
+
+	k1 := chatCacheKey("m", []byte("one"))
+	k2 := chatCacheKey("m", []byte("two"))
+	c.Set(k1, cacheTestEntry("e", "one", time.Time{}), now)
+	c.Set(k2, cacheTestEntry("e", "two", time.Time{}), now)
+	require.Equal(t, 2, cacheLen(c))
+
+	// Overwriting k1 keeps the cache at capacity and must not evict k2.
+	c.Set(k1, cacheTestEntry("e", "one-v2", time.Time{}), now.Add(time.Second))
+	require.Equal(t, 2, cacheLen(c))
+
+	got, ok := c.Get(k2, now)
+	require.True(t, ok)
+	require.Equal(t, "two", string(got.Body))
+
+	got, ok = c.Get(k1, now)
+	require.True(t, ok)
+	require.Equal(t, "one-v2", string(got.Body))
+}