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); +});