From b2ee64c4342b948fa5e85da60bba84f48c0a8791 Mon Sep 17 00:00:00 2001
From: Ajay R <ajayr@formbuddy.com>
Date: Wed, 11 Mar 2026 09:11:30 +0000
Subject: [PATCH] fix(gateway): address health snapshot review feedback

When a caller provides a runtimeSnapshot while a refresh is already
in-flight, await the current refresh and return the follow-up result
instead of returning stale data from the in-flight promise. This ensures
callers that supply fresh runtime state never receive a snapshot that
omits their updates.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/gateway/server/health-state.test.ts | 33 +++++++++++++++++++++++++
 src/gateway/server/health-state.ts      | 12 +++++++++
 2 files changed, 45 insertions(+)

diff --git a/src/gateway/server/health-state.test.ts b/src/gateway/server/health-state.test.ts
index 72fad4b7b54..5a0526b6157 100644
--- a/src/gateway/server/health-state.test.ts
+++ b/src/gateway/server/health-state.test.ts
@@ -150,4 +150,37 @@ describe("refreshGatewayHealthSnapshot", () => {
     // Only one call; no pending snapshot.
     expect(mockGetHealthSnapshot).toHaveBeenCalledTimes(1);
   });
+
+  it("returns the follow-up result (not stale) when runtimeSnapshot is provided during flight", async () => {
+    const snap1 = makeRuntimeSnapshot(false);
+    const snap2 = makeRuntimeSnapshot(true);
+
+    const staleResult = makeHealthSummary({ ts: 1 });
+    const freshResult = makeHealthSummary({ ts: 2 });
+
+    let resolveFirst!: (v: HealthSummary) => void;
+    const firstDone = new Promise<HealthSummary>((r) => {
+      resolveFirst = r;
+    });
+
+    mockGetHealthSnapshot
+      .mockImplementationOnce(() => firstDone)
+      .mockResolvedValueOnce(freshResult);
+
+    // Start first refresh (in-flight).
+    const p1 = refreshGatewayHealthSnapshot({ runtimeSnapshot: snap1 });
+
+    // Second call arrives with a fresher snapshot while first is in-flight.
+    const p2 = refreshGatewayHealthSnapshot({ runtimeSnapshot: snap2 });
+
+    // Complete the first (stale) refresh.
+    resolveFirst(staleResult);
+
+    const [r1, r2] = await Promise.all([p1, p2]);
+
+    // First caller gets the stale result (it started the refresh).
+    expect(r1.ts).toBe(1);
+    // Second caller must get the fresh follow-up result, not the stale one.
+    expect(r2.ts).toBe(2);
+  });
 });
diff --git a/src/gateway/server/health-state.ts b/src/gateway/server/health-state.ts
index b5497e39ea5..0ab3325ffb6 100644
--- a/src/gateway/server/health-state.ts
+++ b/src/gateway/server/health-state.ts
@@ -115,6 +115,18 @@ export async function refreshGatewayHealthSnapshot(opts?: {
         void refreshGatewayHealthSnapshot({ probe: opts?.probe });
       }
     });
+  } else if (opts?.runtimeSnapshot !== undefined) {
+    // Caller provided fresh runtime data that the in-flight refresh won't include.
+    // Wait for the current refresh to complete (its finally block will kick off a
+    // follow-up using pendingRuntimeSnapshot), then return the follow-up result so
+    // the caller receives a snapshot that includes the latest runtime state.
+    await healthRefresh;
+    // After the finally block runs, healthRefresh is either the follow-up promise
+    // (if pendingRuntimeSnapshot was set) or null.
+    if (healthRefresh) {
+      return healthRefresh;
+    }
+    return healthCache!;
   }
   return healthRefresh;
 }