From 2b08e2d83006180787ff4c04e3e1f115731c84a4 Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Fri, 31 Jan 2025 19:19:03 -0800
Subject: [PATCH] Made reliable type passdown from executor to console

---
 blocks/index.ts                         |  2 --
 executor/index.ts                       |  3 +--
 serializer/__tests__/serializer.test.ts | 18 ++++--------------
 serializer/index.ts                     |  9 +++++----
 serializer/types.ts                     |  1 +
 5 files changed, 11 insertions(+), 22 deletions(-)

diff --git a/blocks/index.ts b/blocks/index.ts
index 3fbb08098b..f77948c507 100644
--- a/blocks/index.ts
+++ b/blocks/index.ts
@@ -43,8 +43,6 @@ const toolToBlockType = Object.entries(blocks).reduce((acc, [blockType, config])
 export const getBlock = (type: string): BlockConfig | undefined =>
   blocks[type]
 
-export const getBlockTypeForTool = (toolId: string): string | undefined =>
-  toolToBlockType[toolId]
 
 export const getBlocksByCategory = (category: 'basic' | 'advanced'): BlockConfig[] =>
   Object.values(blocks).filter(block => block.toolbar.category === category)
diff --git a/executor/index.ts b/executor/index.ts
index 1edff77c15..c28d296917 100644
--- a/executor/index.ts
+++ b/executor/index.ts
@@ -18,7 +18,6 @@ import {
   BlockLog
 } from './types'
 import { tools } from '@/tools'
-import { getBlockTypeForTool } from '@/blocks'
 
 export class Executor {
   constructor(
@@ -126,7 +125,7 @@ export class Executor {
           const blockLog: Partial<BlockLog> = {
             blockId: block.id,
             blockTitle: block.metadata?.title,
-            blockType: getBlockTypeForTool(block.config.tool || ''),
+            blockType: block.metadata?.type,
             startedAt: new Date().toISOString(),
           }
 
diff --git a/serializer/__tests__/serializer.test.ts b/serializer/__tests__/serializer.test.ts
index 3c7896a20b..1e537c0e07 100644
--- a/serializer/__tests__/serializer.test.ts
+++ b/serializer/__tests__/serializer.test.ts
@@ -15,18 +15,6 @@ jest.mock('@/components/icons', () => ({
 // Mock blocks
 jest.mock('@/blocks', () => ({
   getBlock: jest.fn(),
-  getBlockTypeForTool: jest.fn((toolId: string) => {
-    switch (toolId) {
-      case 'openai.chat':
-        return 'agent'
-      case 'http.request':
-        return 'api'
-      case 'test-tool':
-        return 'agent'
-      default:
-        return undefined
-    }
-  })
 })) 
 
 describe('Serializer', () => {
@@ -372,8 +360,10 @@ describe('Serializer', () => {
               title: 'Agent Block',
               description: 'Use any LLM',
               category: 'basic',
-              color: '#7F2FFF'
-            }
+              color: '#7F2FFF',
+              type: 'agent'
+            },
+            enabled: true
           }
         ],
         connections: []
diff --git a/serializer/index.ts b/serializer/index.ts
index 0869286e7b..611576c32c 100644
--- a/serializer/index.ts
+++ b/serializer/index.ts
@@ -1,7 +1,7 @@
 import { BlockState, SubBlockState } from '@/stores/workflow/types' 
 import { Edge } from 'reactflow' 
 import { SerializedBlock, SerializedConnection, SerializedWorkflow } from './types' 
-import { getBlock, getBlockTypeForTool } from '@/blocks' 
+import { getBlock } from '@/blocks' 
 
 export class Serializer {
   serializeWorkflow(blocks: Record<string, BlockState>, edges: Edge[]): SerializedWorkflow {
@@ -52,7 +52,8 @@ export class Serializer {
         title: block.name,
         description: blockConfig.toolbar.description,
         category: blockConfig.toolbar.category,
-        color: blockConfig.toolbar.bgColor
+        color: blockConfig.toolbar.bgColor,
+        type: block.type
       },
       enabled: block.enabled
     } 
@@ -108,9 +109,9 @@ export class Serializer {
   }
 
   private deserializeBlock(serializedBlock: SerializedBlock): BlockState {
-    const blockType = getBlockTypeForTool(serializedBlock.config.tool)
+    const blockType = serializedBlock.metadata?.type
     if (!blockType) {
-      throw new Error(`Invalid tool ID: ${serializedBlock.config.tool}`)
+      throw new Error(`Invalid block type: ${serializedBlock.metadata?.type}`)
     }
 
     const blockConfig = getBlock(blockType)
diff --git a/serializer/types.ts b/serializer/types.ts
index 52aab8e4ec..968eca24be 100644
--- a/serializer/types.ts
+++ b/serializer/types.ts
@@ -29,6 +29,7 @@ export interface SerializedBlock {
     category?: string 
     icon?: string 
     color?: string 
+    type: string 
   } 
   enabled: boolean
 }