From d22566702b2dff8f74e55a77621486071e029012 Mon Sep 17 00:00:00 2001
From: Richard Yu <richard95yu@gmail.com>
Date: Fri, 5 Dec 2025 13:49:15 -0800
Subject: [PATCH] Address CodeRabbit feedback: add Zod validation to
 extractWorkflow
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replace unsafe type casts with Zod schema validation for the workflow
container structure. Full workflow validation still happens downstream
via validateComfyWorkflow.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 .../remote/comfyui/jobs/fetchers/fetchJobs.ts | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/platform/remote/comfyui/jobs/fetchers/fetchJobs.ts b/src/platform/remote/comfyui/jobs/fetchers/fetchJobs.ts
index 2b6cde8c9..2d4813303 100644
--- a/src/platform/remote/comfyui/jobs/fetchers/fetchJobs.ts
+++ b/src/platform/remote/comfyui/jobs/fetchers/fetchJobs.ts
@@ -6,6 +6,8 @@
  * All distributions use the /jobs endpoint.
  */
 
+import { z } from 'zod'
+
 import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
 import type { PromptId } from '@/schemas/apiSchema'
 
@@ -143,6 +145,22 @@ export async function fetchJobDetail(
   }
 }
 
+/**
+ * Schema for workflow container structure.
+ * Full workflow validation happens downstream via validateComfyWorkflow.
+ */
+const zWorkflowContainer = z.object({
+  extra_data: z
+    .object({
+      extra_pnginfo: z
+        .object({
+          workflow: z.unknown()
+        })
+        .optional()
+    })
+    .optional()
+})
+
 /**
  * Extracts workflow from job detail response.
  * The workflow is nested at: workflow.extra_data.extra_pnginfo.workflow
@@ -150,11 +168,10 @@ export async function fetchJobDetail(
 export function extractWorkflow(
   job: JobDetail | undefined
 ): ComfyWorkflowJSON | undefined {
-  // Cast is safe - workflow will be validated by loadGraphData -> validateComfyWorkflow
-  const workflowData = job?.workflow as
-    | { extra_data?: { extra_pnginfo?: { workflow?: unknown } } }
-    | undefined
-  return workflowData?.extra_data?.extra_pnginfo?.workflow as
+  const parsed = zWorkflowContainer.safeParse(job?.workflow)
+  if (!parsed.success) return undefined
+  // Full workflow validation happens downstream via validateComfyWorkflow
+  return parsed.data.extra_data?.extra_pnginfo?.workflow as
     | ComfyWorkflowJSON
     | undefined
 }