Fix onConnect handler (#7265)

Signed-off-by: Denis Bykhov <[email protected]>

packages/core/src/client.ts

@@ -85,7 +85,7 @@ export interface ClientConnection extends Storage, FulltextStorage, BackupClient
   isConnected: () => boolean
 
   close: () => Promise<void>
-  onConnect?: (event: ClientConnectEvent) => Promise<void>
+  onConnect?: (event: ClientConnectEvent, data: any) => Promise<void>
 
   // If hash is passed, will return LoadModelResponse
   loadModel: (last: Timestamp, hash?: string) => Promise<Tx[] | LoadModelResponse>
@@ -264,11 +264,11 @@ export async function createClient (
   txHandler(...txBuffer)
   txBuffer = undefined
 
-  const oldOnConnect: ((event: ClientConnectEvent) => Promise<void>) | undefined = conn.onConnect
-  conn.onConnect = async (event) => {
+  const oldOnConnect: ((event: ClientConnectEvent, data: any) => Promise<void>) | undefined = conn.onConnect
+  conn.onConnect = async (event, data) => {
     console.log('Client: onConnect', event)
     if (event === ClientConnectEvent.Maintenance) {
-      await oldOnConnect?.(ClientConnectEvent.Maintenance)
+      await oldOnConnect?.(ClientConnectEvent.Maintenance, data)
       return
     }
     // Find all new transactions and apply
@@ -282,15 +282,15 @@ export async function createClient (
       model = new ModelDb(hierarchy)
 
       await ctx.with('build-model', {}, (ctx) => buildModel(ctx, loadModelResponse, modelFilter, hierarchy, model))
-      await oldOnConnect?.(ClientConnectEvent.Upgraded)
+      await oldOnConnect?.(ClientConnectEvent.Upgraded, data)
 
       // No need to fetch more stuff since upgrade was happened.
       return
     }
 
     if (event === ClientConnectEvent.Connected) {
       // No need to do anything here since we connected.
-      await oldOnConnect?.(event)
+      await oldOnConnect?.(event, data)
       return
     }
 
@@ -318,10 +318,10 @@ export async function createClient (
     if (atxes.length < transactionThreshold && !needFullRefresh) {
       console.log('applying input transactions', atxes.length)
       txHandler(...atxes)
-      await oldOnConnect?.(ClientConnectEvent.Reconnected)
+      await oldOnConnect?.(ClientConnectEvent.Reconnected, data)
     } else {
       // We need to trigger full refresh on queries, etc.
-      await oldOnConnect?.(ClientConnectEvent.Refresh)
+      await oldOnConnect?.(ClientConnectEvent.Refresh, data)
     }
   }
 

plugins/client-resources/src/connection.ts

@@ -101,6 +101,8 @@ class Connection implements ClientConnection {
 
   private helloRecieved: boolean = false
 
+  onConnect?: (event: ClientConnectEvent, data: any) => Promise<void>
+
   rpcHandler = new RPCHandler()
 
   constructor (
@@ -130,6 +132,8 @@ class Connection implements ClientConnection {
       this.sessionId = generateId()
     }
 
+    this.onConnect = opt?.onConnect
+
     this.scheduleOpen(this.ctx, false)
   }
 
@@ -261,7 +265,7 @@ class Connection implements ClientConnection {
     if (resp.id === -1) {
       this.delay = 0
       if (resp.result?.state === 'upgrading') {
-        void this.opt?.onConnect?.(ClientConnectEvent.Maintenance, resp.result.stats)
+        void this.onConnect?.(ClientConnectEvent.Maintenance, resp.result.stats)
         this.upgrading = true
         this.delay = 3
         return
@@ -302,7 +306,7 @@ class Connection implements ClientConnection {
           v.reconnect?.()
         }
 
-        void this.opt?.onConnect?.(
+        void this.onConnect?.(
           (resp as HelloResponse).reconnect === true ? ClientConnectEvent.Reconnected : ClientConnectEvent.Connected,
           this.sessionId
         )

plugins/client-resources/src/index.ts

@@ -118,14 +118,14 @@ export default async () => {
                   reject(new Error(`Connection timeout, and no connection established to ${endpoint}`))
                 }
               }, connectTimeout)
-              newOpt.onConnect = (event) => {
+              newOpt.onConnect = async (event, data) => {
                 // Any event is fine, it means server is alive.
                 clearTimeout(connectTO)
                 resolve()
               }
             })
           }
-          const clientConnection = connect(url, upgradeHandler, tokenPayload.workspace, tokenPayload.email, newOpt)
+          const clientConnection = connect(url, upgradeHandler, tokenPayload.workspace, tokenPayload.email, opt)
           if (connectPromise !== undefined) {
             await connectPromise
           }

plugins/client/src/index.ts

@@ -62,7 +62,7 @@ export interface ClientFactoryOptions {
   onUpgrade?: () => void
   onUnauthorized?: () => void
   onArchived?: () => void
-  onConnect?: (event: ClientConnectEvent, data: any) => void
+  onConnect?: (event: ClientConnectEvent, data: any) => Promise<void>
   ctx?: MeasureContext
   onDialTimeout?: () => void | Promise<void>
 }

plugins/guest-resources/src/connect.ts

@@ -125,7 +125,7 @@ export async function connect (title: string): Promise<Client | undefined> {
       })
     },
     // We need to refresh all active live queries and clear old queries.
-    onConnect: (event: ClientConnectEvent, data: any) => {
+    onConnect: async (event: ClientConnectEvent, data: any) => {
       console.log('WorkbenchClient: onConnect', event)
       try {
         if (event === ClientConnectEvent.Connected) {

plugins/workbench-resources/src/connect.ts

@@ -212,7 +212,7 @@ export async function connect (title: string): Promise<Client | undefined> {
           })
         },
         // We need to refresh all active live queries and clear old queries.
-        onConnect: (event: ClientConnectEvent, data: any) => {
+        onConnect: async (event: ClientConnectEvent, data: any): Promise<void> => {
           console.log('WorkbenchClient: onConnect', event)
           if (event === ClientConnectEvent.Maintenance) {
             if (data != null && data.total !== 0) {

services/github/pod-github/src/client.ts

@@ -18,7 +18,7 @@ import config from './config'
 export async function createPlatformClient (
   workspace: string,
   timeout: number,
-  reconnect?: (event: ClientConnectEvent) => void
+  reconnect?: (event: ClientConnectEvent, data: any) => Promise<void>
 ): Promise<Client> {
   setMetadata(client.metadata.ClientSocketFactory, (url) => {
     return new WebSocket(url, {

services/github/pod-github/src/worker.ts

@@ -1560,7 +1560,7 @@ export class GithubWorker implements IntegrationManager {
     ctx.info('Connecting to', { workspace: workspace.workspaceUrl, workspaceId: workspace.workspaceName })
     let client: Client | undefined
     try {
-      client = await createPlatformClient(workspace.name, 30000, (event: ClientConnectEvent) => {
+      client = await createPlatformClient(workspace.name, 30000, async (event: ClientConnectEvent) => {
         reconnect(workspace.name, event)
       })