diff --git a/packages/gitbook/src/middleware.ts b/packages/gitbook/src/middleware.ts
index 7d24a3b98..a7bce7e66 100644
--- a/packages/gitbook/src/middleware.ts
+++ b/packages/gitbook/src/middleware.ts
@@ -24,6 +24,7 @@ import {
     normalizeVisitorAuthURL,
 } from '@/lib/visitors';
 import { serveResizedImage } from '@/routes/image';
+import { cookies } from 'next/headers';
 import type { SiteURLData } from './lib/context';
 export const config = {
     matcher: [
@@ -545,9 +546,16 @@ function appendQueryParams(url: URL, from: URLSearchParams) {
 /**
  * Write the cookies to a response.
  */
-function writeResponseCookies<R extends NextResponse>(response: R, cookies: ResponseCookies): R {
-    cookies.forEach((cookie) => {
-        response.cookies.set(cookie.name, cookie.value, cookie.options);
+async function writeResponseCookies<R extends NextResponse>(
+    response: R,
+    cookiesToSet: ResponseCookies
+): Promise<R> {
+    const cookiesFn = await cookies();
+    cookiesToSet.forEach((cookie) => {
+        // response.cookies.set(cookie.name, cookie.value, cookie.options);
+        // For some reason we have to use the cookies function instead of response.cookies.set
+        // Without it, it breaks the ai assistant server actions (it thinks it is a static route).
+        cookiesFn.set(cookie.name, cookie.value, cookie.options);
     });
 
     return response;