From 5b17237f0dffd271cc2292615a16e2d8a54c144e Mon Sep 17 00:00:00 2001
From: Levente Kiss <levente.kiss@solarpunk.buzz>
Date: Fri, 31 Oct 2025 18:46:10 +1300
Subject: [PATCH] fix: add active query tracking to QueryOnConnect and await
 its stop()

---
 .../src/query_on_connect/query_on_connect.ts    | 17 +++++++++++++++--
 .../src/reliable_channel/reliable_channel.ts    |  2 +-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/packages/sdk/src/query_on_connect/query_on_connect.ts b/packages/sdk/src/query_on_connect/query_on_connect.ts
index 2fc16e9a4b..7d1c4dddaf 100644
--- a/packages/sdk/src/query_on_connect/query_on_connect.ts
+++ b/packages/sdk/src/query_on_connect/query_on_connect.ts
@@ -54,6 +54,7 @@ export class QueryOnConnect<
 
   private isRunning: boolean = false;
   private abortController?: AbortController;
+  private activeQueryPromise?: Promise<void>;
 
   private boundStoreConnectHandler?: (event: CustomEvent<PeerId>) => void;
   private boundHealthHandler?: (event: CustomEvent<HealthStatus>) => void;
@@ -87,7 +88,7 @@ export class QueryOnConnect<
     this.setupEventListeners();
   }
 
-  public stop(): void {
+  public async stop(): Promise<void> {
     if (!this.isRunning) {
       return;
     }
@@ -99,6 +100,15 @@ export class QueryOnConnect<
       this.abortController = undefined;
     }
 
+    if (this.activeQueryPromise) {
+      log.info("Waiting for active query to complete...");
+      try {
+        await this.activeQueryPromise;
+      } catch (error) {
+        log.warn("Active query failed during stop:", error);
+      }
+    }
+
     this.unsetEventListeners();
   }
 
@@ -130,7 +140,10 @@ export class QueryOnConnect<
       this.lastTimeOffline > this.lastSuccessfulQuery ||
       timeSinceLastQuery > this.forceQueryThresholdMs
     ) {
-      await this.query(peerId);
+      this.activeQueryPromise = this.query(peerId).finally(() => {
+        this.activeQueryPromise = undefined;
+      });
+      await this.activeQueryPromise;
     } else {
       log.info(`no querying`);
     }
diff --git a/packages/sdk/src/reliable_channel/reliable_channel.ts b/packages/sdk/src/reliable_channel/reliable_channel.ts
index 3ed2d6d549..a92567b570 100644
--- a/packages/sdk/src/reliable_channel/reliable_channel.ts
+++ b/packages/sdk/src/reliable_channel/reliable_channel.ts
@@ -534,7 +534,7 @@ export class ReliableChannel<
     }
 
     if (this.queryOnConnect) {
-      this.queryOnConnect.stop();
+      await this.queryOnConnect.stop();
     }
 
     if (this.retryManager) {