Add encryptedBoundArgs to resume data cache

packages/next/src/server/app-render/encryption.ts

@@ -16,7 +16,11 @@ import {
   getServerModuleMap,
   stringToUint8Array,
 } from './encryption-utils'
-import { workUnitAsyncStorage } from './work-unit-async-storage.external'
+import {
+  getPrerenderResumeDataCache,
+  getRenderResumeDataCache,
+  workUnitAsyncStorage,
+} from './work-unit-async-storage.external'
 
 const isEdgeRuntime = process.env.NEXT_RUNTIME === 'edge'
 
@@ -112,25 +116,35 @@ export async function encryptActionBoundArgs(actionId: string, args: any[]) {
 
   const workUnitStore = workUnitAsyncStorage.getStore()
 
-  const prerenderStore =
-    workUnitStore?.type === 'prerender' ? workUnitStore : undefined
+  const prerenderResumeDataCache = workUnitStore
+    ? getPrerenderResumeDataCache(workUnitStore)
+    : null
+  const renderResumeDataCache = workUnitStore
+    ? getRenderResumeDataCache(workUnitStore)
+    : null
+
+  const cacheKey = actionId + serialized
 
   const cachedEncrypted =
-    prerenderStore?.encryptedBoundArgsCache?.get(serialized)
+    prerenderResumeDataCache?.encryptedBoundArgs.get(cacheKey) ??
+    renderResumeDataCache?.encryptedBoundArgs.get(cacheKey)
 
   if (cachedEncrypted) {
     return cachedEncrypted
   }
 
-  prerenderStore?.cacheSignal?.beginRead()
+  const cacheSignal =
+    workUnitStore?.type === 'prerender' ? workUnitStore.cacheSignal : undefined
+
+  cacheSignal?.beginRead()
 
   // Encrypt the serialized string with the action id as the salt.
   // Add a prefix to later ensure that the payload is correctly decrypted, similar
   // to a checksum.
   const encrypted = await encodeActionBoundArg(actionId, serialized)
 
-  prerenderStore?.cacheSignal?.endRead()
-  prerenderStore?.encryptedBoundArgsCache?.set(serialized, encrypted)
+  cacheSignal?.endRead()
+  prerenderResumeDataCache?.encryptedBoundArgs?.set(cacheKey, encrypted)
 
   return encrypted
 }

packages/next/src/server/resume-data-cache/cache-store.ts

@@ -16,31 +16,14 @@ type CacheStore<T> = Pick<Map<string, T>, 'entries' | 'size' | 'get' | 'set'>
 export type FetchCacheStore = CacheStore<CachedFetchValue>
 
 /**
- * Parses fetch cache entries into a FetchCacheStore
- * @param entries - The entries to parse into the store
- * @returns A new FetchCacheStore containing the entries
+ * A cache store for encrypted bound args of inline server functions.
  */
-export function parseFetchCacheStore(
-  entries: Iterable<[string, CachedFetchValue]>
-): FetchCacheStore {
-  return new Map(entries)
-}
-
-/**
- * Stringifies a FetchCacheStore into an array of key-value pairs
- * @param store - The store to stringify
- * @returns A promise that resolves to an array of key-value pairs
- */
-export function stringifyFetchCacheStore(
-  entries: IterableIterator<[string, CachedFetchValue]>
-): [string, CachedFetchValue][] {
-  return Array.from(entries)
-}
+export type EncryptedBoundArgsCacheStore = CacheStore<string>
 
 /**
- * Serialized format for cache entries
+ * Serialized format for "use cache" entries
  */
-interface CacheCacheStoreSerialized {
+interface UseCacheCacheStoreSerialized {
   value: string
   tags: string[]
   stale: number
@@ -53,18 +36,15 @@ interface CacheCacheStoreSerialized {
  * A cache store specifically for "use cache" values that stores promises of
  * cache entries.
  */
-export type UseCacheCacheStore = Pick<
-  Map<string, Promise<CacheEntry>>,
-  'entries' | 'size' | 'get' | 'set'
->
+export type UseCacheCacheStore = CacheStore<Promise<CacheEntry>>
 
 /**
  * Parses serialized cache entries into a UseCacheCacheStore
  * @param entries - The serialized entries to parse
  * @returns A new UseCacheCacheStore containing the parsed entries
  */
 export function parseUseCacheCacheStore(
-  entries: Iterable<[string, CacheCacheStoreSerialized]>
+  entries: Iterable<[string, UseCacheCacheStoreSerialized]>
 ): UseCacheCacheStore {
   const store = new Map<string, Promise<CacheEntry>>()
 
@@ -98,13 +78,13 @@ export function parseUseCacheCacheStore(
 }
 
 /**
- * Stringifies a UseCacheCacheStore into an array of key-value pairs
- * @param store - The store to stringify
+ * Serializes UseCacheCacheStore entries into an array of key-value pairs
+ * @param entries - The store entries to stringify
  * @returns A promise that resolves to an array of key-value pairs with serialized values
  */
-export async function stringifyUseCacheCacheStore(
+export async function serializeUseCacheCacheStore(
   entries: IterableIterator<[string, Promise<CacheEntry>]>
-): Promise<[string, CacheCacheStoreSerialized][]> {
+): Promise<[string, UseCacheCacheStoreSerialized][]> {
   return Promise.all(
     Array.from(entries).map(([key, value]) => {
       return value.then(async (entry) => {
@@ -131,7 +111,7 @@ export async function stringifyUseCacheCacheStore(
             expire: entry.expire,
             revalidate: entry.revalidate,
           },
-        ] as [string, CacheCacheStoreSerialized]
+        ] satisfies [string, UseCacheCacheStoreSerialized]
       })
     })
   )

packages/next/src/server/resume-data-cache/resume-data-cache.test.ts

@@ -26,7 +26,7 @@ describe('stringifyResumeDataCache', () => {
     )
 
     expect(await stringifyResumeDataCache(cache)).toMatchInlineSnapshot(
-      `"{"store":{"fetch":{},"cache":{"key":{"value":"dmFsdWU=","tags":[],"stale":0,"timestamp":0,"expire":0,"revalidate":0}}}}"`
+      `"{"store":{"fetch":{},"cache":{"key":{"value":"dmFsdWU=","tags":[],"stale":0,"timestamp":0,"expire":0,"revalidate":0}},"encryptedBoundArgs":{}}}"`
     )
   })
 })

packages/next/src/server/resume-data-cache/resume-data-cache.ts

@@ -1,10 +1,9 @@
 import {
   type UseCacheCacheStore,
   type FetchCacheStore,
-  stringifyFetchCacheStore,
-  stringifyUseCacheCacheStore,
+  serializeUseCacheCacheStore,
   parseUseCacheCacheStore,
-  parseFetchCacheStore,
+  type EncryptedBoundArgsCacheStore,
 } from './cache-store'
 
 /**
@@ -23,6 +22,12 @@ export interface RenderResumeDataCache {
    * The 'set' operation is omitted to enforce immutability.
    */
   readonly fetch: Omit<FetchCacheStore, 'set'>
+
+  /**
+   * A read-only Map store for encrypted bound args of inline server functions.
+   * The 'set' operation is omitted to enforce immutability.
+   */
+  readonly encryptedBoundArgs: Omit<EncryptedBoundArgsCacheStore, 'set'>
 }
 
 /**
@@ -41,6 +46,12 @@ export interface PrerenderResumeDataCache {
    * Supports both get and set operations to build the cache during pre-rendering.
    */
   readonly fetch: FetchCacheStore
+
+  /**
+   * A mutable Map store for encrypted bound args of inline server functions.
+   * Supports both get and set operations to build the cache during pre-rendering.
+   */
+  readonly encryptedBoundArgs: EncryptedBoundArgsCacheStore
 }
 
 type ResumeStoreSerialized = {
@@ -51,15 +62,20 @@ type ResumeStoreSerialized = {
     fetch: {
       [key: string]: any
     }
+    encryptedBoundArgs: {
+      [key: string]: string
+    }
   }
 }
 
 /**
- * Serializes a resume data cache into a JSON string for storage or transmission.
- * Handles both 'use cache' values and fetch responses.
+ * Serializes a resume data cache into a JSON string for storage or
+ * transmission. Handles 'use cache' values, fetch responses, and encrypted
+ * bound args for inline server functions.
  *
  * @param resumeDataCache - The immutable cache to serialize
- * @returns A Promise that resolves to the serialized cache as a JSON string, or 'null' if empty
+ * @returns A Promise that resolves to the serialized cache as a JSON string, or
+ * 'null' if empty
  */
 export async function stringifyResumeDataCache(
   resumeDataCache: RenderResumeDataCache | PrerenderResumeDataCache
@@ -70,11 +86,12 @@ export async function stringifyResumeDataCache(
 
   const json: ResumeStoreSerialized = {
     store: {
-      fetch: Object.fromEntries(
-        stringifyFetchCacheStore(resumeDataCache.fetch.entries())
-      ),
+      fetch: Object.fromEntries(Array.from(resumeDataCache.fetch.entries())),
       cache: Object.fromEntries(
-        await stringifyUseCacheCacheStore(resumeDataCache.cache.entries())
+        await serializeUseCacheCacheStore(resumeDataCache.cache.entries())
+      ),
+      encryptedBoundArgs: Object.fromEntries(
+        Array.from(resumeDataCache.encryptedBoundArgs.entries())
       ),
     },
   }
@@ -93,6 +110,7 @@ export function createPrerenderResumeDataCache(): PrerenderResumeDataCache {
   return {
     cache: new Map(),
     fetch: new Map(),
+    encryptedBoundArgs: new Map(),
   }
 }
 
@@ -124,6 +142,7 @@ export function createRenderResumeDataCache(
     return {
       cache: new Map(),
       fetch: new Map(),
+      encryptedBoundArgs: new Map(),
     }
   }
 
@@ -132,6 +151,7 @@ export function createRenderResumeDataCache(
   )
   return {
     cache: parseUseCacheCacheStore(Object.entries(json.store.cache)),
-    fetch: parseFetchCacheStore(Object.entries(json.store.fetch)),
+    fetch: new Map(Object.entries(json.store.fetch)),
+    encryptedBoundArgs: new Map(Object.entries(json.store.encryptedBoundArgs)),
   }
 }