feat: add SSE event feed for llama-swap api/events
This commit is contained in:
@@ -0,0 +1,105 @@
|
|||||||
|
import { type LlamaSwapConfig } from "./util";
|
||||||
|
import { parseSse, type SseMessage } from "./sse";
|
||||||
|
import { type FeedEvent, type InflightRequest, type ModelState } from "./inflight-tracker";
|
||||||
|
|
||||||
|
export function decodeEvent(msg: SseMessage): FeedEvent | null {
|
||||||
|
try {
|
||||||
|
const outer = JSON.parse(msg.data) as { type?: string; data?: string };
|
||||||
|
if (!outer.data) return null;
|
||||||
|
const inner = JSON.parse(outer.data) as Record<string, unknown>;
|
||||||
|
|
||||||
|
if (outer.type === "inflight") {
|
||||||
|
const requests: InflightRequest[] =
|
||||||
|
(inner.requests as InflightRequest[]) ?? (inner.request ? [inner.request as InflightRequest] : []);
|
||||||
|
return {
|
||||||
|
type: "inflight",
|
||||||
|
operation: inner.operation as "snapshot" | "add" | "remove",
|
||||||
|
requests,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (outer.type === "modelStatus") {
|
||||||
|
return { type: "modelStatus", models: inner as unknown as ModelState[] };
|
||||||
|
}
|
||||||
|
|
||||||
|
return null;
|
||||||
|
} catch {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const MAX_RETRY_MS = 30000;
|
||||||
|
const INITIAL_RETRY_MS = 1000;
|
||||||
|
|
||||||
|
export class EventFeed {
|
||||||
|
private controller?: AbortController;
|
||||||
|
private closed = false;
|
||||||
|
private retryMs = INITIAL_RETRY_MS;
|
||||||
|
private statusHandler?: (connected: boolean) => void;
|
||||||
|
|
||||||
|
constructor(
|
||||||
|
private cfg: LlamaSwapConfig,
|
||||||
|
private onEvent: (ev: FeedEvent) => void,
|
||||||
|
) {}
|
||||||
|
|
||||||
|
setStatusHandler(fn: (connected: boolean) => void): void {
|
||||||
|
this.statusHandler = fn;
|
||||||
|
}
|
||||||
|
|
||||||
|
start(): void {
|
||||||
|
void this.connect();
|
||||||
|
}
|
||||||
|
|
||||||
|
stop(): void {
|
||||||
|
this.closed = true;
|
||||||
|
this.controller?.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
private async connect(): Promise<void> {
|
||||||
|
while (!this.closed) {
|
||||||
|
this.controller = new AbortController();
|
||||||
|
try {
|
||||||
|
const headers: Record<string, string> = { Accept: "text/event-stream" };
|
||||||
|
if (this.cfg.apiKey) headers["Authorization"] = `Bearer ${this.cfg.apiKey}`;
|
||||||
|
const res = await fetch(`${this.cfg.baseUrl}/api/events`, {
|
||||||
|
headers,
|
||||||
|
signal: this.controller.signal,
|
||||||
|
});
|
||||||
|
if (!res.ok || !res.body) throw new Error(`events HTTP ${res.status}`);
|
||||||
|
this.statusHandler?.(true);
|
||||||
|
this.retryMs = INITIAL_RETRY_MS;
|
||||||
|
await this.readStream(res.body.getReader());
|
||||||
|
} catch {
|
||||||
|
if (this.closed) return;
|
||||||
|
}
|
||||||
|
this.statusHandler?.(false);
|
||||||
|
if (this.closed) return;
|
||||||
|
await sleep(this.retryMs);
|
||||||
|
this.retryMs = Math.min(this.retryMs * 2, MAX_RETRY_MS);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private async readStream(reader: ReadableStreamDefaultReader<Uint8Array>): Promise<void> {
|
||||||
|
const decoder = new TextDecoder();
|
||||||
|
let buffer = "";
|
||||||
|
try {
|
||||||
|
while (!this.closed) {
|
||||||
|
const { done, value } = await reader.read();
|
||||||
|
if (done) break;
|
||||||
|
buffer += decoder.decode(value, { stream: true });
|
||||||
|
const { messages, rest } = parseSse(buffer);
|
||||||
|
buffer = rest;
|
||||||
|
for (const msg of messages) {
|
||||||
|
const ev = decodeEvent(msg);
|
||||||
|
if (ev) this.onEvent(ev);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} finally {
|
||||||
|
reader.releaseLock();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function sleep(ms: number): Promise<void> {
|
||||||
|
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||||
|
}
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { test } from "node:test";
|
||||||
|
import { decodeEvent } from "../src/lib/event-feed";
|
||||||
|
import type { SseMessage } from "../src/lib/sse";
|
||||||
|
|
||||||
|
test("decodeEvent parses an inflight snapshot", () => {
|
||||||
|
const msg: SseMessage = {
|
||||||
|
event: "message",
|
||||||
|
data: JSON.stringify({
|
||||||
|
type: "inflight",
|
||||||
|
data: JSON.stringify({
|
||||||
|
operation: "snapshot",
|
||||||
|
requests: [{ id: "14", model: "Qwen3.8-27B-NVFP4", req_path: "/v1/chat/completions", elapsed_ms: 165510 }],
|
||||||
|
}),
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
const ev = decodeEvent(msg);
|
||||||
|
assert.ok(ev);
|
||||||
|
assert.equal(ev.type, "inflight");
|
||||||
|
if (ev.type === "inflight") {
|
||||||
|
assert.equal(ev.operation, "snapshot");
|
||||||
|
assert.equal(ev.requests.length, 1);
|
||||||
|
assert.equal(ev.requests[0].model, "Qwen3.8-27B-NVFP4");
|
||||||
|
assert.equal(ev.requests[0].elapsed_ms, 165510);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("decodeEvent parses an inflight add with a singular request field", () => {
|
||||||
|
const msg: SseMessage = {
|
||||||
|
event: "message",
|
||||||
|
data: JSON.stringify({
|
||||||
|
type: "inflight",
|
||||||
|
data: JSON.stringify({ operation: "add", request: { id: "9", model: "A" } }),
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
const ev = decodeEvent(msg);
|
||||||
|
assert.ok(ev && ev.type === "inflight");
|
||||||
|
if (ev && ev.type === "inflight") {
|
||||||
|
assert.equal(ev.operation, "add");
|
||||||
|
assert.equal(ev.requests[0].model, "A");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("decodeEvent parses modelStatus", () => {
|
||||||
|
const msg: SseMessage = {
|
||||||
|
event: "message",
|
||||||
|
data: JSON.stringify({
|
||||||
|
type: "modelStatus",
|
||||||
|
data: JSON.stringify([
|
||||||
|
{ id: "DeepSeek-V4-Flash-0731", state: "ready" },
|
||||||
|
{ id: "Qwen3.8-27B-NVFP4", state: "stopped" },
|
||||||
|
]),
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
const ev = decodeEvent(msg);
|
||||||
|
assert.ok(ev && ev.type === "modelStatus");
|
||||||
|
if (ev && ev.type === "modelStatus") {
|
||||||
|
assert.equal(ev.models.length, 2);
|
||||||
|
assert.equal(ev.models[0].state, "ready");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("decodeEvent returns null for unrelated or malformed events", () => {
|
||||||
|
assert.equal(decodeEvent({ event: "message", data: JSON.stringify({ type: "logData", data: "{}" }) }), null);
|
||||||
|
assert.equal(decodeEvent({ event: "message", data: "not json" }), null);
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user