diff --git a/src/lib/inflight-tracker.ts b/src/lib/inflight-tracker.ts new file mode 100644 index 0000000..ad5dbba --- /dev/null +++ b/src/lib/inflight-tracker.ts @@ -0,0 +1,58 @@ +export interface InflightRequest { + id?: string; + model: string; + req_path?: string; + method?: string; + timestamp?: string; + elapsed_ms?: number; +} + +export interface ModelState { + id: string; + state: string; +} + +export type FeedEvent = + | { type: "inflight"; operation: "snapshot" | "add" | "remove"; requests: InflightRequest[] } + | { type: "modelStatus"; models: ModelState[] }; + +export type ModelRuntimeState = "stopped" | "loading" | "ready"; + +export class InflightTracker { + private counts = new Map(); + private states = new Map(); + + apply(event: FeedEvent): void { + if (event.type === "inflight") { + if (event.operation === "snapshot") { + const counts = new Map(); + for (const r of event.requests) counts.set(r.model, (counts.get(r.model) ?? 0) + 1); + this.counts = counts; + } else if (event.operation === "add") { + for (const r of event.requests) this.counts.set(r.model, (this.counts.get(r.model) ?? 0) + 1); + } else if (event.operation === "remove") { + for (const r of event.requests) { + const c = this.counts.get(r.model) ?? 0; + if (c <= 1) this.counts.delete(r.model); + else this.counts.set(r.model, c - 1); + } + } + } else if (event.type === "modelStatus") { + for (const m of event.models) this.states.set(m.id, normalizeState(m.state)); + } + } + + count(modelId: string): number { + return this.counts.get(modelId) ?? 0; + } + + state(modelId: string): ModelRuntimeState | undefined { + return this.states.get(modelId); + } +} + +function normalizeState(state: string): ModelRuntimeState { + if (state === "ready") return "ready"; + if (state === "loading") return "loading"; + return "stopped"; +} diff --git a/src/lib/sse.ts b/src/lib/sse.ts new file mode 100644 index 0000000..01fa21b --- /dev/null +++ b/src/lib/sse.ts @@ -0,0 +1,29 @@ +export interface SseMessage { + event?: string; + data: string; +} + +export function parseSse(chunk: string): { messages: SseMessage[]; rest: string } { + const messages: SseMessage[] = []; + let rest = chunk; + while (true) { + const idx = rest.indexOf("\n\n"); + if (idx === -1) break; + const block = rest.slice(0, idx); + rest = rest.slice(idx + 2); + const message = toMessage(block); + if (message && message.data.length > 0) messages.push(message); + } + return { messages, rest }; +} + +function toMessage(block: string): SseMessage | null { + let event: string | undefined; + const dataLines: string[] = []; + for (const line of block.split("\n")) { + if (line.startsWith("event:")) event = line.slice(6).trim(); + else if (line.startsWith("data:")) dataLines.push(line.slice(5).trimStart()); + } + if (dataLines.length === 0) return null; + return { event, data: dataLines.join("\n") }; +} diff --git a/tests/inflight-tracker.test.ts b/tests/inflight-tracker.test.ts new file mode 100644 index 0000000..0396cbe --- /dev/null +++ b/tests/inflight-tracker.test.ts @@ -0,0 +1,50 @@ +import assert from "node:assert/strict"; +import { test } from "node:test"; +import { InflightTracker, type FeedEvent } from "../src/lib/inflight-tracker"; + +test("snapshot rebuilds counts for all in-flight requests", () => { + const tracker = new InflightTracker(); + tracker.apply({ + type: "inflight", + operation: "snapshot", + requests: [ + { model: "A", id: "1" }, + { model: "A", id: "2" }, + { model: "B", id: "3" }, + ], + }); + assert.equal(tracker.count("A"), 2); + assert.equal(tracker.count("B"), 1); + assert.equal(tracker.count("C"), 0); +}); + +test("add and remove adjust per-model counts", () => { + const tracker = new InflightTracker(); + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "1" }] }); + tracker.apply({ type: "inflight", operation: "add", requests: [{ model: "A", id: "2" }] }); + assert.equal(tracker.count("A"), 2); + tracker.apply({ type: "inflight", operation: "remove", requests: [{ model: "A", id: "1" }] }); + assert.equal(tracker.count("A"), 1); + tracker.apply({ type: "inflight", operation: "remove", requests: [{ model: "A", id: "2" }] }); + assert.equal(tracker.count("A"), 0); + tracker.apply({ type: "inflight", operation: "remove", requests: [{ model: "A", id: "9" }] }); + assert.equal(tracker.count("A"), 0); +}); + +test("modelStatus normalizes states", () => { + const tracker = new InflightTracker(); + tracker.apply({ + type: "modelStatus", + models: [ + { id: "A", state: "ready" }, + { id: "B", state: "loading" }, + { id: "C", state: "stopped" }, + { id: "D", state: "weird" }, + ], + }); + assert.equal(tracker.state("A"), "ready"); + assert.equal(tracker.state("B"), "loading"); + assert.equal(tracker.state("C"), "stopped"); + assert.equal(tracker.state("D"), "stopped"); + assert.equal(tracker.state("missing"), undefined); +}); diff --git a/tests/sse.test.ts b/tests/sse.test.ts new file mode 100644 index 0000000..4c71b43 --- /dev/null +++ b/tests/sse.test.ts @@ -0,0 +1,24 @@ +import assert from "node:assert/strict"; +import { test } from "node:test"; +import { parseSse } from "../src/lib/sse"; + +test("parseSse extracts event + data blocks separated by blank lines", () => { + const { messages, rest } = parseSse('event:message\ndata:{"a":1}\n\nevent:message\ndata:hello\n\n'); + assert.equal(rest, ""); + assert.equal(messages.length, 2); + assert.equal(messages[0].event, "message"); + assert.equal(messages[0].data, '{"a":1}'); + assert.equal(messages[1].data, "hello"); +}); + +test("parseSse keeps a trailing partial event in rest", () => { + const { messages, rest } = parseSse("event:message\ndata:par"); + assert.equal(messages.length, 0); + assert.equal(rest, "event:message\ndata:par"); +}); + +test("parseSse joins multi-line data fields with newline", () => { + const { messages } = parseSse("data:line1\ndata:line2\n\n"); + assert.equal(messages.length, 1); + assert.equal(messages[0].data, "line1\nline2"); +});