From b2dba3981075fb1da5b6c853b52a250a77a8aa83 Mon Sep 17 00:00:00 2001
From: Reinier van der Leer <pwuts@agpt.co>
Date: Fri, 26 Jul 2024 13:26:11 +0200
Subject: [PATCH] fix(builder/api-client): Fix websocket client types (#7603)

- Fix type signatures of `sendWebSocketMessage(..)`, `onWebSocketMessage(..)`, `runGraph(..)` in `autogpt-server-api/client`
  - Add `WebsocketMessageTypeMap`

- Fix type signature of `updateNodesWithExecutionData` in `FlowEditor`
---
 rnd/autogpt_builder/src/components/Flow.tsx      |  4 ++--
 .../src/lib/autogpt-server-api/client.ts         | 16 +++++++++++++---
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/rnd/autogpt_builder/src/components/Flow.tsx b/rnd/autogpt_builder/src/components/Flow.tsx
index 730d14298..95f792cca 100644
--- a/rnd/autogpt_builder/src/components/Flow.tsx
+++ b/rnd/autogpt_builder/src/components/Flow.tsx
@@ -15,7 +15,7 @@ import ReactFlow, {
 import 'reactflow/dist/style.css';
 import CustomNode from './CustomNode';
 import './flow.css';
-import AutoGPTServerAPI, { Block, Graph, ObjectSchema } from '@/lib/autogpt-server-api';
+import AutoGPTServerAPI, { Block, Graph, NodeExecutionResult, ObjectSchema } from '@/lib/autogpt-server-api';
 import { Button } from './ui/button';
 import { Input } from './ui/input';
 import { ChevronRight, ChevronLeft } from "lucide-react";
@@ -440,7 +440,7 @@ const FlowEditor: React.FC<{
 
 
 
-  const updateNodesWithExecutionData = (executionData: any[]) => {
+  const updateNodesWithExecutionData = (executionData: NodeExecutionResult[]) => {
     setNodes((nds) =>
       nds.map((node) => {
         const nodeExecution = executionData.find((exec) => exec.node_id === node.data.backend_id);
diff --git a/rnd/autogpt_builder/src/lib/autogpt-server-api/client.ts b/rnd/autogpt_builder/src/lib/autogpt-server-api/client.ts
index 61b920a0c..6901326ab 100644
--- a/rnd/autogpt_builder/src/lib/autogpt-server-api/client.ts
+++ b/rnd/autogpt_builder/src/lib/autogpt-server-api/client.ts
@@ -182,7 +182,9 @@ export default class AutoGPTServerAPI {
     }
   }
 
-  sendWebSocketMessage(method: string, data: any) {
+  sendWebSocketMessage<M extends keyof WebsocketMessageTypeMap>(
+    method: M, data: WebsocketMessageTypeMap[M]
+  ) {
     if (this.socket && this.socket.readyState === WebSocket.OPEN) {
       this.socket.send(JSON.stringify({ method, data }));
     } else {
@@ -190,7 +192,9 @@ export default class AutoGPTServerAPI {
     }
   }
 
-  onWebSocketMessage(method: string, handler: (data: any) => void) {
+  onWebSocketMessage<M extends keyof WebsocketMessageTypeMap>(
+    method: M, handler: (data: WebsocketMessageTypeMap[M]) => void
+  ) {
     this.messageHandlers[method] = handler;
   }
 
@@ -198,7 +202,7 @@ export default class AutoGPTServerAPI {
     this.sendWebSocketMessage('subscribe', { graph_id: graphId });
   }
 
-  runGraph(graphId: string, data: any = {}) {
+  runGraph(graphId: string, data: WebsocketMessageTypeMap["run_graph"]["data"] = {}) {
     this.sendWebSocketMessage('run_graph', { graph_id: graphId, data });
   }
 }
@@ -212,3 +216,9 @@ type GraphCreateRequestBody = {
 } | {
   graph: GraphCreatable;
 }
+
+type WebsocketMessageTypeMap = {
+  subscribe: { graph_id: string; };
+  run_graph: { graph_id: string; data: { [key: string]: any }; };
+  execution_event: NodeExecutionResult;
+}