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((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; }