Actions: add discriminated union support (#11939)

* feat: discriminated union for form validators

* chore: changeset

4 files changed, 131 insertions, 2 deletions

diff --git a/.changeset/mighty-stingrays-press.md b/.changeset/mighty-stingrays-press.md
new file mode 100644
index 000000000..12c353dcd
--- /dev/null
+++ b/.changeset/mighty-stingrays-press.md
@@ -0,0 +1,63 @@
+---
+'astro': patch
+---
+
+Adds support for Zod discriminated unions on Action form inputs. This allows forms with different inputs to be submitted to the same action, using a given input to decide which object should be used for validation.
+
+This example accepts either a `create` or `update` form submission, and uses the `type` field to determine which object to validate against.
+
+```ts
+import { defineAction } from 'astro:actions';
+import { z } from 'astro:schema';
+
+export const server = {
+  changeUser: defineAction({
+    accept: 'form',
+    input: z.discriminatedUnion('type', [
+      z.object({
+        type: z.literal('create'),
+        name: z.string(),
+        email: z.string().email(),
+      }),
+      z.object({
+        type: z.literal('update'),
+        id: z.number(),
+        name: z.string(),
+        email: z.string().email(),
+      }),
+    ]),
+    async handler(input) {
+      if (input.type === 'create') {
+        // input is { type: 'create', name: string, email: string }
+      } else {
+        // input is { type: 'update', id: number, name: string, email: string }
+      }
+    },
+  }),
+}
+```
+
+The corresponding `create` and `update` forms may look like this:
+
+```astro
+---
+import { actions } from 'astro:actions';
+---
+
+<!--Create-->
+<form action={actions.changeUser} method="POST">
+  <input type="hidden" name="type" value="create" />
+  <input type="text" name="name" required />
+  <input type="email" name="email" required />
+  <button type="submit">Create User</button>
+</form>
+
+<!--Update-->
+<form action={actions.changeUser} method="POST">
+  <input type="hidden" name="type" value="update" />
+  <input type="hidden" name="id" value="user-123" />
+  <input type="text" name="name" required />
+  <input type="email" name="email" required />
+  <button type="submit">Update User</button>
+</form>
+```
diff --git a/packages/astro/src/actions/runtime/virtual/server.ts b/packages/astro/src/actions/runtime/virtual/server.ts
index cd1b4269e..8e5e6bb4f 100644
--- a/packages/astro/src/actions/runtime/virtual/server.ts
+++ b/packages/astro/src/actions/runtime/virtual/server.ts
@@ -92,7 +92,7 @@ function getFormServerHandler<TOutput, TInputSchema extends z.ZodType>(
 
 		if (!inputSchema) return await handler(unparsedInput, context);
 
-		const baseSchema = unwrapSchemaEffects(inputSchema);
+		const baseSchema = unwrapBaseObjectSchema(inputSchema, unparsedInput);
 		const parsed = await inputSchema.safeParseAsync(
 			baseSchema instanceof z.ZodObject
 				? formDataToObject(unparsedInput, baseSchema)
@@ -191,7 +191,7 @@ function handleFormDataGet(
 	return validator instanceof z.ZodNumber ? Number(value) : value;
 }
 
-function unwrapSchemaEffects(schema: z.ZodType) {
+function unwrapBaseObjectSchema(schema: z.ZodType, unparsedInput: FormData) {
 	while (schema instanceof z.ZodEffects || schema instanceof z.ZodPipeline) {
 		if (schema instanceof z.ZodEffects) {
 			schema = schema._def.schema;
@@ -200,5 +200,15 @@ function unwrapSchemaEffects(schema: z.ZodType) {
 			schema = schema._def.in;
 		}
 	}
+	if (schema instanceof z.ZodDiscriminatedUnion) {
+		const typeKey = schema._def.discriminator;
+		const typeValue = unparsedInput.get(typeKey);
+		if (typeof typeValue !== 'string') return schema;
+
+		const objSchema = schema._def.optionsMap.get(typeValue);
+		if (!objSchema) return schema;
+
+		return objSchema;
+	}
 	return schema;
 }
diff --git a/packages/astro/test/actions.test.js b/packages/astro/test/actions.test.js
index 334e07a17..17758e82c 100644
--- a/packages/astro/test/actions.test.js
+++ b/packages/astro/test/actions.test.js
@@ -395,6 +395,39 @@ describe('Astro Actions', () => {
 			assert.ok(value.date instanceof Date);
 			assert.ok(value.set instanceof Set);
 		});
+
+		it('Supports discriminated union for different form fields', async () => {
+			const formData = new FormData();
+			formData.set('type', 'first-chunk');
+			formData.set('alt', 'Cool image');
+			formData.set('image', new File([''], 'chunk-1.png'));
+			const reqFirst = new Request('http://example.com/_actions/imageUploadInChunks', {
+				method: 'POST',
+				body: formData,
+			});
+
+			const resFirst = await app.render(reqFirst);
+			assert.equal(resFirst.status, 200);
+			assert.equal(resFirst.headers.get('Content-Type'), 'application/json+devalue');
+			const data = devalue.parse(await resFirst.text());
+			const uploadId = data?.uploadId;
+			assert.ok(uploadId);
+
+			const formDataRest = new FormData();
+			formDataRest.set('type', 'rest-chunk');
+			formDataRest.set('uploadId', 'fake');
+			formDataRest.set('image', new File([''], 'chunk-2.png'));
+			const reqRest = new Request('http://example.com/_actions/imageUploadInChunks', {
+				method: 'POST',
+				body: formDataRest,
+			});
+
+			const resRest = await app.render(reqRest);
+			assert.equal(resRest.status, 200);
+			assert.equal(resRest.headers.get('Content-Type'), 'application/json+devalue');
+			const dataRest = devalue.parse(await resRest.text());
+			assert.equal('fake', dataRest?.uploadId);
+		});
 	});
 });
 
diff --git a/packages/astro/test/fixtures/actions/src/actions/index.ts b/packages/astro/test/fixtures/actions/src/actions/index.ts
index ed7692799..4e6120309 100644
--- a/packages/astro/test/fixtures/actions/src/actions/index.ts
+++ b/packages/astro/test/fixtures/actions/src/actions/index.ts
@@ -7,6 +7,29 @@ const passwordSchema = z
 	.max(128, 'Password length exceeded. Max 128 chars.');
 
 export const server = {
+	imageUploadInChunks: defineAction({
+		accept: 'form',
+		input: z.discriminatedUnion('type', [
+			z.object({
+				type: z.literal('first-chunk'),
+				image: z.instanceof(File),
+				alt: z.string(),
+			}),
+			z.object({ type: z.literal('rest-chunk'), image: z.instanceof(File), uploadId: z.string() }),
+		]),
+		handler: async (data) => {
+			if (data.type === 'first-chunk') {
+				const uploadId = Math.random().toString(36).slice(2);
+				return {
+					uploadId,
+				};
+			} else {
+				return {
+					uploadId: data.uploadId,
+				};
+			}
+		},
+	}),
 	subscribe: defineAction({
 		input: z.object({ channel: z.string() }),
 		handler: async ({ channel }) => {