From f98043c21bc225b7603a22629f5bc744817add25 Mon Sep 17 00:00:00 2001 From: Bryce Zuccaro Date: Sat, 15 Aug 2026 14:26:28 -0600 Subject: [PATCH] fix: self-heal in-flight counts by pruning stale request ids On a long-lived SSE connection a missed remove leaves a stale request id in the tracker forever; llama-swap only sends a fresh snapshot on connect, so counts drifted up and stayed stuck. The tracker now timestamps each request id (upserts refresh it) and prunes ids not updated within 120s, on every inflight event and on a 30s runtime interval. Removes also accept numeric ids defensively. --- src/lib/event-feed.ts | 2 +- src/lib/inflight-tracker.ts | 18 +++++++++++++----- src/lib/runtime.ts | 6 ++++++ tests/event-feed.test.ts | 9 +++++++++ tests/inflight-tracker.test.ts | 31 +++++++++++++++++++++++++++++++ 5 files changed, 60 insertions(+), 6 deletions(-) diff --git a/src/lib/event-feed.ts b/src/lib/event-feed.ts index 1bebbf2..f2644a4 100644 --- a/src/lib/event-feed.ts +++ b/src/lib/event-feed.ts @@ -17,7 +17,7 @@ export function decodeEvent(msg: SseMessage): FeedEvent | null { if (outer.type === "inflight") { const operation = inner.operation as string; if (operation === "remove") { - const id = typeof inner.id === "string" ? inner.id : undefined; + const id = typeof inner.id === "string" ? inner.id : typeof inner.id === "number" ? String(inner.id) : undefined; if (id) return { type: "inflight", operation: "remove", id }; return null; } diff --git a/src/lib/inflight-tracker.ts b/src/lib/inflight-tracker.ts index 5406626..0732a6c 100644 --- a/src/lib/inflight-tracker.ts +++ b/src/lib/inflight-tracker.ts @@ -21,27 +21,35 @@ export type FeedEvent = export type ModelRuntimeState = "stopped" | "loading" | "ready"; export class InflightTracker { - private requests = new Map(); + private requests = new Map(); private states = new Map(); + constructor(private now: () => number = () => Date.now()) {} + apply(event: FeedEvent): void { if (event.type === "inflight") { if (event.operation === "snapshot") { - this.requests = new Map(); - for (const r of event.requests) if (r.id) this.requests.set(r.id, r.model); + this.requests = new Map(); + for (const r of event.requests) if (r.id) this.requests.set(r.id, { model: r.model, seen: this.now() }); } else if (event.operation === "add") { - for (const r of event.requests) if (r.id) this.requests.set(r.id, r.model); + for (const r of event.requests) if (r.id) this.requests.set(r.id, { model: r.model, seen: this.now() }); } else if (event.operation === "remove") { this.requests.delete(event.id); } + this.prune(); } else if (event.type === "modelStatus") { for (const m of event.models) this.states.set(m.id, normalizeState(m.state)); } } + prune(maxAgeMs = 120_000): void { + const cutoff = this.now() - maxAgeMs; + for (const [id, entry] of this.requests) if (entry.seen < cutoff) this.requests.delete(id); + } + count(modelId: string): number { let n = 0; - for (const model of this.requests.values()) if (model === modelId) n++; + for (const entry of this.requests.values()) if (entry.model === modelId) n++; return n; } diff --git a/src/lib/runtime.ts b/src/lib/runtime.ts index ad89877..6e4ecfb 100644 --- a/src/lib/runtime.ts +++ b/src/lib/runtime.ts @@ -12,6 +12,7 @@ class Runtime { private feed?: EventFeed; private poller?: MetricsPoller; private statsCache?: StatsCache; + private pruneTimer?: ReturnType; private listeners = new Set<() => void>(); ensureConnections(cfg: LlamaSwapConfig): void { @@ -45,6 +46,9 @@ class Runtime { this.statsCache = new StatsCache(cfg); this.statsCache.onChange(() => this.emit()); } + if (!this.pruneTimer) { + this.pruneTimer = setInterval(() => this.tracker.prune(), PRUNE_INTERVAL_MS); + } } get pollerInstance(): MetricsPoller | undefined { @@ -73,3 +77,5 @@ class Runtime { } export const runtime = new Runtime(); + +const PRUNE_INTERVAL_MS = 30_000; diff --git a/tests/event-feed.test.ts b/tests/event-feed.test.ts index dc66157..d89b1d5 100644 --- a/tests/event-feed.test.ts +++ b/tests/event-feed.test.ts @@ -67,6 +67,15 @@ test("decodeEvent parses an inflight remove with a bare id", () => { assert.deepEqual(ev, { type: "inflight", operation: "remove", id: "817" }); }); +test("decodeEvent accepts a numeric id on remove", () => { + const msg: SseMessage = { + event: "message", + data: JSON.stringify({ type: "inflight", data: JSON.stringify({ operation: "remove", id: 820 }) }), + }; + const ev = decodeEvent(msg); + assert.deepEqual(ev, { type: "inflight", operation: "remove", id: "820" }); +}); + test("decodeEvent parses modelStatus", () => { const msg: SseMessage = { event: "message", diff --git a/tests/inflight-tracker.test.ts b/tests/inflight-tracker.test.ts index fe43e18..f00862d 100644 --- a/tests/inflight-tracker.test.ts +++ b/tests/inflight-tracker.test.ts @@ -73,3 +73,34 @@ test("total sums in-flight requests across all models", () => { tracker.apply({ type: "inflight", operation: "remove", id: "1" }); assert.equal(tracker.total(), 2); }); + +test("prune drops request ids that stopped receiving updates", () => { + let now = 0; + const tracker = new InflightTracker(() => now); + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] }); + now = 200_000; + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "2" }] }); + tracker.prune(60_000); + assert.equal(tracker.total(), 1); + assert.equal(tracker.count("A"), 1); +}); + +test("active requests refreshed by upserts survive pruning", () => { + let now = 0; + const tracker = new InflightTracker(() => now); + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] }); + now = 200_000; + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] }); + tracker.prune(60_000); + assert.equal(tracker.total(), 1); +}); + +test("prune is applied on every inflight event", () => { + let now = 0; + const tracker = new InflightTracker(() => now); + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "stale" }] }); + now = 200_000; + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "B", id: "live" }] }); + assert.equal(tracker.total(), 1); + assert.equal(tracker.count("B"), 1); +});