From 02acbe1fb1f78120faf51783d6402976d0dec4e5 Mon Sep 17 00:00:00 2001
From: ephraimduncan <ephraimduncan68@gmail.com>
Date: Thu, 30 Jul 2026 16:48:09 +0000
Subject: [PATCH] fix(elements): bound code block token cache with LRU eviction

Streaming tool input/output re-highlights every intermediate version of a
code string, and each unique string was cached forever in the module-scoped
tokensCache, growing tab memory without limit during long conversations.

Cap the cache at 100 entries: refresh recency on hit, evict the oldest
insertion when full. The language highlighter cache stays unbounded since
it is naturally limited by the set of loaded languages.
---
 .../elements/__tests__/code-block.test.tsx    | 41 ++++++++++++++++++-
 packages/elements/src/code-block.tsx          | 34 ++++++++++-----
 2 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/packages/elements/__tests__/code-block.test.tsx b/packages/elements/__tests__/code-block.test.tsx
index b05fb472..7b1c6fe3 100644
--- a/packages/elements/__tests__/code-block.test.tsx
+++ b/packages/elements/__tests__/code-block.test.tsx
@@ -1,7 +1,11 @@
 import { render, screen, waitFor } from "@testing-library/react";
 import { userEvent } from "@testing-library/user-event";
 
-import { CodeBlock, CodeBlockCopyButton } from "../src/code-block";
+import {
+  CodeBlock,
+  CodeBlockCopyButton,
+  highlightCode,
+} from "../src/code-block";
 
 describe("codeBlock", () => {
   it("renders code content", async () => {
@@ -148,3 +152,38 @@ describe("codeBlockCopyButton", () => {
     });
   });
 });
+
+// Must match MAX_TOKENS_CACHE_SIZE in src/code-block.tsx
+const TOKENS_CACHE_BOUND = 100;
+
+const awaitHighlight = (code: string) =>
+  new Promise<void>((resolve) => {
+    const result = highlightCode(code, "javascript", () => resolve());
+    if (result) {
+      resolve();
+    }
+  });
+
+describe("token cache eviction", () => {
+  it("bounds the cache and keeps recently used entries", async () => {
+    const batchA = Array.from(
+      { length: TOKENS_CACHE_BOUND },
+      (_, i) => `const a${i} = ${i};`
+    );
+    await Promise.all(batchA.map((code) => awaitHighlight(code)));
+
+    // Every entry of the fill batch is retained; refresh recency of the first
+    expect(highlightCode(batchA[0], "javascript")).not.toBeNull();
+
+    // Overflow the cache with new unique entries
+    const batchB = Array.from(
+      { length: TOKENS_CACHE_BOUND - 1 },
+      (_, i) => `const b${i} = ${i};`
+    );
+    await Promise.all(batchB.map((code) => awaitHighlight(code)));
+
+    // The refreshed entry survived eviction; unrefreshed entries did not
+    expect(highlightCode(batchA[0], "javascript")).not.toBeNull();
+    expect(highlightCode(batchA[1], "javascript")).toBeNull();
+  });
+});
diff --git a/packages/elements/src/code-block.tsx b/packages/elements/src/code-block.tsx
index 820142d2..e02ff1a3 100644
--- a/packages/elements/src/code-block.tsx
+++ b/packages/elements/src/code-block.tsx
@@ -1,5 +1,13 @@
 "use client";
 
+import type { ComponentProps, CSSProperties, HTMLAttributes } from "react";
+import type {
+  BundledLanguage,
+  BundledTheme,
+  HighlighterGeneric,
+  ThemedToken,
+} from "shiki";
+
 import { Button } from "@repo/shadcn-ui/components/ui/button";
 import {
   Select,
@@ -10,7 +18,6 @@ import {
 } from "@repo/shadcn-ui/components/ui/select";
 import { cn } from "@repo/shadcn-ui/lib/utils";
 import { CheckIcon, CopyIcon } from "lucide-react";
-import type { ComponentProps, CSSProperties, HTMLAttributes } from "react";
 import {
   createContext,
   memo,
@@ -21,12 +28,6 @@ import {
   useRef,
   useState,
 } from "react";
-import type {
-  BundledLanguage,
-  BundledTheme,
-  HighlighterGeneric,
-  ThemedToken,
-} from "shiki";
 import { createHighlighter } from "shiki";
 
 // Shiki uses bitflags for font styles: 1=italic, 2=bold, 4=underline
@@ -135,7 +136,9 @@ const highlighterCache = new Map<
   Promise<HighlighterGeneric<BundledLanguage, BundledTheme>>
 >();
 
-// Token cache
+// Token cache — bounded LRU so streamed intermediate versions of a code
+// string (each producing a unique key) don't accumulate for the page lifetime
+const MAX_TOKENS_CACHE_SIZE = 100;
 const tokensCache = new Map<string, TokenizedCode>();
 
 // Subscribers for async token updates
@@ -189,9 +192,11 @@ export const highlightCode = (
 ): TokenizedCode | null => {
   const tokensCacheKey = getTokensCacheKey(code, language);
 
-  // Return cached result if available
+  // Return cached result if available, refreshing its LRU recency
   const cached = tokensCache.get(tokensCacheKey);
   if (cached) {
+    tokensCache.delete(tokensCacheKey);
+    tokensCache.set(tokensCacheKey, cached);
     return cached;
   }
 
@@ -224,7 +229,16 @@ export const highlightCode = (
         tokens: result.tokens,
       };
 
-      // Cache the result
+      // Cache the result, evicting the least recently used entry when full
+      if (
+        !tokensCache.has(tokensCacheKey) &&
+        tokensCache.size >= MAX_TOKENS_CACHE_SIZE
+      ) {
+        const oldestKey = tokensCache.keys().next().value;
+        if (oldestKey !== undefined) {
+          tokensCache.delete(oldestKey);
+        }
+      }
       tokensCache.set(tokensCacheKey, tokenized);
 
       // Notify all subscribers