From 6c4ddb91ee06a9f764fe10e0a2f3a9ecf0934750 Mon Sep 17 00:00:00 2001
From: Matthew Phillips <matthew@skypack.dev>
Date: Mon, 10 Jun 2024 14:43:31 -0400
Subject: [PATCH] In Vercel Edge, include cookies set by Astro.cookies.set

---
 .changeset/shaggy-camels-dream.md             |  5 ++++
 packages/integrations/vercel/package.json     |  2 +-
 .../vercel/src/serverless/adapter.ts          |  1 +
 .../vercel/src/serverless/middleware.ts       | 23 ++++++++++++++-----
 .../vercel/test/edge-middleware.test.js       |  9 ++++++++
 .../src/middleware.js                         |  1 +
 6 files changed, 34 insertions(+), 7 deletions(-)
 create mode 100644 .changeset/shaggy-camels-dream.md

diff --git a/.changeset/shaggy-camels-dream.md b/.changeset/shaggy-camels-dream.md
new file mode 100644
index 000000000000..daeaeb5eec90
--- /dev/null
+++ b/.changeset/shaggy-camels-dream.md
@@ -0,0 +1,5 @@
+---
+'@astrojs/vercel': patch
+---
+
+In Vercel Edge, include cookies set by Astro.cookies.set
diff --git a/packages/integrations/vercel/package.json b/packages/integrations/vercel/package.json
index 2fc1f1eb5d7b..6f7d2da3d24b 100644
--- a/packages/integrations/vercel/package.json
+++ b/packages/integrations/vercel/package.json
@@ -47,7 +47,7 @@
     "build": "astro-scripts build \"src/**/*.ts\" && tsc",
     "build:ci": "astro-scripts build \"src/**/*.ts\"",
     "dev": "astro-scripts dev \"src/**/*.ts\"",
-    "test": "astro-scripts test --timeout 50000 \"test/**/!(hosted|edge-middleware).test.js\"",
+    "test": "astro-scripts test --timeout 50000 \"test/**/!(hosted).test.js\"",
     "test:hosted": "astro-scripts test --timeout 30000 \"test/hosted/*.test.js\""
   },
   "dependencies": {
diff --git a/packages/integrations/vercel/src/serverless/adapter.ts b/packages/integrations/vercel/src/serverless/adapter.ts
index 65f5fafd798f..4481275129ff 100644
--- a/packages/integrations/vercel/src/serverless/adapter.ts
+++ b/packages/integrations/vercel/src/serverless/adapter.ts
@@ -508,6 +508,7 @@ class VercelBuilder {
 
 		await generateEdgeMiddleware(
 			entry,
+			this.config.root,
 			new URL(VERCEL_EDGE_MIDDLEWARE_FILE, this.config.srcDir),
 			new URL('./middleware.mjs', functionFolder),
 			middlewareSecret,
diff --git a/packages/integrations/vercel/src/serverless/middleware.ts b/packages/integrations/vercel/src/serverless/middleware.ts
index f232b1060262..e8e6adf9f90f 100644
--- a/packages/integrations/vercel/src/serverless/middleware.ts
+++ b/packages/integrations/vercel/src/serverless/middleware.ts
@@ -22,6 +22,7 @@ import {
  */
 export async function generateEdgeMiddleware(
 	astroMiddlewareEntryPointPath: URL,
+	root: URL,
 	vercelEdgeMiddlewareHandlerPath: URL,
 	outPath: URL,
 	middlewareSecret: string,
@@ -39,7 +40,7 @@ export async function generateEdgeMiddleware(
 	await esbuild.build({
 		stdin: {
 			contents: code,
-			resolveDir: process.cwd(),
+			resolveDir: fileURLToPath(root),
 		},
 		target: 'es2020',
 		platform: 'browser',
@@ -96,18 +97,28 @@ export default async function middleware(request, context) {
 	});
 	ctx.locals = { vercel: { edge: context }, ...${handlerTemplateCall} };
 	const { origin } = new URL(request.url);
-	const next = () => {
+	const next = async () => {
 		const { vercel, ...locals } = ctx.locals;
-		return fetch(new URL('/${NODE_PATH}', request.url), {
+		const response = await fetch(new URL('/${NODE_PATH}', request.url), {
 			headers: {
 				...Object.fromEntries(request.headers.entries()),
 				'${ASTRO_MIDDLEWARE_SECRET_HEADER}': '${middlewareSecret}',
 				'${ASTRO_PATH_HEADER}': request.url.replace(origin, ''),
 				'${ASTRO_LOCALS_HEADER}': trySerializeLocals(locals)
 			}
-		})
-	}
+		});
+		return new Response(response.body, {
+			status: response.status,
+			statusText: response.statusText,
+			headers: response.headers,
+		});
+	};
 
-	return onRequest(ctx, next);
+	const response = await onRequest(ctx, next);
+	// Append cookies from Astro.cookies
+	for(const setCookieHeaderValue of ctx.cookies.headers()) {
+		response.headers.append('set-cookie', setCookieHeaderValue);
+	}
+	return response;
 }`;
 }
diff --git a/packages/integrations/vercel/test/edge-middleware.test.js b/packages/integrations/vercel/test/edge-middleware.test.js
index 668d92808577..54d58c02e202 100644
--- a/packages/integrations/vercel/test/edge-middleware.test.js
+++ b/packages/integrations/vercel/test/edge-middleware.test.js
@@ -30,6 +30,15 @@ describe('Vercel edge middleware', () => {
 		);
 	});
 
+	it('edge sets Set-Cookie headers', async () => {
+		let entry = new URL('../.vercel/output/functions/_middleware.func/middleware.mjs', build.config.outDir);
+		const module = await import(entry);
+		const request = new Request('http://example.com/foo');
+		const response = await module.default(request, {});
+		assert.equal(response.headers.get('set-cookie'), 'foo=bar');
+		assert.ok((await response.text()).length, 'Body is included');
+	});
+
 	// TODO: The path here seems to be inconsistent?
 	it.skip('with edge handle file, should successfully build the middleware', async () => {
 		const fixture = await loadFixture({
diff --git a/packages/integrations/vercel/test/fixtures/middleware-with-edge-file/src/middleware.js b/packages/integrations/vercel/test/fixtures/middleware-with-edge-file/src/middleware.js
index 349a0aa79581..b1a67f550965 100644
--- a/packages/integrations/vercel/test/fixtures/middleware-with-edge-file/src/middleware.js
+++ b/packages/integrations/vercel/test/fixtures/middleware-with-edge-file/src/middleware.js
@@ -3,6 +3,7 @@
  */
 export const onRequest = async (context, next) => {
 	const test = 'something';
+	context.cookies.set('foo', 'bar');
 	const response = await next();
 	return response;
 };