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 = + 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) {