Respond with all tags of all cache used for page generation

Author: SamyPesse

packages/gitbook/src/cloudflare-entrypoint.ts

@@ -1,6 +1,7 @@
 // @ts-ignore
 import nextOnPagesHandler from '@cloudflare/next-on-pages/fetch-handler';
 
+import { withResponseCacheTags } from './lib/cache';
 import { withMiddlewareHeadersStorage } from './lib/middleware';
 
 /**
@@ -9,8 +10,8 @@ import { withMiddlewareHeadersStorage } from './lib/middleware';
  */
 export default {
     async fetch(request, env, ctx) {
-        const response = await withMiddlewareHeadersStorage(() =>
-            nextOnPagesHandler.fetch(request, env, ctx)
+        const response = await withResponseCacheTags(() =>
+            withMiddlewareHeadersStorage(() => nextOnPagesHandler.fetch(request, env, ctx))
         );
 
         return response;

packages/gitbook/src/lib/cache/cache.ts

@@ -7,6 +7,7 @@ import { assertIsNotV2 } from '../v2';
 import { waitUntil } from '../waitUntil';
 import { cacheBackends } from './backends';
 import { memoryCache } from './memory';
+import { addResponseCacheTag } from './response';
 import type { CacheBackend, CacheEntry } from './types';
 
 export type CacheFunctionOptions = {
@@ -134,6 +135,11 @@ export function cache<Args extends any[], Result>(
             let result: readonly [CacheEntry, string] | null = null;
             const tag = cacheDef.tag?.(...args);
 
+            // Add the cache tag to the HTTP response
+            if (tag) {
+                addResponseCacheTag(tag);
+            }
+
             // Try the memory backend, independently of the other backends as it doesn't have a network cost
             const memoryEntry = await memoryCache.get({ key, tag });
             if (memoryEntry) {

packages/gitbook/src/lib/cache/index.ts

@@ -1,3 +1,4 @@
 export * from './cache';
 export * from './http';
 export * from './revalidateTags';
+export * from './response';

packages/gitbook/src/lib/cache/response.ts

@@ -0,0 +1,44 @@
+import { AsyncLocalStorage } from 'node:async_hooks';
+
+const symbolResponseCacheTagsLocalStorage = Symbol.for('__gitbook-response-cache-tags__');
+
+/**
+ * Wrap a response handler with a storage to store cache tags and write them as header to the response.
+ */
+export async function withResponseCacheTags(handler: () => Promise<Response>): Promise<Response> {
+    const responseCacheTagsLocalStorage =
+        // @ts-ignore
+        (globalThis[symbolResponseCacheTagsLocalStorage] as AsyncLocalStorage<Set<string>>) ??
+        new AsyncLocalStorage<Set<string>>();
+    // @ts-ignore
+    globalThis[symbolResponseCacheTagsLocalStorage] = responseCacheTagsLocalStorage;
+
+    const responseCacheTags = new Set<string>();
+    const response = await responseCacheTagsLocalStorage.run(responseCacheTags, handler);
+
+    if (responseCacheTags.size > 0) {
+        const headerCacheTag = Array.from(responseCacheTags).join(',');
+        response.headers.set('cache-tag', headerCacheTag);
+        response.headers.set('x-gitbook-cache-tag', headerCacheTag);
+    }
+
+    return response;
+}
+
+/**
+ * Add a cache tag to the HTTP response.
+ * This is a hack for v1, for v2, we use ISR and do not cache the whole response.
+ */
+export function addResponseCacheTag(...tags: string[]) {
+    const responseCacheTagsLocalStorage =
+        // @ts-ignore
+        (globalThis[symbolResponseCacheTagsLocalStorage] as AsyncLocalStorage<Set<string>>) ??
+        new AsyncLocalStorage<Set<string>>();
+    // @ts-ignore
+    globalThis[symbolResponseCacheTagsLocalStorage] = responseCacheTagsLocalStorage;
+
+    const store = responseCacheTagsLocalStorage.getStore();
+    tags.forEach((tag) => {
+        store?.add(tag);
+    });
+}

packages/gitbook/src/lib/middleware.ts

@@ -1,5 +1,9 @@
 import { AsyncLocalStorage } from 'node:async_hooks';
 
+/*
+ * This code is ONLY for v1.
+ */
+
 /**
  * Set a header on the middleware response.
  * We do this because of https://github.com/opennextjs/opennextjs-cloudflare/issues/92

packages/gitbook/src/middleware.ts

@@ -33,6 +33,7 @@ import {
 import { joinPath, normalizePathname } from '@/lib/paths';
 import { getProxyModeBasePath } from '@/lib/proxy';
 import { MiddlewareHeaders } from '@v2/lib/middleware';
+import { addResponseCacheTag } from './lib/cache';
 
 export const config = {
     matcher:
@@ -287,10 +288,8 @@ export async function middleware(request: NextRequest) {
         }
     }
 
-    if (resolved.cacheTags && resolved.cacheTags.length > 0) {
-        const headerCacheTag = resolved.cacheTags.join(',');
-        setMiddlewareHeader(response, 'cache-tag', headerCacheTag);
-        setMiddlewareHeader(response, 'x-gitbook-cache-tag', headerCacheTag);
+    if (resolved.cacheTags) {
+        addResponseCacheTag(...resolved.cacheTags);
     }
 
     return response;