@elizaos/plugin-local-inference 2.0.0-beta.1 → 2.0.3-beta.3

package/src/services/context-fit.test.ts

@@ -0,0 +1,121 @@
+import { describe, expect, it } from "vitest";
+import { computeRuntimeContextFit } from "./context-fit";
+
+describe("computeRuntimeContextFit", () => {
+	it("shrinks context to the largest 4k window that fits the q8_0 KV budget", () => {
+		const fit = computeRuntimeContextFit({
+			params: "2B",
+			weightMb: 1434,
+			usableMb: 2560,
+			nativeContext: 131072,
+		});
+
+		expect(fit).not.toBeNull();
+		expect(fit?.contextDownscaled).toBe(true);
+		expect(fit?.contextSize).toBeGreaterThanOrEqual(8192);
+		expect(fit?.contextSize).toBeLessThan(131072);
+		expect(fit?.contextSize % 4096).toBe(0);
+		expect(fit?.kvBytesPerToken).toBeGreaterThan(0);
+	});
+
+	it("keeps the native context when the KV budget has headroom", () => {
+		const fit = computeRuntimeContextFit({
+			params: "9B",
+			weightMb: 5529,
+			usableMb: 23 * 1024,
+			nativeContext: 131072,
+		});
+
+		expect(fit?.contextSize).toBe(131072);
+		expect(fit?.contextDownscaled).toBe(false);
+		expect(fit?.maxFittingContext).toBeGreaterThanOrEqual(131072);
+	});
+
+	it("returns null when not even the minimum local context fits", () => {
+		const fit = computeRuntimeContextFit({
+			params: "2B",
+			weightMb: 2200,
+			usableMb: 2300,
+			nativeContext: 131072,
+		});
+
+		expect(fit).toBeNull();
+	});
+
+	it("defaults to q8_0 KV", () => {
+		const fit = computeRuntimeContextFit({
+			params: "9B",
+			weightMb: 5529,
+			usableMb: 64 * 1024,
+			nativeContext: 131072,
+		});
+		expect(fit?.kvQuant).toBe("q8_0");
+	});
+
+	it("upgrades to f16 KV when there's headroom and it's opted in (#8809 AC#4)", () => {
+		const fit = computeRuntimeContextFit({
+			params: "9B",
+			weightMb: 5529,
+			usableMb: 64 * 1024,
+			nativeContext: 131072,
+			preferAccurateKvWhenHeadroom: true,
+		});
+		expect(fit?.kvQuant).toBe("f16");
+		// f16 still affords the full native window on a roomy host.
+		expect(fit?.contextSize).toBe(131072);
+		expect(fit?.contextDownscaled).toBe(false);
+	});
+
+	it("never trades context for f16 KV precision (stays q8_0 when f16 would shrink the window)", () => {
+		// kvBudget here lets q8_0 reach the full native window but f16 (≈1.88×)
+		// would not — so the opt-in must keep q8_0 rather than downscale context.
+		const fit = computeRuntimeContextFit({
+			params: "9B",
+			weightMb: 5529,
+			usableMb: 8053,
+			nativeContext: 131072,
+			preferAccurateKvWhenHeadroom: true,
+		});
+		expect(fit?.contextSize).toBe(131072);
+		expect(fit?.kvQuant).toBe("q8_0");
+	});
+
+	it("scales the selected context with the device RAM class (4/8/16/24/128 GB) (#8809 AC#4)", () => {
+		// One representative tier across the five canonical device classes.
+		// `weightMb` is the resident weight; `usableMb` is the post-headroom
+		// budget the caller (active-model.ts) passes. The q8_0 KV rate comes from
+		// the tier params, so the selected window tracks free RAM.
+		const deviceClassesGb = [4, 8, 16, 24, 128];
+		const fits = deviceClassesGb.map((gb) =>
+			computeRuntimeContextFit({
+				params: "9B",
+				weightMb: 6500,
+				usableMb: gb * 1024,
+				nativeContext: 131072,
+			}),
+		);
+
+		// 4 GB can't hold the weights + a minimum KV window → no fit.
+		expect(fits[0]).toBeNull();
+
+		// Every class that fits gets a 4k-aligned window within [min, native].
+		for (const fit of fits.slice(1)) {
+			expect(fit).not.toBeNull();
+			expect((fit?.contextSize ?? 0) % 4096).toBe(0);
+			expect(fit?.contextSize).toBeGreaterThanOrEqual(8192);
+			expect(fit?.contextSize).toBeLessThanOrEqual(131072);
+		}
+
+		// Context never shrinks as RAM grows (headroom → larger KV window).
+		const sizes = fits.map((f) => f?.contextSize ?? 0);
+		for (let i = 1; i < sizes.length; i++) {
+			expect(sizes[i]).toBeGreaterThanOrEqual(sizes[i - 1]);
+		}
+
+		// Tight-but-fitting (8 GB) downscales below native; roomy (128 GB) hits it.
+		expect(fits[1]?.contextDownscaled).toBe(true);
+		expect(fits[1]?.contextSize).toBeLessThan(131072);
+		expect(fits.at(-1)?.contextSize).toBe(131072);
+		expect(fits.at(-1)?.contextDownscaled).toBe(false);
+	});
+});

package/src/services/context-fit.ts

@@ -0,0 +1,113 @@
+import { ELIZA_1_MIN_LOCAL_CONTEXT } from "@elizaos/shared/local-inference";
+import { estimateQuantizedKvBytesPerToken } from "./kv-spill";
+
+const BYTES_PER_MIB = 1024 * 1024;
+const CONTEXT_STEP = 4096;
+const DEFAULT_WORKING_SET_MB = 1024;
+
+// q8_0 stores the KV cache at 34 bytes / 32 elements; f16 at 2 bytes / element.
+// f16 KV therefore costs ~1.88× the q8_0 per-token rate the estimate is keyed to.
+const F16_OVER_Q8_0_KV_RATIO = 2 / (34 / 32);
+
+export interface RuntimeContextFitInput {
+	params: string;
+	weightMb: number;
+	usableMb: number;
+	nativeContext: number;
+	minContext?: number;
+	workingSetMb?: number;
+	contextStep?: number;
+	/**
+	 * When the host has enough headroom to run the more accurate f16 KV cache at
+	 * (at least) the same window q8_0 would give, prefer f16 instead of leaving
+	 * precision on the table. Opt-in — q8_0 stays the default per the device-fit
+	 * contract; this only ever *upgrades* precision and never trades away context
+	 * (#8809 AC#4). See CONTEXT_SCALING.md §5.
+	 */
+	preferAccurateKvWhenHeadroom?: boolean;
+}
+
+export interface RuntimeContextFit {
+	contextSize: number;
+	contextDownscaled: boolean;
+	maxFittingContext: number;
+	kvBytesPerToken: number;
+	workingSetMb: number;
+	/** The KV cache precision the chosen window was sized against. */
+	kvQuant: "q8_0" | "f16";
+}
+
+function roundDownToStep(value: number, step: number): number {
+	return Math.max(0, Math.floor(value / step) * step);
+}
+
+/**
+ * Choose the runtime context window that fits the current host budget.
+ *
+ * The admission gate still decides whether the model may load at all. This
+ * helper only sizes the q8_0 KV window for an admitted Eliza-1 tier so a tight
+ * host gets the largest safe window instead of blindly taking the catalog
+ * ceiling.
+ */
+export function computeRuntimeContextFit(
+	input: RuntimeContextFitInput,
+): RuntimeContextFit | null {
+	const minContext = input.minContext ?? ELIZA_1_MIN_LOCAL_CONTEXT;
+	const step = input.contextStep ?? CONTEXT_STEP;
+	const workingSetMb = input.workingSetMb ?? DEFAULT_WORKING_SET_MB;
+	if (
+		!Number.isFinite(input.weightMb) ||
+		!Number.isFinite(input.usableMb) ||
+		!Number.isFinite(input.nativeContext) ||
+		input.weightMb <= 0 ||
+		input.usableMb <= 0 ||
+		input.nativeContext < minContext ||
+		step <= 0
+	) {
+		return null;
+	}
+
+	const kvBytesPerToken = estimateQuantizedKvBytesPerToken(input.params);
+	if (!Number.isFinite(kvBytesPerToken) || kvBytesPerToken <= 0) return null;
+
+	const kvBudgetMb = input.usableMb - input.weightMb - workingSetMb;
+	if (kvBudgetMb <= 0) return null;
+	const kvBudgetBytes = kvBudgetMb * BYTES_PER_MIB;
+
+	const q8MaxFittingContext = roundDownToStep(
+		kvBudgetBytes / kvBytesPerToken,
+		step,
+	);
+	if (q8MaxFittingContext < minContext) return null;
+	const q8ContextSize = Math.min(input.nativeContext, q8MaxFittingContext);
+
+	// Default: q8_0 KV, sized to the host. Opt-in headroom upgrade: if f16 KV
+	// still affords at least the q8_0-selected window, use it — more precise, and
+	// never at the cost of context.
+	let kvQuant: "q8_0" | "f16" = "q8_0";
+	let kvBytesPerTokenChosen = kvBytesPerToken;
+	let maxFittingContext = q8MaxFittingContext;
+	let contextSize = q8ContextSize;
+	if (input.preferAccurateKvWhenHeadroom) {
+		const f16BytesPerToken = kvBytesPerToken * F16_OVER_Q8_0_KV_RATIO;
+		const f16MaxFittingContext = roundDownToStep(
+			kvBudgetBytes / f16BytesPerToken,
+			step,
+		);
+		if (f16MaxFittingContext >= q8ContextSize) {
+			kvQuant = "f16";
+			kvBytesPerTokenChosen = f16BytesPerToken;
+			maxFittingContext = f16MaxFittingContext;
+			contextSize = Math.min(input.nativeContext, f16MaxFittingContext);
+		}
+	}
+
+	return {
+		contextSize,
+		contextDownscaled: contextSize < input.nativeContext,
+		maxFittingContext,
+		kvBytesPerToken: kvBytesPerTokenChosen,
+		workingSetMb,
+		kvQuant,
+	};
+}

package/src/services/conversation-registry.test.ts

@@ -0,0 +1,235 @@
+import { describe, expect, it } from "vitest";
+import {
+	ConversationRegistry,
+	conversationRegistry,
+} from "./conversation-registry";
+
+describe("ConversationRegistry.open", () => {
+	it("returns the same handle for repeated opens of the same conversation", () => {
+		const registry = new ConversationRegistry();
+		const a = registry.open({
+			conversationId: "room-1",
+			modelId: "eliza-1-9b",
+			parallel: 4,
+		});
+		const b = registry.open({
+			conversationId: "room-1",
+			modelId: "eliza-1-9b",
+			parallel: 4,
+		});
+		expect(b).toBe(a);
+		expect(registry.size()).toBe(1);
+	});
+
+	it("treats different model ids as distinct handles", () => {
+		const registry = new ConversationRegistry();
+		const a = registry.open({
+			conversationId: "room-1",
+			modelId: "model-a",
+			parallel: 4,
+		});
+		const b = registry.open({
+			conversationId: "room-1",
+			modelId: "model-b",
+			parallel: 4,
+		});
+		expect(b).not.toBe(a);
+		expect(registry.size()).toBe(2);
+	});
+
+	it("requires non-empty conversationId and modelId", () => {
+		const registry = new ConversationRegistry();
+		expect(() => registry.open({ conversationId: "", modelId: "m" })).toThrow();
+		expect(() => registry.open({ conversationId: "c", modelId: "" })).toThrow();
+	});
+
+	it("pins the handle to slot 0 when parallel <= 1", () => {
+		const registry = new ConversationRegistry();
+		const handle = registry.open({
+			conversationId: "x",
+			modelId: "m",
+			parallel: 1,
+		});
+		expect(handle.slotId).toBe(0);
+	});
+
+	it("spreads concurrent opens across slots, lowest-loaded first", () => {
+		const registry = new ConversationRegistry();
+		const slots = new Set<number>();
+		for (let i = 0; i < 4; i += 1) {
+			const handle = registry.open({
+				conversationId: `room-${i}`,
+				modelId: "m",
+				parallel: 4,
+			});
+			slots.add(handle.slotId);
+		}
+		expect(slots.size).toBe(4);
+	});
+});
+
+describe("ConversationRegistry.close", () => {
+	it("frees the slot and is idempotent", () => {
+		const registry = new ConversationRegistry();
+		const handle = registry.open({
+			conversationId: "x",
+			modelId: "m",
+			parallel: 4,
+		});
+		expect(handle.closed).toBe(false);
+		registry.close("x", "m");
+		registry.close("x", "m"); // idempotent — must not throw
+		expect(registry.get("x", "m")).toBeNull();
+	});
+
+	it("frees a slot for reuse on next open", () => {
+		const registry = new ConversationRegistry();
+		const a = registry.open({
+			conversationId: "a",
+			modelId: "m",
+			parallel: 2,
+		});
+		const b = registry.open({
+			conversationId: "b",
+			modelId: "m",
+			parallel: 2,
+		});
+		expect(a.slotId).not.toBe(b.slotId);
+		registry.close("a", "m");
+		const c = registry.open({
+			conversationId: "c",
+			modelId: "m",
+			parallel: 2,
+		});
+		// c should land on the freed slot (a's slot)
+		expect(c.slotId).toBe(a.slotId);
+	});
+});
+
+describe("ConversationRegistry.get", () => {
+	it("returns null for unknown or closed handles", () => {
+		const registry = new ConversationRegistry();
+		expect(registry.get("nope", "m")).toBeNull();
+		registry.open({ conversationId: "x", modelId: "m", parallel: 4 });
+		registry.close("x", "m");
+		expect(registry.get("x", "m")).toBeNull();
+	});
+});
+
+describe("ConversationRegistry.evictIdle", () => {
+	it("drops handles whose ttl has elapsed", () => {
+		const registry = new ConversationRegistry();
+		registry.open({
+			conversationId: "x",
+			modelId: "m",
+			parallel: 4,
+			ttlMs: 1_000,
+		});
+		expect(registry.size()).toBe(1);
+		const dropped = registry.evictIdle(Date.now() + 5_000);
+		expect(dropped).toEqual(["x"]);
+		expect(registry.size()).toBe(0);
+	});
+
+	it("keeps handles whose ttl has NOT elapsed", () => {
+		const registry = new ConversationRegistry();
+		registry.open({
+			conversationId: "x",
+			modelId: "m",
+			parallel: 4,
+			ttlMs: 60_000,
+		});
+		const dropped = registry.evictIdle(Date.now() + 10_000);
+		expect(dropped).toEqual([]);
+		expect(registry.size()).toBe(1);
+	});
+});
+
+describe("ConversationRegistry.highWater", () => {
+	it("tracks the largest concurrent open count", () => {
+		const registry = new ConversationRegistry();
+		expect(registry.highWater()).toBe(0);
+		registry.open({ conversationId: "a", modelId: "m", parallel: 8 });
+		registry.open({ conversationId: "b", modelId: "m", parallel: 8 });
+		registry.open({ conversationId: "c", modelId: "m", parallel: 8 });
+		expect(registry.highWater()).toBe(3);
+		registry.close("a", "m");
+		registry.close("b", "m");
+		// High-water mark must NOT decrease — it's a max over the lifetime
+		expect(registry.highWater()).toBe(3);
+	});
+});
+
+describe("ConversationRegistry.recommendedParallel (--parallel auto-resize decision)", () => {
+	it("returns the running count when the high-water mark hasn't outgrown it", () => {
+		const registry = new ConversationRegistry();
+		// 2 concurrent, headroom max(2, ceil(2*0.25)=1) = 2 → desired 4.
+		registry.open({ conversationId: "a", modelId: "m", parallel: 4 });
+		registry.open({ conversationId: "b", modelId: "m", parallel: 4 });
+		expect(registry.highWater()).toBe(2);
+		expect(registry.recommendedParallel(4)).toBe(4); // 4 already covers it
+		expect(registry.recommendedParallel(8)).toBe(8); // larger running wins
+	});
+
+	it("recommends high-water + 25%-headroom when it exceeds the running count", () => {
+		const registry = new ConversationRegistry();
+		for (let i = 0; i < 20; i += 1) {
+			registry.open({ conversationId: `c-${i}`, modelId: "m", parallel: 4 });
+		}
+		expect(registry.highWater()).toBe(20);
+		// 20 + max(2, ceil(20*0.25)=5) = 25.
+		expect(registry.recommendedParallel(4)).toBe(25);
+	});
+
+	it("headroom floors at 2 (small high-water marks still get a buffer)", () => {
+		const registry = new ConversationRegistry();
+		for (let i = 0; i < 5; i += 1) {
+			registry.open({ conversationId: `c-${i}`, modelId: "m", parallel: 2 });
+		}
+		expect(registry.highWater()).toBe(5);
+		// ceil(5*0.25) = 2 → headroom 2 → desired 7.
+		expect(registry.recommendedParallel(2)).toBe(7);
+	});
+
+	it("is monotonic: closing conversations does not shrink the recommendation", () => {
+		const registry = new ConversationRegistry();
+		const handles = Array.from({ length: 10 }, (_, i) =>
+			registry.open({ conversationId: `c-${i}`, modelId: "m", parallel: 4 }),
+		);
+		expect(registry.recommendedParallel(4)).toBe(13); // 10 + ceil(10*.25)=3
+		for (const h of handles) registry.close(h.conversationId, h.modelId);
+		expect(registry.size()).toBe(0);
+		expect(registry.recommendedParallel(4)).toBe(13); // unchanged
+	});
+});
+
+describe("ConversationRegistry.__resetForTests", () => {
+	it("drops every handle and resets the high-water mark", () => {
+		const registry = new ConversationRegistry();
+		const a = registry.open({ conversationId: "a", modelId: "m", parallel: 4 });
+		registry.open({ conversationId: "b", modelId: "m", parallel: 4 });
+		expect(registry.size()).toBe(2);
+		expect(registry.highWater()).toBe(2);
+		registry.__resetForTests();
+		expect(registry.size()).toBe(0);
+		expect(registry.highWater()).toBe(0);
+		expect(registry.recommendedParallel(4)).toBe(4);
+		// The dropped handle is marked closed (further use is rejected by the engine).
+		expect(a.closed).toBe(true);
+		// A slot freed by reset is reusable from slot 0 again.
+		const handle = registry.open({
+			conversationId: "c",
+			modelId: "m",
+			parallel: 4,
+		});
+		expect(handle.slotId).toBe(0);
+	});
+
+	it("isolates the module singleton across test files", () => {
+		conversationRegistry.__resetForTests();
+		conversationRegistry.open({ conversationId: "leak", modelId: "m" });
+		expect(conversationRegistry.size()).toBe(1);
+		conversationRegistry.__resetForTests();
+		expect(conversationRegistry.size()).toBe(0);
+	});
+});