From 9dc961ca42d10923c0b7d27d4dbd15bc83bb7f00 Mon Sep 17 00:00:00 2001
From: platfowner <seo@comcom.ai>
Date: Tue, 13 Aug 2024 10:22:37 +0900
Subject: [PATCH] Add custom client id parameter to connect()

---
 src/ainize.ts           | 16 +++++++++++-----
 src/handlers/handler.ts |  8 ++++----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/ainize.ts b/src/ainize.ts
index d85a180..24d19bc 100644
--- a/src/ainize.ts
+++ b/src/ainize.ts
@@ -36,21 +36,27 @@ export default class Ainize {
 
   /**
    * Login to ainize using AI Network account private key.
-   * @param {string} privateKey 
+   * @param {string} privateKey The private key to initialize the AinModule with.
+   * @param {ConnectionCallback} ConnectionCallback The connection callback function.
+   * @param {DisconnectionCallback} disconnectionCallback The disconnection callback function.
+   * @param {string} customClientId The custom client id to set.
    */
-  async login(privateKey: string, connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback) {
+  async login(privateKey: string, connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback, customClientId?: string) {
     this.ain.setDefaultAccount(privateKey);
-    await this.handler.connect(connectionCb, disconnectionCb);
+    await this.handler.connect(connectionCb, disconnectionCb, customClientId);
     console.log('login success! address:', await this.ain.getAddress());
   }
 
   /**
    * Login to ainize using AIN Wallet Signer.
+   * @param {ConnectionCallback} ConnectionCallback The connection callback function.
+   * @param {DisconnectionCallback} disconnectionCallback The disconnection callback function.
+   * @param {string} customClientId The custom client id to set.
    */
-  async loginWithSigner(connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback) {
+  async loginWithSigner(connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback, customClientId?: string) {
     const signer = new AinWalletSigner;
     this.ain.setSigner(signer);
-    await this.handler.connect(connectionCb, disconnectionCb);
+    await this.handler.connect(connectionCb, disconnectionCb, customClientId);
     console.log('login success! address: ', await this.ain.getAddress());
   }
 
diff --git a/src/handlers/handler.ts b/src/handlers/handler.ts
index 6724c21..bbde682 100644
--- a/src/handlers/handler.ts
+++ b/src/handlers/handler.ts
@@ -24,9 +24,9 @@ export default class Handler {
     return this.ain.getEventManager().isConnected();
   }
 
-  async connect(connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback) {
+  async connect(connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback, customClientId?: string) {
     this.checkEventManager();
-    await this.ain.getEventManager().connect(connectionCb, this.connectionRetryCb.bind(this, connectionCb, disconnectionCb));
+    await this.ain.getEventManager().connect(connectionCb, this.connectionRetryCb.bind(this, connectionCb, disconnectionCb, customClientId));
     console.log('connected');
   };
   
@@ -36,7 +36,7 @@ export default class Handler {
     console.log('Disconnected');
   }
 
-  private async connectionRetryCb(connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback, webSocket?: any) {
+  private async connectionRetryCb(connectionCb?: ConnectionCallback, disconnectionCb?: DisconnectionCallback, customClientId?: string, webSocket?: any) {
     try {
       if (disconnectionCb) {
         disconnectionCb(webSocket);
@@ -44,7 +44,7 @@ export default class Handler {
       const address = await AinModule.getInstance().getAddress();
       if (address) {
         console.log('Disconnected. Reconnecting...');
-        await this.connect(connectionCb, disconnectionCb);
+        await this.connect(connectionCb, disconnectionCb, customClientId);
       }
     } catch (_) {
       return;