fix: allow cookies to be set in rewritten responses (#11280)

* fix: allow cookies to be set in rewritten responses

* Merge cookies

* Add support for endpoints and more tests

9 files changed, 118 insertions, 5 deletions

diff --git a/.changeset/big-eyes-share.md b/.changeset/big-eyes-share.md
new file mode 100644
index 000000000..f08f87c50
--- /dev/null
+++ b/.changeset/big-eyes-share.md
@@ -0,0 +1,5 @@
+---
+'astro': patch
+---
+
+Fixes a bug that prevented cookies from being set when using experimental rewrites
diff --git a/packages/astro/src/core/cookies/cookies.ts b/packages/astro/src/core/cookies/cookies.ts
index 069afc796..c176fc757 100644
--- a/packages/astro/src/core/cookies/cookies.ts
+++ b/packages/astro/src/core/cookies/cookies.ts
@@ -192,6 +192,19 @@ class AstroCookies implements AstroCookiesInterface {
 	}
 
 	/**
+	 * Merges a new AstroCookies instance into the current instance. Any new cookies
+	 * will be added to the current instance, overwriting any existing cookies with the same name.
+	 */
+	merge(cookies: AstroCookies) {
+		const outgoing = cookies.#outgoing;
+		if (outgoing) {
+			for (const [key, value] of outgoing) {
+				this.#ensureOutgoingMap().set(key, value);
+			}
+		}
+	}
+
+	/**
 	 * Astro.cookies.header() returns an iterator for the cookies that have previously
 	 * been set by either Astro.cookies.set() or Astro.cookies.delete().
 	 * This method is primarily used by adapters to set the header on outgoing responses.
diff --git a/packages/astro/src/core/cookies/response.ts b/packages/astro/src/core/cookies/response.ts
index c4dd38893..26f032fdd 100644
--- a/packages/astro/src/core/cookies/response.ts
+++ b/packages/astro/src/core/cookies/response.ts
@@ -10,7 +10,7 @@ export function responseHasCookies(response: Response): boolean {
 	return Reflect.has(response, astroCookiesSymbol);
 }
 
-function getFromResponse(response: Response): AstroCookies | undefined {
+export function getFromResponse(response: Response): AstroCookies | undefined {
 	let cookies = Reflect.get(response, astroCookiesSymbol);
 	if (cookies != null) {
 		return cookies as AstroCookies;
diff --git a/packages/astro/src/core/render-context.ts b/packages/astro/src/core/render-context.ts
index d46f80988..413922042 100644
--- a/packages/astro/src/core/render-context.ts
+++ b/packages/astro/src/core/render-context.ts
@@ -27,6 +27,7 @@ import {
 	responseSentSymbol,
 } from './constants.js';
 import { AstroCookies, attachCookiesToResponse } from './cookies/index.js';
+import { getFromResponse } from './cookies/response.js';
 import { AstroError, AstroErrorData } from './errors/index.js';
 import { callMiddleware } from './middleware/callMiddleware.js';
 import { sequence } from './middleware/index.js';
@@ -151,14 +152,17 @@ export class RenderContext {
 					);
 				}
 			}
+			let response: Response;
+
 			switch (this.routeData.type) {
-				case 'endpoint':
-					return renderEndpoint(componentInstance as any, ctx, serverLike, logger);
+				case 'endpoint': {
+					response = await renderEndpoint(componentInstance as any, ctx, serverLike, logger);
+					break;
+				}
 				case 'redirect':
 					return renderRedirect(this);
 				case 'page': {
 					const result = await this.createResult(componentInstance!);
-					let response: Response;
 					try {
 						response = await renderPage(
 							result,
@@ -185,12 +189,19 @@ export class RenderContext {
 					) {
 						response.headers.set(REROUTE_DIRECTIVE_HEADER, 'no');
 					}
-					return response;
+					break;
 				}
 				case 'fallback': {
 					return new Response(null, { status: 500, headers: { [ROUTE_TYPE_HEADER]: 'fallback' } });
 				}
 			}
+			// We need to merge the cookies from the response back into this.cookies
+			// because they may need to be passed along from a rewrite.
+			const responseCookies = getFromResponse(response);
+			if (responseCookies) {
+				cookies.merge(responseCookies);
+			}
+			return response;
 		};
 
 		const response = await callMiddleware(
diff --git a/packages/astro/test/astro-cookies.test.js b/packages/astro/test/astro-cookies.test.js
index 9d7136c4f..482474b54 100644
--- a/packages/astro/test/astro-cookies.test.js
+++ b/packages/astro/test/astro-cookies.test.js
@@ -52,6 +52,31 @@ describe('Astro.cookies', () => {
 				assert.equal(response.headers.has('set-cookie'), true);
 			}
 		});
+
+		it('can set cookies in a rewritten page request', async () => {
+			const response = await fixture.fetch('/from');
+			assert.equal(response.status, 200);
+
+			assert.match(response.headers.get('set-cookie'), /my_cookie=value/);
+		});
+
+		it('overwrites cookie values set in the source page with values from the target page', async () => {
+			const response = await fixture.fetch('/from');
+			assert.equal(response.status, 200);
+			assert.match(response.headers.get('set-cookie'), /another=set-in-target/);
+		});
+
+		it('allows cookies to be set in the source page', async () => {
+			const response = await fixture.fetch('/from');
+			assert.equal(response.status, 200);
+			assert.match(response.headers.get('set-cookie'), /set-in-from=yes/);
+		});
+
+		it('can set cookies in a rewritten endpoint request', async () => {
+			const response = await fixture.fetch('/from-endpoint');
+			assert.equal(response.status, 200);
+			assert.match(response.headers.get('set-cookie'), /test=value/);
+		});
 	});
 
 	describe('Production', () => {
@@ -140,5 +165,34 @@ describe('Astro.cookies', () => {
 			assert.equal(typeof data, 'object');
 			assert.equal(data.mode, 'dark');
 		});
+
+		it('can set cookies in a rewritten page request', async () => {
+			const request = new Request('http://example.com/from');
+			const response = await app.render(request, { addCookieHeader: true });
+			assert.equal(response.status, 200);
+
+			assert.match(response.headers.get('Set-Cookie'), /my_cookie=value/);
+		});
+
+		it('overwrites cookie values set in the source page with values from the target page', async () => {
+			const request = new Request('http://example.com/from');
+			const response = await app.render(request, { addCookieHeader: true });
+			assert.equal(response.status, 200);
+			assert.match(response.headers.get('Set-Cookie'), /another=set-in-target/);
+		});
+
+		it('allows cookies to be set in the source page', async () => {
+			const request = new Request('http://example.com/from');
+			const response = await app.render(request, { addCookieHeader: true });
+			assert.equal(response.status, 200);
+			assert.match(response.headers.get('Set-Cookie'), /set-in-from=yes/);
+		});
+
+		it('can set cookies in a rewritten endpoint request', async () => {
+			const request = new Request('http://example.com/from-endpoint');
+			const response = await app.render(request, { addCookieHeader: true });
+			assert.equal(response.status, 200);
+			assert.match(response.headers.get('Set-Cookie'), /test=value/);
+		});
 	});
 });
diff --git a/packages/astro/test/fixtures/astro-cookies/src/pages/from-endpoint.ts b/packages/astro/test/fixtures/astro-cookies/src/pages/from-endpoint.ts
new file mode 100644
index 000000000..08c586dcf
--- /dev/null
+++ b/packages/astro/test/fixtures/astro-cookies/src/pages/from-endpoint.ts
@@ -0,0 +1,3 @@
+export async function GET(context) {
+	return context.rewrite('/to-endpoint');
+}
diff --git a/packages/astro/test/fixtures/astro-cookies/src/pages/from.astro b/packages/astro/test/fixtures/astro-cookies/src/pages/from.astro
new file mode 100644
index 000000000..893858a28
--- /dev/null
+++ b/packages/astro/test/fixtures/astro-cookies/src/pages/from.astro
@@ -0,0 +1,5 @@
+---
+Astro.cookies.set('another','set-in-from');
+Astro.cookies.set('set-in-from','yes');
+return Astro.rewrite('/rewrite-target');
+---
diff --git a/packages/astro/test/fixtures/astro-cookies/src/pages/rewrite-target.astro b/packages/astro/test/fixtures/astro-cookies/src/pages/rewrite-target.astro
new file mode 100644
index 000000000..72b62ef86
--- /dev/null
+++ b/packages/astro/test/fixtures/astro-cookies/src/pages/rewrite-target.astro
@@ -0,0 +1,18 @@
+---
+Astro.cookies.set('my_cookie', 'value')
+Astro.cookies.set('another','set-in-target');
+
+---
+
+<html lang="en">
+	<head>
+		<meta charset="utf-8" />
+		<link rel="icon" type="image/svg+xml" href="/favicon.svg" />
+		<meta name="viewport" content="width=device-width" />
+		<meta name="generator" content={Astro.generator} />
+		<title>Page 2</title>
+	</head>
+  <body>
+    <h1>Page 2</h1>
+  </body>
+</html>
diff --git a/packages/astro/test/fixtures/astro-cookies/src/pages/to-endpoint.ts b/packages/astro/test/fixtures/astro-cookies/src/pages/to-endpoint.ts
new file mode 100644
index 000000000..26b8ed46c
--- /dev/null
+++ b/packages/astro/test/fixtures/astro-cookies/src/pages/to-endpoint.ts
@@ -0,0 +1,4 @@
+export async function GET(context) {
+	context.cookies.set('test', 'value');
+	return Response.json({hi: "world"})
+}