diff --git a/server/src/workflow-management/classes/Generator.ts b/server/src/workflow-management/classes/Generator.ts
index e2faf92c..9368b2f8 100644
--- a/server/src/workflow-management/classes/Generator.ts
+++ b/server/src/workflow-management/classes/Generator.ts
@@ -357,5 +357,15 @@ export class WorkflowGenerator {
     }
   };
 
+  /**
+   * Updates the socket used for communication with the client.
+   * @param socket The socket to be used for communication.
+   * @returns void
+   */
+  public updateSocket = (socket: Socket) : void => {
+    this.socket = socket;
+    this.registerEventHandlers(socket);
+  };
+
   
 }