From 71c6d62c7e2671c4cc745b8d00c8e9c72bdc8108 Mon Sep 17 00:00:00 2001
From: Richard Yu <richard95yu@gmail.com>
Date: Mon, 12 Jan 2026 17:56:04 -0800
Subject: [PATCH] refactor: use Zod validation for extractWorkflow instead of
 type assertion
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Address review comment about type assertion. Now uses validateComfyWorkflow
with safeParse to properly validate the workflow schema instead of casting.
Logs validation failures for debugging while returning undefined for graceful
degradation.

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

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 .../remote/comfyui/jobs/fetchJobs.test.ts     | 47 +++++++++++++++----
 src/platform/remote/comfyui/jobs/fetchJobs.ts | 24 +++++++---
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/platform/remote/comfyui/jobs/fetchJobs.test.ts b/src/platform/remote/comfyui/jobs/fetchJobs.test.ts
index f3b2ad5a8..56f7fdef6 100644
--- a/src/platform/remote/comfyui/jobs/fetchJobs.test.ts
+++ b/src/platform/remote/comfyui/jobs/fetchJobs.test.ts
@@ -257,35 +257,66 @@ describe('fetchJobs', () => {
   })
 
   describe('extractWorkflow', () => {
-    it('extracts workflow from nested structure', () => {
+    const validWorkflow = {
+      version: 0.4,
+      last_node_id: 1,
+      last_link_id: 0,
+      nodes: [],
+      links: []
+    }
+
+    it('extracts and validates workflow from nested structure', async () => {
       const jobDetail = {
         ...createMockJob('job1', 'completed'),
         workflow: {
           extra_data: {
             extra_pnginfo: {
-              workflow: { nodes: [], links: [] }
+              workflow: validWorkflow
             }
           }
         }
       }
 
-      const workflow = extractWorkflow(jobDetail)
+      const workflow = await extractWorkflow(jobDetail)
 
-      expect(workflow).toEqual({ nodes: [], links: [] })
+      expect(workflow).toEqual(validWorkflow)
     })
 
-    it('returns undefined if workflow not present', () => {
+    it('returns undefined if workflow not present', async () => {
       const jobDetail = createMockJob('job1', 'completed')
 
-      const workflow = extractWorkflow(jobDetail)
+      const workflow = await extractWorkflow(jobDetail)
 
       expect(workflow).toBeUndefined()
     })
 
-    it('returns undefined for undefined input', () => {
-      const workflow = extractWorkflow(undefined)
+    it('returns undefined for undefined input', async () => {
+      const workflow = await extractWorkflow(undefined)
 
       expect(workflow).toBeUndefined()
     })
+
+    it('returns undefined for invalid workflow and logs warning', async () => {
+      const consoleSpy = vi.spyOn(console, 'warn').mockImplementation(() => {})
+      const jobDetail = {
+        ...createMockJob('job1', 'completed'),
+        workflow: {
+          extra_data: {
+            extra_pnginfo: {
+              workflow: { invalid: 'data' }
+            }
+          }
+        }
+      }
+
+      const workflow = await extractWorkflow(jobDetail)
+
+      expect(workflow).toBeUndefined()
+      expect(consoleSpy).toHaveBeenCalledWith(
+        '[extractWorkflow] Workflow validation failed:',
+        expect.any(String)
+      )
+      consoleSpy.mockRestore()
+    })
   })
 })
diff --git a/src/platform/remote/comfyui/jobs/fetchJobs.ts b/src/platform/remote/comfyui/jobs/fetchJobs.ts
index 41d5aeaa9..f5facceff 100644
--- a/src/platform/remote/comfyui/jobs/fetchJobs.ts
+++ b/src/platform/remote/comfyui/jobs/fetchJobs.ts
@@ -7,6 +7,7 @@
  */
 
 import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
+import { validateComfyWorkflow } from '@/platform/workflow/validation/schemas/workflowSchema'
 import type { PromptId } from '@/schemas/apiSchema'
 
 import type {
@@ -136,16 +137,25 @@ export async function fetchJobDetail(
 }
 
 /**
- * Extracts workflow from job detail response.
+ * Extracts and validates workflow from job detail response.
  * The workflow is nested at: workflow.extra_data.extra_pnginfo.workflow
- * Full workflow validation happens downstream via validateComfyWorkflow.
+ *
+ * Uses Zod validation via validateComfyWorkflow to ensure the workflow
+ * conforms to the expected schema. Logs validation failures for debugging
+ * but still returns undefined to allow graceful degradation.
  */
-export function extractWorkflow(
+export async function extractWorkflow(
   job: JobDetail | undefined
-): ComfyWorkflowJSON | undefined {
+): Promise<ComfyWorkflowJSON | undefined> {
   const parsed = zWorkflowContainer.safeParse(job?.workflow)
   if (!parsed.success) return undefined
-  return parsed.data.extra_data?.extra_pnginfo?.workflow as
-    | ComfyWorkflowJSON
-    | undefined
+
+  const rawWorkflow = parsed.data.extra_data?.extra_pnginfo?.workflow
+  if (!rawWorkflow) return undefined
+
+  const validated = await validateComfyWorkflow(rawWorkflow, (error) => {
+    console.warn('[extractWorkflow] Workflow validation failed:', error)
+  })
+
+  return validated ?? undefined
 }