auto-model-router 0.2.25 → 0.2.27

package/omp-extension/router-embed.ts

@@ -19,11 +19,11 @@
  *     - /path/to/auto-model-router/omp-extension/router-toast.ts
  */
 
-import { readFileSync } from "node:fs";
+import { appendFileSync, readFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { join } from "node:path";
 
-import { ompModelsPath, syncModelsYml } from "../src/cli/config-cmd.ts";
+import { ompModelsPath } from "../src/cli/config-cmd.ts";
 import { loadConfig } from "../src/config/load.ts";
 import { startServer } from "../src/server/http.ts";
 import type { StartedServer } from "../src/server/http.ts";
@@ -52,6 +52,39 @@ function readModelsYml(): string {
 	}
 }
 
+/**
+ * Appends one line to `$AUTO_MODEL_ROUTER_HOME/embed.log`.
+ *
+ * A FILE, deliberately: `console.*` from an extension does not reach omp's
+ * session log, so the first attempt at this diagnostic left no trace anywhere
+ * and the port lifecycle had to be reconstructed from netstat and mtimes. Best
+ * effort — a logging failure must never break a session.
+ */
+function writeEmbedLog(line: string): void {
+	try {
+		appendFileSync(join(routerHome(), "embed.log"), `${new Date().toISOString()} ${line}\n`, "utf8");
+	} catch {
+		// Unwritable home: the router still works, we just lose the breadcrumb.
+	}
+}
+
+/**
+ * Stops the router when the PROCESS ends — never when a session does.
+ * Idempotent: registered once, however many sessions this process hosts.
+ */
+let exitHooked = false;
+function trackProcessExit(): void {
+	if (exitHooked) return;
+	exitHooked = true;
+	// `exit` cannot await, and does not need to: the OS reclaims the socket.
+	// The signal hooks exist so a Ctrl-C releases the port promptly.
+	for (const signal of ["SIGINT", "SIGTERM"] as const) {
+		process.once(signal, () => {
+			void app?.stop().catch(() => {});
+		});
+	}
+}
+
 /**
  * Registers the auto-model-router provider (and its virtual models) into omp's model
  * registry at a specific bound port.
@@ -94,16 +127,29 @@ function registerRouterProvider(pi: ExtensionAPI, port: number, cfg: RouterConfi
 	});
 }
 
+/**
+ * The router bound by THIS PROCESS, and the port it serves.
+ *
+ * Module scope on purpose: Bun caches the module per process, so when omp loads
+ * the extension into a second host (its provider-refresh / reload path does
+ * exactly that) these stay visible. A second host then REUSES this router
+ * instead of binding another port and orphaning every model handle omp already
+ * resolved against the first one.
+ */
+let app: StartedServer | null = null;
+let boundPort: number | null = null;
+
+/** `$AUTO_MODEL_ROUTER_HOME`, tilde-expanded, defaulting to ~/.auto-model-router. */
+function routerHome(): string {
+	const raw = process.env.AUTO_MODEL_ROUTER_HOME ?? join(homedir(), ".auto-model-router");
+	return raw === "~" || raw.startsWith("~/") || raw.startsWith("~\\") ? join(homedir(), raw.slice(1)) : raw;
+}
+
 export default function (pi: ExtensionAPI): void {
 	pi.setLabel("auto-model-router embed");
 
 	// Shared port file, written only by the main session's router.
-	const homeRaw = process.env.AUTO_MODEL_ROUTER_HOME ?? join(homedir(), ".auto-model-router");
-	const home =
-		homeRaw === "~" || homeRaw.startsWith("~/") || homeRaw.startsWith("~\\")
-			? join(homedir(), homeRaw.slice(1))
-			: homeRaw;
-	const portFile = embedPortPath(home);
+	const portFile = embedPortPath(routerHome());
 	// Each interactive session binds its OWN router on an ephemeral port, so
 	// sessions stay independent: no shared process to contend over, and no
 	// session left broken because another one exited. `server.port` from
@@ -113,12 +159,22 @@ export default function (pi: ExtensionAPI): void {
 	const requestedPort = resolveEmbedPort(process.env.AUTO_MODEL_ROUTER_PORT);
 	const cfg = loadConfig({ overrides: { server: { host: "127.0.0.1", port: requestedPort } } });
 
-	let app: StartedServer | null = null;
-
 	pi.on("session_start", async (_event, ctx) => {
 		// The omp UI session id tags every request so the toast can scope its
 		// notifications to that exact session (see router-toast.ts).
 		const sessionId = ctx.sessionManager.getSessionId();
+
+		// This module is cached per PROCESS, so `app` and `boundPort` are
+		// process-global even when omp loads the extension into more than one
+		// host. A router already bound in this process is therefore reusable:
+		// re-register it for the new session id and return. Never rebind — a
+		// second bind would take a different port and orphan every model handle
+		// omp already resolved against the first one.
+		if (app !== null && boundPort !== null) {
+			registerRouterProvider(pi, boundPort, cfg, sessionId);
+			return;
+		}
+
 		if (!ctx.hasUI) {
 			// Subagents and headless (-p) sessions prefer the main session's
 			// shared router: one process, one ledger, one place to inspect.
@@ -136,11 +192,8 @@ export default function (pi: ExtensionAPI): void {
 			const started = startServer(cfg);
 			if (started.server.port === undefined) return;
 			app = started;
-			registerRouterProvider(pi, started.server.port, cfg, sessionId);
-			pi.on("session_shutdown", () => {
-				void app?.stop().catch(() => {});
-				app = null;
-			});
+			boundPort = started.server.port;
+			registerRouterProvider(pi, boundPort, cfg, sessionId);
 			return;
 		}
 
@@ -148,68 +201,91 @@ export default function (pi: ExtensionAPI): void {
 		// the provider against the exact bound port. Registration happens only
 		// here — never at factory load, where a stale shared port would be
 		// captured into omp's model registry and defeat the correct bound URL.
-		if (app) return;
 
 		// A fixed port was asked for (env var) and a live router already answers
 		// there: share it rather than failing to bind. Ephemeral ports — the
 		// default — never take this path, so sessions stay independent.
 		if (requestedPort !== 0 && (await probeEmbed(requestedPort))) {
 			writeEmbedPort(portFile, requestedPort);
-			syncModelsYml(cfg, requestedPort);
 			registerRouterProvider(pi, requestedPort, cfg, sessionId);
 			pi.setLabel(`auto-model-router embed (shared :${requestedPort})`);
 			return;
 		}
 
-		// Bind. If a FIXED port was requested and something else holds it, fall
-		// back to an ephemeral one rather than leaving this session with no
-		// provider at all. An ephemeral request that fails is a real error.
+		// ADOPT THE ADVERTISED PORT. This is the fix for "provider error: Unable
+		// to connect" on every real turn while utility calls kept working.
+		//
+		// omp resolves `modelRoles.default` (auto-model-router/auto) from
+		// models.yml during STARTUP — before this extension loads, so before we
+		// can bind or register anything. That resolved handle is a SNAPSHOT: a
+		// later registerProvider replaces the registry entry but cannot rewrite
+		// a handle omp already built. models.yml names the port of the LAST
+		// session that wrote it, which after a normal restart is the session the
+		// user just closed — a dead socket. Utility calls (title generation,
+		// auto-thinking) resolve AFTER our registration and so hit the live port,
+		// which is exactly the asymmetry that made this look like a router fault.
+		// Measured in the field via embed.log:
+		//   embed ready pid=61872 port=54985 models.yml-advertised=50596
+		//
+		// So bind the port omp already resolved against, whenever nothing holds
+		// it. Each session still runs its OWN router — this only chooses which
+		// port that router listens on. A live peer holding it means the handle
+		// works anyway (that peer serves it), and we fall back to ephemeral.
+		const advertised = modelsYmlPort(readModelsYml());
+		if (requestedPort === 0 && advertised !== null) cfg.server.port = advertised;
+
+		// Fall back to an ephemeral port when the preferred one is taken, rather
+		// than leaving this session with no provider at all.
 		let started: StartedServer;
 		try {
 			started = startServer(cfg);
-		} catch (err) {
-			if (requestedPort === 0) throw err;
+		} catch {
 			cfg.server.port = 0;
 			started = startServer(cfg);
 		}
 		const actualPort = started.server.port;
 		if (actualPort === undefined) return;
 		app = started;
+		boundPort = actualPort;
 
 		// Publish the port; subagents and the toast read it from here.
 		writeEmbedPort(portFile, actualPort);
 
-		// Keep models.yml pointing at this port. Headless runs (`-p`) and
-		// subagent processes resolve models from models.yml in a FRESH registry
-		// — extension registration does not reach them — so without this they
-		// fail with "Model not found" when no interactive session is live
-		// (the print-mode gap the external benchmark hit).
-		const advertised = modelsYmlPort(readModelsYml());
-		const syncAction = syncModelsYml(cfg, actualPort);
+		// models.yml is deliberately NOT written. The port belongs to THIS
+		// process and changes every launch, so persisting it into a file omp
+		// reads at STARTUP — before this extension loads — makes a dead port
+		// authoritative for the next session's `modelRoles.default`, and that
+		// handle is a snapshot no later registration can repair. That is the
+		// regression this whole class of failure came from. The provider is
+		// registered dynamically below instead, which is what worked before the
+		// file was ever written.
+		//
+		// `auto-model-router config --write` still exists for anyone who wants a
+		// static block on purpose; the adoption above keeps such a block
+		// harmless by binding whatever port it names when that port is free.
 
 		// Register BEFORE any await: everything omp resolves after this point
 		// picks up the live URL, so the registration must not sit behind I/O.
 		registerRouterProvider(pi, actualPort, cfg, sessionId);
-		pi.setLabel(`auto-model-router embed :${actualPort}${syncAction === null ? "" : ` (models.yml ${syncAction})`}`);
+		pi.setLabel(`auto-model-router embed :${actualPort}`);
 
-		pi.on("session_shutdown", () => {
-			void app?.stop().catch(() => {});
-			app = null;
-		});
+		// NO `session_shutdown` teardown. That event is emitted from session
+		// DISPOSAL — including omp's provider-refresh / extension-reload path,
+		// which runs in a throwaway extension host while the real session keeps
+		// going. Because this module is cached per process, such a handler stops
+		// the LIVE router: every subsequent turn then fails with Bun's
+		// "Unable to connect", while utility calls that resolve after a later
+		// rebind still work — the exact asymmetry observed in the field. The
+		// router's lifetime is the PROCESS, and the OS reclaims the socket when
+		// the process exits.
+		trackProcessExit();
 
-		// Diagnostic, not a guess. "provider error: Unable to connect" on real
-		// turns while utility calls keep working means omp dialled an address
-		// this router does not serve — and that text is Bun's fetch error, which
-		// names no URL. Log the facts that distinguish the cases: the port bound,
-		// the port models.yml advertised BEFORE the sync above, and whether our
-		// own socket answers.
-		const selfProbeOk = await probeEmbed(actualPort);
-		console.info(
-			`[auto-model-router] embed ready: serving :${actualPort}` +
-				` | models.yml advertised :${advertised ?? "none"} at startup` +
-				` | self-probe ${selfProbeOk ? "ok" : "FAILED"}` +
-				` | session ${sessionId}`,
+		writeEmbedLog(
+			`embed ready pid=${process.pid} port=${actualPort}` +
+				` models.yml-advertised=${advertised ?? "none"}` +
+				` models.yml-untouched` +
+				` self-probe=${(await probeEmbed(actualPort)) ? "ok" : "FAILED"}` +
+				` session=${sessionId}`,
 		);
-		if (!selfProbeOk) pi.setLabel(`auto-model-router: BOUND :${actualPort} BUT NOT ANSWERING`);
 	});
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "auto-model-router",
-  "version": "0.2.25",
+  "version": "0.2.27",
   "private": false,
   "description": "Local cost/complexity-aware model router for Oh My Pi, backed by OpenRouter",
   "type": "module",

package/src/cli/config-cmd.ts

@@ -317,44 +317,6 @@ export function ompModelsPath(): string {
 	if (agentDir !== undefined && agentDir !== "") return join(agentDir, "models.yml");
 	return join(homedir(), ".omp", "agent", "models.yml");
 }
-/**
- * Keeps omp's models.yml carrying an up-to-date `auto-model-router` provider
- * block pointing at `port`, silently. The embedded extension calls this every
- * time the main session binds its router, so headless (`-p`) runs, subagent
- * processes, and any other consumer that builds a FRESH model registry (which
- * extension registration does NOT reach) still resolve
- * `auto-model-router/auto` — they read models.yml, not the live registry.
- *
- * Same splice + validation as `config --write`, minus the console output and
- * the backup: this runs on every session start, and a `.bak` per launch would
- * churn the directory. Failure to write is logged by the caller, never thrown —
- * a read-only models.yml must degrade to "registration only", which is the
- * pre-existing behavior.
- *
- * Returns the splice action, or null when nothing was written (already
- * current, or the write failed).
- */
-export function syncModelsYml(cfg: RouterConfig, port: number, target = ompModelsPath()): SpliceResult["action"] | null {
-	try {
-		const pointed: RouterConfig = { ...cfg, server: { ...cfg.server, port } };
-		const block = renderProviderBlock(pointed, null);
-		const existing = existsSync(target) ? readFileSync(target, "utf8") : "";
-		const result = spliceProviderBlock(existing, block);
-		if (result.action === "replaced") {
-			// The guards make "replaced" cheap to detect but the text may still be
-			// byte-identical (same port, same costs): skip the write so the file
-			// mtime stays stable for tools watching it.
-			if (result.text === existing) return null;
-		}
-		assertUsableModelsYaml(result.text);
-		mkdirSync(dirname(target), { recursive: true });
-		writeFileSync(target, result.text, "utf8");
-		return result.action;
-	} catch {
-		return null;
-	}
-}
-
 /** omp's models.yml location: `$PI_CODING_AGENT_DIR` relocates the whole agent dir. */
 
 export async function configCommand(args: CliArgs): Promise<void> {

package/test/embed-lifecycle.test.ts

@@ -0,0 +1,170 @@
+import { afterAll, beforeAll, describe, expect, test } from "bun:test";
+import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+import type { ExtensionAPI, ExtensionContext, ProviderRegistration } from "@oh-my-pi/pi-coding-agent";
+
+/**
+ * How the embedded router must behave, pinned against the two failures that
+ * produced "provider error: Unable to connect" on every real turn while utility
+ * calls kept working.
+ *
+ * 1. THE PORT COMES FROM THIS PROCESS, NOT FROM A FILE. omp resolves
+ *    `modelRoles.default` from models.yml during STARTUP — before this extension
+ *    loads — and that handle is a snapshot no later registerProvider can
+ *    rewrite. So the extension must not persist its ephemeral port into
+ *    models.yml (a dead port then becomes authoritative for the NEXT session),
+ *    and when a block already exists it adopts the port that block names so the
+ *    handle omp built is valid. Measured in the field:
+ *      embed ready pid=61872 port=54985 models.yml-advertised=50596
+ *
+ * 2. THE ROUTER'S LIFETIME IS THE PROCESS. omp emits `session_shutdown` from
+ *    session disposal, which includes its provider-refresh / extension-reload
+ *    path running in a throwaway host while the real session continues. This
+ *    module is cached per process, so a teardown there stopped the LIVE router.
+ *
+ * One boot shared by every test, deliberately: the module is cached per process
+ * in production too, so this is the real shape.
+ */
+
+const registrations: { id: string; baseUrl: string }[] = [];
+const handlers = new Map<string, ((event: unknown, ctx: ExtensionContext) => void | Promise<void>)[]>();
+
+const pi: ExtensionAPI = {
+	setLabel: () => {},
+	on: (event, handler) => {
+		const list = handlers.get(event) ?? [];
+		list.push(handler);
+		handlers.set(event, list);
+	},
+	registerProvider: (id: string, registration: ProviderRegistration) => {
+		registrations.push({ id, baseUrl: registration.baseUrl });
+	},
+	unregisterProvider: () => {},
+	registerCommand: () => {},
+};
+
+async function fire(event: string, sessionId: string): Promise<void> {
+	const ctx = { hasUI: true, sessionManager: { getSessionId: () => sessionId } } as ExtensionContext;
+	for (const handler of handlers.get(event) ?? []) await handler({ type: event }, ctx);
+}
+
+/** Health-probes the router: asserts the socket's state, never used as a wait. */
+async function alive(port: number): Promise<boolean> {
+	try {
+		const res = await fetch(`http://127.0.0.1:${port}/health`, { signal: AbortSignal.timeout(2_000) });
+		return res.ok;
+	} catch {
+		return false;
+	}
+}
+
+const portOfLatestRegistration = (): number =>
+	Number.parseInt(/:(\d+)\//.exec(registrations.at(-1)?.baseUrl ?? "")?.[1] ?? "0", 10);
+
+let home = "";
+let modelsYmlPath = "";
+let advertised = 0;
+let modelsYmlBefore = "";
+
+beforeAll(async () => {
+	home = mkdtempSync(join(tmpdir(), "embed-life-"));
+	const agentDir = join(home, "agent");
+	mkdirSync(agentDir, { recursive: true });
+	modelsYmlPath = join(agentDir, "models.yml");
+
+	// A port nothing listens on — the shape a restart leaves behind, since
+	// models.yml names the session the user just closed.
+	const probe = Bun.serve({ port: 0, hostname: "127.0.0.1", fetch: () => new Response("x") });
+	advertised = probe.port as number;
+	probe.stop(true);
+
+	writeFileSync(
+		modelsYmlPath,
+		`providers:
+    # BEGIN auto-model-router
+    auto-model-router:
+      baseUrl: http://127.0.0.1:${advertised}/v1
+      api: openai-completions
+      auth: none
+      models:
+        - id: auto
+          name: Auto (auto-model-router)
+    # END auto-model-router
+`,
+		"utf8",
+	);
+	modelsYmlBefore = readFileSync(modelsYmlPath, "utf8");
+
+	process.env.AUTO_MODEL_ROUTER_HOME = home;
+	process.env.AUTO_MODEL_ROUTER_DB = join(home, "router.db");
+	process.env.PI_CODING_AGENT_DIR = agentDir;
+	delete process.env.AUTO_MODEL_ROUTER_PORT;
+
+	const mod = (await import("../omp-extension/router-embed.ts")) as { default: (api: ExtensionAPI) => void };
+	mod.default(pi);
+	await fire("session_start", "session-one");
+});
+
+afterAll(async () => {
+	// Release the socket the way the extension intends: on a process signal.
+	process.emit("SIGTERM");
+	// Poll the real condition instead of guessing a duration; each probe yields.
+	for (let i = 0; i < 50 && (await alive(portOfLatestRegistration())); i++) {
+		/* keep probing until the socket stops answering */
+	}
+	delete process.env.PI_CODING_AGENT_DIR;
+	try {
+		rmSync(home, { recursive: true, force: true });
+	} catch {
+		/* SQLite may still hold the file on Windows; the temp dir is disposable */
+	}
+});
+
+describe("embedded router: port selection", () => {
+	test("adopts the port models.yml advertises, so omp's pre-resolved handle is valid", () => {
+		expect(portOfLatestRegistration()).toBe(advertised);
+	});
+
+	test("the adopted port actually serves", async () => {
+		expect(await alive(advertised)).toBe(true);
+	});
+
+	test("does NOT write its port into models.yml", () => {
+		// Persisting an ephemeral port makes it authoritative for the NEXT
+		// session's startup resolution, which is where the dead handle came from.
+		expect(readFileSync(modelsYmlPath, "utf8")).toBe(modelsYmlBefore);
+	});
+
+	test("publishes the port for subagents and the toast", () => {
+		const portFile = join(home, "embed.port");
+		expect(existsSync(portFile)).toBe(true);
+		expect(readFileSync(portFile, "utf8").trim()).toBe(String(advertised));
+	});
+});
+
+describe("embedded router: lifetime", () => {
+	test("registers NO session_shutdown teardown", () => {
+		// omp fires that from a throwaway host during provider refresh, so a
+		// teardown there kills a router the live session is still using.
+		expect(handlers.get("session_shutdown") ?? []).toHaveLength(0);
+	});
+
+	test("a session_shutdown leaves the router running", async () => {
+		await fire("session_shutdown", "session-one");
+		expect(await alive(advertised)).toBe(true);
+	});
+
+	test("a second session reuses the same port instead of rebinding", async () => {
+		// Rebinding would take a different port and orphan every handle omp had
+		// already resolved against the first one.
+		await fire("session_start", "session-two");
+		expect(portOfLatestRegistration()).toBe(advertised);
+		expect(await alive(advertised)).toBe(true);
+	});
+
+	test("each session still gets its own registration, so per-session tagging survives reuse", () => {
+		expect(registrations.length).toBeGreaterThanOrEqual(2);
+	});
+});