From fc2a5a42ce66a5cd3209d118c0bc735ad86d9075 Mon Sep 17 00:00:00 2001
From: Allen Hutchison <adh@google.com>
Date: Sun, 14 Dec 2025 19:54:28 -0800
Subject: [PATCH] fix: use zod for safety check result validation (#15026)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 packages/core/src/safety/checker-runner.ts | 30 ++++++++++++++--------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/packages/core/src/safety/checker-runner.ts b/packages/core/src/safety/checker-runner.ts
index e9748cd5c3..02f824d980 100644
--- a/packages/core/src/safety/checker-runner.ts
+++ b/packages/core/src/safety/checker-runner.ts
@@ -15,6 +15,23 @@ import type { SafetyCheckInput, SafetyCheckResult } from './protocol.js';
 import { SafetyCheckDecision } from './protocol.js';
 import type { CheckerRegistry } from './registry.js';
 import type { ContextBuilder } from './context-builder.js';
+import { z } from 'zod';
+
+const SafetyCheckResultSchema: z.ZodType<SafetyCheckResult> =
+  z.discriminatedUnion('decision', [
+    z.object({
+      decision: z.literal(SafetyCheckDecision.ALLOW),
+      reason: z.string().optional(),
+    }),
+    z.object({
+      decision: z.literal(SafetyCheckDecision.DENY),
+      reason: z.string().min(1),
+    }),
+    z.object({
+      decision: z.literal(SafetyCheckDecision.ASK_USER),
+      reason: z.string().min(1),
+    }),
+  ]);
 
 /**
  * Configuration for the checker runner.
@@ -212,17 +229,8 @@ export class CheckerRunner {
 
         // Try to parse the output
         try {
-          const result: SafetyCheckResult = JSON.parse(stdout);
-
-          // Validate the result structure
-          if (
-            !result.decision ||
-            !Object.values(SafetyCheckDecision).includes(result.decision)
-          ) {
-            throw new Error(
-              'Invalid result: missing or invalid "decision" field',
-            );
-          }
+          const rawResult = JSON.parse(stdout);
+          const result = SafetyCheckResultSchema.parse(rawResult);
 
           resolve(result);
         } catch (parseError) {