pi-mtplx 0.1.8 → 0.2.0

package/README.md

@@ -18,14 +18,14 @@ Restart Pi after installation.
 2. Download the model(s) you want, e.g. `mtplx install Youssofal/Qwen3.8-27B-MTPLX-Optimized-Quality`.
 3. Register a downloaded model with Pi via `/mtplx` → **Models** (or add it manually to `~/.pi/agent/mtplx-models.json`), then open `/model` to activate it.
 
-If you pick an MTPLX model that isn't installed, the server won't start — Pi will warn you. Use `mtplx list` to see what you've downloaded.
+`/mtplx` → **Models** only scans models installed through the local MTPLX CLI. It never downloads model weights and does not discover models available only on a configured remote endpoint. If you pick an MTPLX model that isn't installed locally, Pi cannot auto-start it. Use `mtplx list` to see what you've downloaded.
 
 ## What happens on first run
 
 Once installed, Pi automatically manages your MTPLX workflow:
 
-- **Model discovery** — `/mtplx` → **Models** scans what you've downloaded (`mtplx list`) and registers any of them with Pi. No model is bundled or pre-hardcoded: each Pi model id and its capabilities (context window, vision, reasoning) are derived live from the installed artifact, so models beyond the MTPLX stock set work too.
-- **Auto-start** — The MTPLX server starts when you switch to an `mtplx` model and shuts down cleanly when Pi exits.
+- **Model discovery** — `/mtplx` → **Models** scans what you've downloaded (`mtplx list`) and registers any of them with Pi. Each Pi model ID comes from MTPLX's own `quickstart --dry-run` plan, while capabilities (context window, vision, reasoning) are read from the installed artifact. Models beyond the MTPLX stock set work too.
+- **Auto-start** — The MTPLX server starts when you switch to an `mtplx` model. Auto-start and auto-shutdown are on by default and can be changed in `/mtplx`.
 - **Token speed** — A `⚡N.N tk/s` indicator appears in the footer showing the generation speed of the last assistant turn.
 
 ## Commands
@@ -34,10 +34,15 @@ Run `/mtplx` to open an interactive menu:
 
 | Option | What it does |
 | -------- | ------------- |
-| **Toggle (on/off)** | Start or stop the MTPLX server |
+| **Server** | Read-only status showing the model currently served at the configured endpoint, or that the endpoint is unavailable |
+| **Toggle (on/off)** | Start or stop an MTPLX server launched by this Pi session; separately managed servers are left running |
+| **API Key** | View a masked identifier for, or replace, the API key Pi uses for the local MTPLX server |
+| **Endpoint** | Set the OpenAI-compatible MTPLX base URL; useful for a separately started server |
 | **Fan Curves** | Set the thermal profile (`default`, `smart`, `max`) |
+| **Auto Shutdown Pi-Owned Server** | Choose whether Pi stops an MTPLX server it launched on `/quit` or a normal terminal-close shutdown (on by default); separately managed and remote servers are always left running |
+| **Auto Start** | Choose whether Pi starts or switches MTPLX when loading an MTPLX model (on by default) |
 | **SSD Session Cache** | Enable or disable MTPLX's SSD-backed session cache for subsequent server starts (on by default) |
-| **Models** | Register or unregister models — ✓ means registered (click to unregister), ✗ means available (click to register) |
+| **Models** | Register or unregister locally installed models — ✓ means registered (click to unregister), ✗ means available (click to register); it does not fetch or discover remote models |
 | **Uninstall** | Remove the `mtplx` provider from Pi's config |
 
 > **To activate a model** after registering or unregistering it, open **`/model`** (or `/scoped-models`). pi-mtplx writes Pi's config files immediately, but Pi loads them into memory on startup — so a **newly registered** MTPLX model only appears in `/model` after you restart Pi with **`/quit`** and relaunch it (`pi`). Unregistering/re-registering an existing model is picked up by opening `/model`, but a brand-new model id requires the restart.
@@ -46,21 +51,23 @@ Run `/mtplx` to open an interactive menu:
 
 ### Model registry
 
-Models are registered in `~/.pi/agent/mtplx-models.json`. Each entry maps a Pi model ID to an MTPLX artifact ref:
+Models are registered in `~/.pi/agent/mtplx-models.json`. Each entry maps MTPLX's canonical served model ID to an MTPLX artifact ref:
 
 ```json
 {
-  "mtplx-qwen3.8-27b-mtplx-optimized-quality": {
+  "mtplx-qwen38-27b-optimized-quality": {
     "ref": "Youssofal/Qwen3.8-27B-MTPLX-Optimized-Quality"
   }
 }
 ```
 
-Register new models via the `/mtplx` → **Models** menu, or add entries manually to this file. To activate a **newly added** model, restart Pi with `/quit` and relaunch it (`pi`), then open `/model` (or `/scoped-models`) — Pi reads this file at startup, so a brand-new model id won't show up in `/model` otherwise.
+Register new models via the `/mtplx` → **Models** menu; it asks MTPLX for the exact served ID. Opening this menu also migrates older ref-derived IDs to their canonical MTPLX names. Restart Pi after a migration or registration, then open `/model` (or `/scoped-models`) to activate the model.
+
+Local Forge outputs use their installed paths as artifact refs; downloaded repositories retain their `owner/model` refs. If an artifact cannot produce a serving plan, the menu shows MTPLX's error and skips that artifact while keeping the other models available.
 
 ### Fan mode
 
-Controls the thermal profile of your MTPLX server. Saved to `~/.pi/agent/mtplx-fanmode.json` and persists across restarts.
+Controls the thermal profile used when Pi starts an MTPLX server. Saved to `~/.pi/agent/mtplx-fanmode.json` and persists across restarts. When Pi discovers a server that is already running—whether it is local, remote, or Pi-owned—it does not automatically change the fan curve; if the server reports a supported curve that differs from Pi's saved default, Pi adopts that curve as its new default for future Pi-started servers. Choosing `/mtplx` → **Fan Curves** is an explicit override and changes the fan curve of the active healthy server.
 
 | Mode | Behavior |
 | ------ | ---------- |
@@ -68,12 +75,41 @@ Controls the thermal profile of your MTPLX server. Saved to `~/.pi/agent/mtplx-f
 | `smart` | Default — adaptive thermal management |
 | `max` | Maximum fan, fastest inference |
 
+### Server authentication
+
+pi-mtplx always uses an API key. It uses `providers.mtplx.apiKey` from `~/.pi/agent/models.json` when configured; otherwise it uses the local default `mtplx-local`. The resolved key is passed to Pi-managed MTPLX startup and sent with health, fan-control, and inference requests. If Pi already has a stored MTPLX API-key credential, pi-mtplx synchronizes it with this key at startup and whenever `/mtplx` → **API Key** saves a change. Restart Pi after changing the key so its in-memory inference provider reloads the configuration. The menu only shows a masked suffix for custom keys, never the full secret.
+
+### Standalone MTPLX
+
+To use an MTPLX server that you start outside Pi, set `/mtplx` → **Auto Start** to off, then set `/mtplx` → **Endpoint** to that server's OpenAI base URL (for example, `http://127.0.0.1:8001/v1`). Restart Pi after changing it. The endpoint, health checks, and fan controls all use `providers.mtplx.baseUrl`; its default is `http://127.0.0.1:8000/v1`. pi-mtplx never stops a server it did not launch, including on normal Pi shutdown.
+
+The served model ID must be the same ID selected in Pi. Before every MTPLX request, pi-mtplx verifies this through `/health` and reports both IDs if they differ. A separately managed server is never switched; start it with `--model-id <Pi model ID>` if its artifact's default identity differs. `/mtplx` → **Models** cannot discover a remote-only model, so add its exact ID to Pi's `mtplx` provider configuration manually before selecting it. Use the same API key in both the MTPLX launch command and `/mtplx` → **API Key**.
+
+### Auto start
+
+Auto Start is on by default. Pi first checks the selected endpoint before starting anything. If a healthy MTPLX server is already there and serves the selected model, Pi uses it without replacing it. If it serves a different model, Pi switches it only when that server was started by the current Pi session; separately managed and remote servers are left alone, and Pi blocks the request with both model IDs. Pi starts MTPLX only when no MTPLX server is available at a loopback endpoint. With Auto Start off, Pi performs the same health and model-match check but never starts or switches a server; it instead tells you to start MTPLX at the configured endpoint.
+
+### Auto shutdown Pi-owned server
+
+Auto shutdown is on by default. When enabled, pi-mtplx stops an MTPLX server it launched during the current Pi session when Pi exits via `/quit` or a normal terminal-close shutdown. It never stops a separately managed or remote server. Turn it off in `/mtplx` → **Auto Shutdown Pi-Owned Server** to leave the Pi-launched server running after Pi exits. It cannot handle abrupt termination such as `SIGKILL` or a power loss.
+
+## Remote endpoints
+
+I do not have a remote server, so I was unable to personally test how pi-mtplx interacts with remote endpoints.
+
+## Reporting bugs
+
+If you encounter a bug, please open a [GitHub Issue](../../issues) and let me know.
+
 ## Troubleshooting
 
 | Problem | Fix |
 | --------- | ----- |
 | **"MTPLX not started"** — Pi warns when you ask an MTPLX model to respond | Run `/mtplx` → **Toggle** to start the server |
 | **"No MTPLX models registered"** | Run `/mtplx` → **Models** to discover and register one |
+| **"MTPLX is already running ... but Pi requested ..."** | A manually started MTPLX server is serving a different, unmanaged model. Stop it in the MTPLX app, or find it with `lsof -nP -iTCP:8000 -sTCP:LISTEN` and run `kill -TERM <PID>`, then retry so Pi can start and manage the selected model. |
+| **"MTPLX rejected the API key"** | The running server requires a different key. Use `/mtplx` → **API Key** to enter its current key, then retry. |
+| **"Connection error"** with standalone MTPLX | Pi's configured endpoint has no listener. Set `/mtplx` → **Endpoint** to the standalone server's exact URL (including `/v1`), then restart Pi. |
 | **MTPLX startup timed out after 180s** | Run `mtplx status --deep` for MTPLX-side diagnostics (model validation, memory, thermal) |
 
 ## License

package/extensions/mtplx.ts

@@ -6,33 +6,69 @@
  * and shared helpers (src/utils.ts) into Pi via before_agent_start,
  * agent_end and session_shutdown, plus the /mtplx command.
  */
-import { acquire, release, stopServer } from "../src/mtplx-process.ts";
-import { getFanMode, health, setFanMode, setFanModeValue } from "../src/mtplx-client.ts";
+import { acquire, isOwnedByThisSession, release, stopServer, validateServer } from "../src/mtplx-process.ts";
+import { authenticationFailureMessage, getFanMode, health, healthProbe, setFanMode, setFanModeValue } from "../src/mtplx-client.ts";
 import { MTPLX_PROVIDER, manageModels, removeModel, removePiMtplxProvider } from "../src/model-discovery.ts";
-import { FAN_MODES, isMtplxModel, loadSsdSessionCache, saveFanMode, saveSsdSessionCache, type FanMode } from "../src/utils.ts";
+import { FAN_MODES, isMtplxModel, loadAutoShutdown, loadAutoStart, loadMtplxApiKey, loadMtplxEndpoint, loadSsdSessionCache, maskMtplxApiKey, mtplxEndpointFromBaseUrl, saveAutoShutdown, saveAutoStart, saveFanMode, saveMtplxApiKey, saveMtplxEndpoint, saveSsdSessionCache, syncMtplxStoredCredential, type FanMode } from "../src/utils.ts";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 
+function syncFanModeFromRunningServer(fanMode: string | undefined): { previous: FanMode; current: FanMode } | undefined {
+	if (typeof fanMode !== "string" || !(FAN_MODES as readonly string[]).includes(fanMode)) return;
+	const current = fanMode as FanMode;
+	const previous = getFanMode();
+	if (current === previous) return;
+	setFanModeValue(current);
+	saveFanMode(current);
+	return { previous, current };
+}
+
 export default function mtplxAutostart(pi: ExtensionAPI): void {
+	let lastServerNotice: string | undefined;
+
+	// Pi's runtime may already have loaded auth.json for this process, but
+	// syncing here ensures its next launch uses the same key as models.json.
+	syncMtplxStoredCredential();
+
 	pi.registerCommand("mtplx", {
-		description: "MTPLX — toggle the server, configure fan and SSD cache, manage models, or uninstall",
+		description: "MTPLX — view server status, configure its endpoint and lifecycle, manage models, or uninstall",
 		handler: async (_args, ctx) => {
 			const current = await health();
 			const status = current ? "on" : "off";
+			const fanSync = current ? syncFanModeFromRunningServer(current.fan_mode) : undefined;
+			const server = current
+				? `Server (serving: ${current.model})`
+				: "Server (unavailable — check endpoint or API key)";
 			await ctx.ui.setStatus("mtplx", `MTPLX: ${status}`);
+			if (fanSync) {
+				ctx.ui.notify(
+					`MTPLX was already running with fan curve ${fanSync.current}, while Pi's saved default was ${fanSync.previous}. Pi left the server unchanged and updated its saved default to ${fanSync.current}.`,
+					"info",
+				);
+			}
 			const ssdSessionCache = loadSsdSessionCache();
+			const autoShutdown = loadAutoShutdown();
+			const autoStart = loadAutoStart();
+			const apiKeyLabel = maskMtplxApiKey(loadMtplxApiKey());
+			const endpoint = loadMtplxEndpoint();
 			const topChoices = [
+				server,
 				`Toggle (${status})`,
+				`API Key (current: ${apiKeyLabel})`,
+				`Endpoint (current: ${endpoint.baseUrl})`,
 				`Fan Curves (current: ${getFanMode()})`,
+				`Auto Shutdown Pi-Owned Server (current: ${autoShutdown ? "on" : "off"})`,
+				`Auto Start (current: ${autoStart ? "on" : "off"})`,
 				`SSD Session Cache (current: ${ssdSessionCache ? "on" : "off"})`,
 				"Models",
 				"Uninstall (remove provider)",
 			];
 			const top = await ctx.ui.select("MTPLX", topChoices, undefined);
 			if (!top) return;
+			if (top === server) return;
 			if (top.startsWith("Toggle")) {
 				if (current) {
-					await stopServer();
-					ctx.ui.notify("MTPLX server stopped", "info");
+					const stopped = await stopServer();
+					ctx.ui.notify(stopped ? "MTPLX server stopped" : "MTPLX is managed separately; pi-mtplx left it running.", stopped ? "info" : "warning");
 				} else if (ctx.model && isMtplxModel(ctx.model)) {
 					await acquire(ctx.model.id);
 					ctx.ui.notify(`MTPLX server started (${ctx.model.id})`, "info");
@@ -41,6 +77,57 @@ export default function mtplxAutostart(pi: ExtensionAPI): void {
 				}
 				return;
 			}
+			if (top.startsWith("Endpoint")) {
+				const baseUrl = await ctx.ui.input("MTPLX endpoint", `Current: ${endpoint.baseUrl} — e.g. http://127.0.0.1:8001/v1`);
+				if (baseUrl === undefined) return;
+				try {
+					const url = new URL(baseUrl.trim());
+					if (url.protocol !== "http:" && url.protocol !== "https:") throw new Error("unsupported protocol");
+				} catch {
+					ctx.ui.notify("MTPLX endpoint was not updated. Enter an http(s) URL.", "error");
+					return;
+				}
+				if (mtplxEndpointFromBaseUrl(baseUrl).baseUrl === endpoint.baseUrl) {
+					ctx.ui.notify("MTPLX endpoint is unchanged.", "info");
+					return;
+				}
+				if (isOwnedByThisSession()) {
+					try {
+						await stopServer();
+					} catch (error) {
+						ctx.ui.notify(
+							`MTPLX endpoint was not changed because Pi could not stop its running server: ${error instanceof Error ? error.message : String(error)}`,
+							"error",
+						);
+						return;
+					}
+					ctx.ui.notify("Stopped Pi-owned MTPLX server before changing the endpoint.", "info");
+				}
+				if (!saveMtplxEndpoint(baseUrl)) {
+					ctx.ui.notify("MTPLX endpoint was not updated. Enter an http(s) URL.", "error");
+					return;
+				}
+				const probe = await healthProbe();
+				ctx.ui.notify(probe.health ? "MTPLX endpoint saved and verified. Restart Pi before inference uses it." : "MTPLX endpoint saved, but health verification failed. Check its host, port, and API key; then restart Pi.", probe.health ? "info" : "warning");
+				return;
+			}
+			if (top.startsWith("API Key")) {
+				const apiKey = await ctx.ui.input(`MTPLX API Key — current: ${apiKeyLabel}`, "Paste a custom key, or leave blank for the default (visible while typing)");
+				if (apiKey === undefined) return;
+				if (!saveMtplxApiKey(apiKey)) {
+					ctx.ui.notify("MTPLX API key was not updated. Check that models.json is valid.", "error");
+					return;
+				}
+				const probe = await healthProbe();
+				if (probe.health) {
+					ctx.ui.notify("MTPLX API key updated and verified against the running server. Restart Pi before inference requests use the new key.", "info");
+				} else if (probe.authenticationRejected) {
+					ctx.ui.notify(`MTPLX API key was saved, but verification failed: ${authenticationFailureMessage()} Restart Pi after correcting it.`, "error");
+				} else {
+					ctx.ui.notify("MTPLX API key updated. No running server was available to verify it. Restart Pi before inference requests use the new key.", "warning");
+				}
+				return;
+			}
 			if (top.startsWith("Fan Curves")) {
 				const choices = FAN_MODES.map((mode) => (mode === getFanMode() ? `${mode} (current)` : mode));
 				const picked = await ctx.ui.select("MTPLX — fan mode for autostart", choices, undefined);
@@ -49,10 +136,35 @@ export default function mtplxAutostart(pi: ExtensionAPI): void {
 				if (!(FAN_MODES as readonly string[]).includes(mode)) return;
 				setFanModeValue(mode);
 				saveFanMode(mode);
-				if (current) {
-					await setFanMode();
-				}
-				ctx.ui.notify(`MTPLX fan mode set to ${getFanMode()}`, "info");
+				if (current) await setFanMode();
+				ctx.ui.notify(
+					current
+						? `MTPLX fan curve set to ${getFanMode()} and saved as Pi's default for future Pi-started servers.`
+						: `MTPLX fan default saved as ${getFanMode()}. It applies the next time Pi starts MTPLX.`,
+					"info",
+				);
+			}
+			if (top.startsWith("Auto Shutdown")) {
+				const choices = [
+					autoShutdown ? "Off" : "Off (current)",
+					autoShutdown ? "On (current)" : "On",
+				];
+				const picked = await ctx.ui.select("MTPLX — stop the Pi-owned server when Pi exits", choices, undefined);
+				if (!picked) return;
+				const enabled = picked.startsWith("On");
+				saveAutoShutdown(enabled);
+				ctx.ui.notify(`Pi will ${enabled ? "stop its own MTPLX server" : "leave its own MTPLX server running"} when Pi exits.`, "info");
+			}
+			if (top.startsWith("Auto Start")) {
+				const choices = [
+					autoStart ? "Off" : "Off (current)",
+					autoStart ? "On (current)" : "On",
+				];
+				const picked = await ctx.ui.select("MTPLX — autostart MTPLX server when loading an MTPLX model", choices, undefined);
+				if (!picked) return;
+				const enabled = picked.startsWith("On");
+				saveAutoStart(enabled);
+				ctx.ui.notify(`MTPLX will ${enabled ? "autostart" : "not autostart"} when loading an MTPLX model.`, "info");
 			}
 			if (top.startsWith("SSD Session Cache")) {
 				const choices = [
@@ -84,13 +196,36 @@ export default function mtplxAutostart(pi: ExtensionAPI): void {
 		},
 	});
 
-	pi.on("before_agent_start", async (_event, ctx) => {
+	// This must run in `input`, not `before_agent_start`: Pi reports errors
+	// from before_agent_start but still sends the prompt to the provider.
+	pi.on("input", async (_event, ctx) => {
 		if (!ctx.model || !isMtplxModel(ctx.model)) return;
+		const autoStart = loadAutoStart();
 		try {
-			await acquire(ctx.model.id);
+			const server = autoStart
+				? await acquire(ctx.model.id)
+				: await validateServer(ctx.model.id);
+			const fanSync = server.alreadyRunning ? syncFanModeFromRunningServer(server.fanMode) : undefined;
+			ctx.ui.setStatus("mtplx", `MTPLX: ${server.model}`);
+			if (server.alreadyRunning) {
+				const message = autoStart
+					? `MTPLX is already running at ${server.endpoint.baseUrl}, serving ${JSON.stringify(server.model)}. Auto Start will not start MTPLX. The model selected in Pi matches; proceeding.`
+					: `MTPLX is available at ${server.endpoint.baseUrl}, serving ${JSON.stringify(server.model)}. It matches Pi's selected model; Auto Start is off, so Pi will proceed without managing the server.`;
+				const fanMessage = fanSync
+					? ` Its fan curve is ${fanSync.current}, while Pi's saved default was ${fanSync.previous}; Pi left the server unchanged and updated its saved default to ${fanSync.current}.`
+					: "";
+				const notice = message + fanMessage;
+				if (notice !== lastServerNotice) ctx.ui.notify(notice, "info");
+				lastServerNotice = notice;
+			} else {
+				lastServerNotice = undefined;
+			}
 		} catch (error) {
-			throw new Error(`MTPLX request blocked: ${error instanceof Error ? error.message : String(error)}`);
+			lastServerNotice = undefined;
+			ctx.ui.notify(`MTPLX request blocked: ${error instanceof Error ? error.message : String(error)}`, "error");
+			return { action: "handled" };
 		}
+		return { action: "continue" };
 	});
 
 	pi.on("agent_end", async (_event, ctx) => {
@@ -98,11 +233,11 @@ export default function mtplxAutostart(pi: ExtensionAPI): void {
 	});
 
 	pi.on("session_shutdown", async (event) => {
-		if (event.reason !== "quit") return;
+		if (event.reason !== "quit" || !loadAutoShutdown()) return;
 		try {
 			await stopServer();
 		} catch (error) {
-			console.error(`MTPLX cleanup on quit failed: ${error instanceof Error ? error.message : String(error)}`);
+			console.error(`MTPLX auto-shutdown failed: ${error instanceof Error ? error.message : String(error)}`);
 		}
 	});
 }

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-mtplx",
-	"version": "0.1.8",
+	"version": "0.2.0",
 	"description": "Zero-config MTPLX integration for Pi coding agent",
 	"keywords": [
 		"pi-package",

package/src/model-discovery.ts

@@ -1,20 +1,18 @@
 /**
  * MTPLX model registry: Pi model id → installed MTPLX artifact ref.
  *
- * The refs are the artifact identifiers reported by `mtplx models --json`;
- * `--model-id` makes /health and /v1/models report Pi's id.
+ * The refs are the artifact identifiers reported by `mtplx list --json`.
+ * MTPLX's `quickstart --dry-run --json` supplies the canonical served id used
+ * by Pi, /health, and /v1/models.
  *
  * The extension ships no model weights and hard-codes no model. The registry
  * holds only models registered through `/mtplx`, persisted to
- * ~/.pi/agent/mtplx-models.json; each Pi model id is derived live from the
- * artifact's ref (see modelIdFromRef) so models.json, enabledModels and this
- * registry always agree on the id.
+ * ~/.pi/agent/mtplx-models.json; models.json, enabledModels and this registry
+ * use the same MTPLX-supplied id.
  *
- * Pi's model catalog is the USER's own ~/.pi/agent/models.json — a provider
- * config that pre-existed this package (created by `mtplx start pi` / the
- * /mtplx UI). The canonical provider name is MTPLX's own `mtplx`
- * (PI_PROVIDER_ID in MTPLX's mtplx/pi.py); this extension reads and writes
- * only that provider entry and leaves every other provider untouched.
+ * Pi's model catalog lives in ~/.pi/agent/models.json. The canonical provider
+ * name is `mtplx`; this extension reads and writes only that provider entry
+ * and leaves every other provider untouched.
  */
 import { readFileSync, writeFileSync, mkdirSync, existsSync } from "node:fs";
 import { homedir } from "node:os";
@@ -22,12 +20,11 @@ import { join } from "node:path";
 import { execFile } from "node:child_process";
 import { promisify } from "node:util";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { modelIdFromRef, displayNameFromId } from "./utils.ts";
+import { commandError, displayNameFromId } from "./utils.ts";
 
 const execFileAsync = promisify(execFile);
 
-// Matches MTPLX's own PI_PROVIDER_ID ("mtplx") so the extension operates on
-// the exact provider block that `mtplx start pi` creates.
+// Keep the provider name consistent across model registration and lifecycle handling.
 export const MTPLX_PROVIDER = "mtplx";
 const MODELS_FILE = join(homedir(), ".pi", "agent", "mtplx-models.json");
 const SETTINGS_FILE = join(homedir(), ".pi", "agent", "settings.json");
@@ -101,6 +98,9 @@ export function disableModelInSettings(modelId: string): boolean {
 	}
 }
 type MtplxListedModel = { repo_id?: unknown; path?: unknown; name?: unknown };
+type ResolvedMtplxModel = { model: MtplxListedModel; ref: string; id: string };
+
+const servedIdCache = new Map<string, string>();
 
 export async function listMtplxModels(): Promise<MtplxListedModel[]> {
 	try {
@@ -114,7 +114,64 @@ export async function listMtplxModels(): Promise<MtplxListedModel[]> {
 }
 
 export function listedIdentity(model: MtplxListedModel): string {
-	return typeof model.repo_id === "string" && model.repo_id ? model.repo_id : typeof model.path === "string" ? model.path : "";
+	const repo = typeof model.repo_id === "string" ? model.repo_id : "";
+	// Forge artifacts have a bare directory name in repo_id. Quickstart treats
+	// that as a path relative to Pi's cwd, so use the installed path instead.
+	if (repo.includes("/")) return repo;
+	return typeof model.path === "string" && model.path ? model.path : repo;
+}
+
+/**
+ * MTPLX itself chooses the default served id. `mtplx list` reports artifact
+ * refs, not that id, so ask quickstart for its dry-run plan rather than trying
+ * to reproduce its naming rules in Pi.
+ */
+export async function servedModelId(ref: string): Promise<string> {
+	const cached = servedIdCache.get(ref);
+	if (cached) return cached;
+	try {
+		const { stdout } = await execFileAsync("mtplx", ["quickstart", "--model", ref, "--dry-run", "--json"], { timeout: 30_000 });
+		const id = servedModelIdFromDryRun(JSON.parse(stdout) as unknown);
+		servedIdCache.set(ref, id);
+		return id;
+	} catch (error) {
+		// MTPLX reports CLI errors as JSON on stdout, which execFile's default
+		// error.message omits. Keep that diagnostic visible in the Models menu.
+		const stdout = (error as { stdout?: string } | null)?.stdout;
+		let detail: unknown;
+		try {
+			const payload = JSON.parse(stdout ?? "") as { detail?: unknown; error?: unknown };
+			detail = payload.detail || payload.error;
+		} catch { /* Fall back to stderr / process error below. */ }
+		throw new Error(typeof detail === "string" ? detail : commandError(error), { cause: error });
+	}
+}
+
+/** A broken or unfinished artifact must not block management of other models. */
+export async function resolveInstalledModels(
+	installed: readonly MtplxListedModel[],
+	resolveId: (ref: string) => Promise<string> = servedModelId,
+): Promise<{ resolved: ResolvedMtplxModel[]; failures: string[] }> {
+	const resolved: ResolvedMtplxModel[] = [];
+	const failures: string[] = [];
+	for (const model of installed) {
+		const ref = listedIdentity(model);
+		if (!ref) continue;
+		try {
+			resolved.push({ model, ref, id: await resolveId(ref) });
+		} catch (error) {
+			failures.push(`${ref}: ${error instanceof Error ? error.message : String(error)}`);
+		}
+	}
+	return { resolved, failures };
+}
+
+/** Parse the model id from MTPLX's machine-readable quickstart plan. */
+export function servedModelIdFromDryRun(plan: unknown): string {
+	if (typeof plan !== "object" || plan === null || typeof (plan as { model_id?: unknown }).model_id !== "string" || !(plan as { model_id: string }).model_id) {
+		throw new Error("MTPLX did not return a served model id");
+	}
+	return (plan as { model_id: string }).model_id;
 }
 
 /**
@@ -135,6 +192,68 @@ function catalogModelIds(): string[] {
 	}
 }
 
+/**
+ * Replace legacy ref-derived ids with MTPLX's own served ids. This runs only
+ * from the explicit Models menu, where the resulting Pi restart is expected.
+ */
+function migrateRegisteredModelIds(models: readonly ResolvedMtplxModel[]): string[] {
+	const servedIds = new Map(models.map(({ ref, id }) => [ref, id]));
+	const changes = new Map<string, string>();
+	for (const [oldId, entry] of Object.entries(MTPLX_MODELS)) {
+		const newId = servedIds.get(entry.ref);
+		if (!newId || newId === oldId) continue;
+		const collision = MTPLX_MODELS[newId];
+		if (collision && collision.ref !== entry.ref) {
+			console.warn(`pi-mtplx: cannot migrate ${oldId} to ${newId}; that id is already assigned to a different artifact.`);
+			continue;
+		}
+		changes.set(oldId, newId);
+	}
+	if (changes.size === 0) return [];
+
+	try {
+		const modelsJsonPath = join(homedir(), ".pi", "agent", "models.json");
+		const catalog = JSON.parse(readFileSync(modelsJsonPath, "utf8")) as { providers?: Record<string, unknown> };
+		const provider = catalog.providers?.[MTPLX_PROVIDER] as { models?: Record<string, unknown>[] } | undefined;
+		if (provider && Array.isArray(provider.models)) {
+			const seen = new Set<string>();
+			provider.models = provider.models.flatMap((model) => {
+				const id = typeof model.id === "string" ? model.id : undefined;
+				const replacement = id ? changes.get(id) : undefined;
+				const nextId = replacement ?? id;
+				if (!nextId || seen.has(nextId)) return [];
+				seen.add(nextId);
+				return [{ ...model, id: nextId, ...(replacement ? { name: displayNameFromId(nextId) } : {}) }];
+			});
+		}
+		writeFileSync(modelsJsonPath, JSON.stringify(catalog, null, 2) + "\n");
+
+		if (existsSync(SETTINGS_FILE)) {
+			const settings = JSON.parse(readFileSync(SETTINGS_FILE, "utf8")) as { enabledModels?: unknown };
+			if (Array.isArray(settings.enabledModels)) {
+				settings.enabledModels = [...new Set(settings.enabledModels.map((entry) => {
+					if (typeof entry !== "string") return entry;
+					const id = entry.startsWith(`${MTPLX_PROVIDER}/`) ? entry.slice(MTPLX_PROVIDER.length + 1) : undefined;
+					return id && changes.has(id) ? `${MTPLX_PROVIDER}/${changes.get(id)}` : entry;
+				}))];
+				writeFileSync(SETTINGS_FILE, JSON.stringify(settings, null, 2) + "\n");
+			}
+		}
+
+		for (const [oldId, newId] of changes) {
+			const entry = MTPLX_MODELS[oldId];
+			if (!entry) continue;
+			MTPLX_MODELS[newId] = entry;
+			delete MTPLX_MODELS[oldId];
+		}
+		saveRegisteredModels();
+		return [...changes].map(([oldId, newId]) => `${oldId} → ${newId}`);
+	} catch (error) {
+		console.error(`pi-mtplx could not migrate MTPLX model ids: ${error instanceof Error ? error.message : String(error)}`);
+		return [];
+	}
+}
+
 type ModelProfile = {
 	contextWindow: number;
 	maxTokens: number;
@@ -197,13 +316,17 @@ export async function manageModels(ctx: ExtensionContext): Promise<boolean> {
 		ctx.ui.notify("No MTPLX models found. Install one with `mtplx install`.", "warning");
 		return false;
 	}
+	const { resolved, failures } = await resolveInstalledModels(installed);
+	if (failures.length > 0) {
+		ctx.ui.notify(`Skipped ${failures.length} MTPLX model(s) whose served ID could not be read:\n${failures.join("\n")}`, "warning");
+	}
+	if (resolved.length === 0) return false;
+	const migrated = migrateRegisteredModelIds(resolved);
+	if (migrated.length > 0) {
+		ctx.ui.notify(`Updated registered model IDs to MTPLX's canonical names. Restart Pi before selecting them.`, "info");
+	}
 	const catalog = new Set(catalogModelIds());
-	const choices = installed.map((model) => {
-		const ref = listedIdentity(model);
-		// Each artifact maps to exactly one Pi model id, derived live from its ref.
-		// No hard-coded ids — the same id is written to models.json, enabledModels and
-		// the registry so the three never drift apart.
-		const id = modelIdFromRef(ref);
+	const choices = resolved.map(({ ref, id }) => {
 		// ✓ means registered in Pi's catalog (models.json); ✗ means just installed in MTPLX.
 		const mark = catalog.has(id) ? "✓" : "✗";
 		return `${mark} ${id} — ${ref}`;
@@ -212,9 +335,10 @@ export async function manageModels(ctx: ExtensionContext): Promise<boolean> {
 	const picked = await ctx.ui.select("MTPLX models — ✓ registered in Pi · ✗ available in MTPLX — registering a new model needs a Pi restart (/quit → pi) to show up in /model", choices, undefined);
 	if (!picked || picked === "Cancel") return false;
 	const ref = picked.split(" — ").slice(1).join(" — ");
-	const modelId = modelIdFromRef(ref);
 	// The installed-model record for the chosen artifact (used to read its live metadata).
-	const chosen = installed.find((m) => listedIdentity(m) === ref);
+	const selected = resolved.find((entry) => entry.ref === ref);
+	if (!selected) return false;
+	const { model: chosen, id: modelId } = selected;
 
 	if (catalog.has(modelId)) {
 		// Unregister: remove from models.json, enabledModels, and the id→ref registry.
@@ -349,4 +473,4 @@ export function removePiMtplxProvider(): boolean {
 		console.error(`pi-mtplx uninstall could not update models.json: ${error instanceof Error ? error.message : String(error)}`);
 		return false;
 	}
-}
+}

package/src/mtplx-client.ts

@@ -1,12 +1,13 @@
 /**
- * Minimal HTTP client for the locally-running MTPLX server (OpenAI-compatible
- * endpoint on 127.0.0.1:8000). Only touches endpoints that MTPLX exposes:
+ * Minimal HTTP client for the configured MTPLX OpenAI-compatible endpoint.
+ * Only touches endpoints that MTPLX exposes:
  *   GET  /health
  *   POST /v1/mtplx/thermal/fan_mode
  */
-import { HOST, PORT, type FanMode, loadFanMode } from "./utils.ts";
+import { type FanMode, type MtplxEndpoint, loadFanMode, loadMtplxEndpoint, loadResolvedMtplxApiKey } from "./utils.ts";
 
 export type Health = { ok: true; model: string; model_path?: string; fan_mode?: string };
+export type HealthProbe = { health: Health | undefined; authenticationRejected?: true };
 
 // Current fan curve for autostart and live updates. Loaded from disk so the
 // choice survives Pi restarts (see FANMODE_FILE in utils.ts).
@@ -20,37 +21,66 @@ export function setFanModeValue(mode: FanMode): void {
 	fanMode = mode;
 }
 
-export async function health(): Promise<Health | undefined> {
+/**
+ * Pi's own model requests use the configured key, or pi-mtplx's deterministic
+ * local fallback. Mirror that header for every lifecycle endpoint.
+ */
+function authRequest(): { headers: Record<string, string> } {
+	return { headers: { authorization: `Bearer ${loadResolvedMtplxApiKey()}` } };
+}
+
+export function authenticationFailureMessage(): string {
+	return "MTPLX rejected Pi's API key. Update it via /mtplx → API Key.";
+}
+
+export async function healthProbe(endpoint = loadMtplxEndpoint()): Promise<HealthProbe> {
 	const controller = new AbortController();
 	const timeout = setTimeout(() => controller.abort(), 1_500);
 	try {
-		const response = await fetch(`http://${HOST}:${PORT}/health`, { signal: controller.signal });
-		if (!response.ok) return undefined;
+		const auth = authRequest();
+		const response = await fetch(`${endpoint.origin}/health`, {
+			headers: auth.headers,
+			signal: controller.signal,
+		});
+		if (response.status === 401 || response.status === 403) {
+			return { health: undefined, authenticationRejected: true };
+		}
+		if (!response.ok) return { health: undefined };
 		const body = (await response.json()) as { ok?: unknown; model?: unknown; model_path?: unknown; fan_mode?: unknown };
-		if (body.ok !== true || typeof body.model !== "string") return undefined;
+		if (body.ok !== true || typeof body.model !== "string") return { health: undefined };
 		return {
-			ok: true,
-			model: body.model,
-			model_path: typeof body.model_path === "string" ? body.model_path : undefined,
-			fan_mode: typeof body.fan_mode === "string" ? body.fan_mode : undefined,
+			health: {
+				ok: true,
+				model: body.model,
+				model_path: typeof body.model_path === "string" ? body.model_path : undefined,
+				fan_mode: typeof body.fan_mode === "string" ? body.fan_mode : undefined,
+			},
 		};
 	} catch {
-		return undefined;
+		return { health: undefined };
 	} finally {
 		clearTimeout(timeout);
 	}
 }
 
+export async function health(endpoint?: MtplxEndpoint): Promise<Health | undefined> {
+	return (await healthProbe(endpoint)).health;
+}
+
 export async function setFanMode(): Promise<void> {
 	const controller = new AbortController();
 	const timeout = setTimeout(() => controller.abort(), 10_000);
 	try {
-		const response = await fetch(`http://${HOST}:${PORT}/v1/mtplx/thermal/fan_mode`, {
+		const auth = authRequest();
+		const response = await fetch(`${loadMtplxEndpoint().origin}/v1/mtplx/thermal/fan_mode`, {
 			method: "POST",
-			headers: { "content-type": "application/json" },
+			headers: { "content-type": "application/json", ...auth.headers },
 			body: JSON.stringify({ mode: fanMode }),
 			signal: controller.signal,
 		});
+		if (response.status === 401 || response.status === 403) {
+			throw new Error(authenticationFailureMessage());
+		}
 		const body = (await response.json()) as { verified?: unknown; current_mode?: unknown; error?: unknown };
 		if (!response.ok || body.verified !== true || body.current_mode !== fanMode) {
 			throw new Error(`MTPLX fan mode could not be set to ${fanMode}: ${typeof body.error === "string" ? body.error : "unverified response"}`);
@@ -60,4 +90,4 @@ export async function setFanMode(): Promise<void> {
 	} finally {
 		clearTimeout(timeout);
 	}
-}
+}

package/src/mtplx-process.ts

@@ -7,21 +7,18 @@
  *     Pi's model id — a positive fingerprint that a given healthy server on
  *     the port is one this extension owns.
  *   - Cleanup is therefore precise: a server only gets shut down if the
- *     extension spawned it in this session, or if `/health` positively
- *     identifies it as an MTPLX server serving one of the extension's model
- *     ids. A foreign service on the port is never touched, and a manually
- *     started MTPLX server (no matching --model-id fingerprint) is never
- *     killed — it is simply served from as-is.
- *   - Stop goes through `mtplx stop --host --port --json` (MTPLX's own
- *     graceful-stop mechanism: SIGTERM → grace → SIGKILL), which targets the
- *     server answering on that exact host:port, not arbitrary processes.
+ *     extension spawned it in this Pi session. A manually or separately
+ *     started server is never killed, even when it serves a registered model.
+ *   - A server launched by this Pi session is stopped through its own detached
+ *     process group. This avoids relying on a second `mtplx` CLI invocation
+ *     during shutdown.
  */
 import { spawn, execFile } from "node:child_process";
 import { promisify } from "node:util";
 import type { ChildProcess } from "node:child_process";
-import { getFanMode, health, setFanMode, type Health } from "./mtplx-client.ts";
+import { authenticationFailureMessage, getFanMode, health, healthProbe, setFanMode } from "./mtplx-client.ts";
 import { MTPLX_MODELS } from "./model-discovery.ts";
-import { HOST, PORT, READY_TIMEOUT_MS, POLL_MS, loadSsdSessionCache, sleep, commandError, portIsOccupied } from "./utils.ts";
+import { READY_TIMEOUT_MS, POLL_MS, loadMtplxEndpoint, loadResolvedMtplxApiKey, loadSsdSessionCache, sleep, commandError, portIsOccupied, type MtplxEndpoint } from "./utils.ts";
 
 const execFileAsync = promisify(execFile);
 
@@ -42,84 +39,132 @@ export function startupError(cause: Error): Error {
 }
 
 /**
- * Shut down the server on HOST:PORT — but only if we can positively identify
- * it as MTPLX (via /health). A healthy, non-MTPLX listener is an error, never
- * a kill target.
+ * Shut down the server at the configured endpoint, but only when this Pi
+ * session launched it. A healthy external listener is never a kill target.
  */
-export async function stopServer(): Promise<void> {
-	const current = await health();
+export async function stopServer(): Promise<boolean> {
+	const endpoint = loadMtplxEndpoint();
+	const probe = await healthProbe();
+	const current = probe.health;
 	if (!current) {
-		if (await portIsOccupied()) {
-			throw new Error(`MTPLX cannot use ${HOST}:${PORT}: another, non-MTPLX service is listening there.`);
+		if (await portIsOccupied(endpoint)) {
+			if (probe.authenticationRejected) throw new Error(authenticationFailureMessage());
+			throw new Error(`MTPLX cannot use ${endpoint.host}:${endpoint.port}: another, non-MTPLX service is listening there.`);
 		}
 		// Nothing listening (or not MTPLX): nothing to stop. Drop any stale handle.
 		ownedChild = undefined;
-		return;
+		return true;
 	}
 
-	// /health says an MTPLX server is answering here.
-	if (!isOwnedByThisSession() && !currentModelIsOurs(current.model)) {
-		// An MTPLX server the extension does not own (e.g. started manually by
-		// the user). MTPLX exposes no reliable way to distinguish it from a
-		// Pi-managed one at the port level beyond --model-id, so do not kill
-		// it. If a Pi-owned one is still needed, the user can run /mtplx →
-		// Toggle.
+	if (!isOwnedByThisSession()) {
+	// A separately started server can deliberately use the same model id as a
+	// Pi-managed one, so model identity is not ownership.
 		console.warn(
-			`pi-mtplx: leaving MTPLX server on ${HOST}:${PORT} (model ${JSON.stringify(current.model)}) untouched — not owned by this Pi session. Use /mtplx → Toggle to stop it.`,
+			`pi-mtplx: leaving MTPLX server on ${endpoint.host}:${endpoint.port} (model ${JSON.stringify(current.model)}) untouched — not owned by this Pi session.`,
 		);
-		return;
+		return false;
 	}
 
-	try {
-		await execFileAsync("mtplx", ["stop", "--host", HOST, "--port", String(PORT), "--json"], { timeout: 20_000 });
-	} catch (error) {
-		throw new Error(`MTPLX shutdown failed: ${commandError(error)}`);
+	let stopError: unknown;
+	const childPid = ownedChild?.pid;
+	let signalledOwnedProcess = false;
+	if (isOwnedByThisSession() && childPid) {
+		try {
+			// `detached: true` gives this child its own POSIX process group. Signal
+			// that group so a quickstart wrapper and its server are stopped together.
+			process.kill(-childPid, "SIGTERM");
+			signalledOwnedProcess = true;
+		} catch (error) {
+			// The child may have already exited while its server survived. In that
+			// case, fall through to MTPLX's host-and-port stop command below.
+			if (!(error && typeof error === "object" && "code" in error && error.code === "ESRCH")) {
+				stopError = error;
+			}
+		}
+	}
+
+	if (!signalledOwnedProcess) {
+		try {
+			if (!endpoint.isLoopback) throw new Error("MTPLX endpoint is remote; pi-mtplx will not stop a server it did not launch.");
+			await execFileAsync("mtplx", ["stop", "--host", endpoint.host, "--port", String(endpoint.port), "--json"], { timeout: 20_000 });
+			stopError = undefined;
+		} catch (error) {
+			// Some MTPLX CLI failures occur after it has already signalled the server.
+			// Confirm the listener state before reporting shutdown as failed.
+			stopError = error;
+		}
 	}
-	ownedChild = undefined;
 
 	const deadline = Date.now() + 20_000;
 	while (Date.now() < deadline) {
-		if (!(await health())) return;
+		if (!(await portIsOccupied(endpoint))) {
+			ownedChild = undefined;
+			return true;
+		}
 		await sleep(POLL_MS);
 	}
-	throw new Error(`MTPLX shutdown timed out; ${HOST}:${PORT} is still healthy.`);
+	if (stopError) throw new Error(`MTPLX shutdown failed: ${commandError(stopError)}`);
+	if (signalledOwnedProcess) {
+		throw new Error(`MTPLX shutdown timed out after signalling Pi's server process; ${endpoint.host}:${endpoint.port} is still occupied.`);
+	}
+	throw new Error(`MTPLX shutdown timed out; ${endpoint.host}:${endpoint.port} is still occupied.`);
+}
+
+function modelMismatchError(runningModel: string, requestedModel: string, endpoint: MtplxEndpoint): Error {
+	return new Error(
+		`MTPLX is already running at ${endpoint.baseUrl}, serving ${JSON.stringify(runningModel)}, but Pi has ${JSON.stringify(requestedModel)} selected. ` +
+		`Pi will not replace the running server. Select ${JSON.stringify(runningModel)} in Pi, or restart MTPLX with \`--model-id ${requestedModel}\`, then try again.`,
+	);
 }
 
-/** Positive ownership fingerprint: /health model id belongs to this extension's registry. */
-function currentModelIsOurs(model: string): boolean {
-	return Object.keys(MTPLX_MODELS).includes(model);
+function autoStartDisabledUnavailableError(endpoint: MtplxEndpoint): Error {
+	return new Error(
+		`MTPLX is unavailable at ${endpoint.baseUrl}. Auto Start is off, so Pi will not start a server. ` +
+		"Start MTPLX at this endpoint or enable Auto Start, then try again.",
+	);
 }
 
+export type ReadyServer = {
+	endpoint: MtplxEndpoint;
+	model: string;
+	fanMode?: string;
+	alreadyRunning: boolean;
+};
+
 /**
  * Spawn the MTPLX server for a registered model and wait until /health
  * confirms it serves exactly that model. The child is spawned detached so it
  * outlives Pi's event-loop teardown, and unref'd so it never blocks Pi exit.
  */
 export async function startServer(modelId: string): Promise<void> {
+	const endpoint = loadMtplxEndpoint();
 	const configured = MTPLX_MODELS[modelId];
 	if (!configured) {
 		throw new Error(`MTPLX model ${JSON.stringify(modelId)} is not mapped to an installed MTPLX artifact. Update the pi-mtplx model registry after adding it to Pi.`);
 	}
 
+	const args = [
+		"quickstart",
+		"--model",
+		configured.ref,
+		"--model-id",
+		modelId,
+		"--fan-mode",
+		getFanMode(),
+		"--host",
+		endpoint.host,
+		"--port",
+		String(endpoint.port),
+		// Explicitly pass the user's persisted choice; this extension defaults it to on.
+		"--ssd-session-cache",
+		loadSsdSessionCache() ? "on" : "off",
+	];
+	args.push("--api-key", loadResolvedMtplxApiKey());
+
 	let exited: Error | undefined;
 	const child = spawn(
 		"mtplx",
-		[
-			"quickstart",
-			"--model",
-			configured.ref,
-			"--model-id",
-			modelId,
-			"--fan-mode",
-			getFanMode(),
-			"--host",
-			HOST,
-			"--port",
-			String(PORT),
-			// Explicitly pass the user's persisted choice; this extension defaults it to on.
-			"--ssd-session-cache",
-			loadSsdSessionCache() ? "on" : "off",
-		],
+		args,
 		{ detached: true, stdio: "ignore" },
 	);
 	ownedChild = child;
@@ -146,30 +191,58 @@ export async function startServer(modelId: string): Promise<void> {
 }
 
 /**
- * Ensure the server on HOST:PORT serves `modelId`, transparently switching
- * models: stop the current (identified) server, then start the requested one.
+ * Confirm that the configured endpoint is healthy and serves `modelId`.
+ * Unlike `ensureServer`, this never starts, stops, or switches a server.
  */
-export async function ensureServer(modelId: string): Promise<void> {
-	const current = await health();
+export async function validateServer(modelId: string, endpoint = loadMtplxEndpoint()): Promise<ReadyServer> {
+	const probe = await healthProbe(endpoint);
+	const current = probe.health;
 	if (current?.model === modelId) {
-		if (current.fan_mode !== getFanMode()) await setFanMode();
-		return;
+		return { endpoint, model: current.model, fanMode: current.fan_mode, alreadyRunning: true };
+	}
+	if (current) {
+		throw modelMismatchError(current.model, modelId, endpoint);
+	}
+	if (probe.authenticationRejected) throw new Error(authenticationFailureMessage());
+	throw autoStartDisabledUnavailableError(endpoint);
+}
+
+/**
+ * Ensure the configured endpoint serves `modelId`. A healthy separately
+ * managed server is authoritative: Pi reuses a matching one and never
+ * replaces a different model. A server started by this Pi session can be
+ * stopped and switched to the selected model. Pi starts a server only when
+ * the endpoint is unavailable or after switching its own server.
+ */
+export async function ensureServer(modelId: string, endpoint = loadMtplxEndpoint()): Promise<ReadyServer> {
+	const probe = await healthProbe(endpoint);
+	const current = probe.health;
+	if (current) {
+		if (current.model === modelId) {
+			return { endpoint, model: current.model, fanMode: current.fan_mode, alreadyRunning: true };
+		}
+		if (!isOwnedByThisSession()) throw modelMismatchError(current.model, modelId, endpoint);
+		await stopServer();
+	}
+	if (await portIsOccupied(endpoint)) {
+		if (probe.authenticationRejected) throw new Error(authenticationFailureMessage());
+		throw new Error(`MTPLX cannot start because ${endpoint.host}:${endpoint.port} is occupied by a non-MTPLX service.`);
 	}
-	if (current) await stopServer();
-	else if (await portIsOccupied()) {
-		throw new Error(`MTPLX cannot start because ${HOST}:${PORT} is occupied by a non-MTPLX service.`);
+	if (!endpoint.isLoopback) {
+		throw new Error(`MTPLX endpoint ${endpoint.baseUrl} is unavailable. pi-mtplx only auto-starts loopback servers; start this endpoint separately.`);
 	}
 	await startServer(modelId);
 	await setFanMode();
+	return { endpoint, model: modelId, fanMode: getFanMode(), alreadyRunning: false };
 }
 
-let transition: Promise<void> | undefined;
+let transition: Promise<ReadyServer> | undefined;
 let transitionModel: string | undefined;
 let activeModel: string | undefined;
 let activeAgents = 0;
 let idleWaiters: Array<() => void> = [];
 
-export function ensureOnce(modelId: string): Promise<void> {
+export function ensureOnce(modelId: string): Promise<ReadyServer> {
 	if (transition) {
 		if (transitionModel === modelId) return transition;
 		return transition.then(() => ensureOnce(modelId));
@@ -182,7 +255,7 @@ export function ensureOnce(modelId: string): Promise<void> {
 	return transition;
 }
 
-export async function acquire(modelId: string): Promise<void> {
+export async function acquire(modelId: string): Promise<ReadyServer> {
 	// Never replace a model while an already-admitted MTPLX agent is running.
 	// Same-model requests share both this lease and any in-flight start Promise.
 	while (activeAgents > 0 && activeModel !== modelId) {
@@ -194,7 +267,7 @@ export async function acquire(modelId: string): Promise<void> {
 	activeModel = modelId;
 	activeAgents += 1;
 	try {
-		await ensureOnce(modelId);
+		return await ensureOnce(modelId);
 	} catch (error) {
 		release();
 		throw error;

package/src/utils.ts

@@ -1,7 +1,7 @@
 /**
  * Shared constants and small helpers for the pi-mtplx extension.
  */
-import { readFileSync, writeFileSync, mkdirSync } from "node:fs";
+import { existsSync, readFileSync, writeFileSync, mkdirSync } from "node:fs";
 import { createConnection } from "node:net";
 import { homedir } from "node:os";
 import { join } from "node:path";
@@ -10,10 +10,55 @@ export const HOST = "127.0.0.1";
 export const PORT = 8000;
 export const READY_TIMEOUT_MS = 180_000;
 export const POLL_MS = 500;
+export const PI_MODELS_FILE = join(homedir(), ".pi", "agent", "models.json");
+export const PI_AUTH_FILE = join(homedir(), ".pi", "agent", "auth.json");
 
 export type FanMode = "default" | "smart" | "max";
 export const FAN_MODES: readonly FanMode[] = ["default", "smart", "max"];
 export const DEFAULT_SSD_SESSION_CACHE = true;
+export const DEFAULT_AUTO_SHUTDOWN = true;
+export const DEFAULT_AUTO_START = true;
+export const DEFAULT_MTPLX_API_KEY = "mtplx-local";
+
+export type MtplxEndpoint = {
+	baseUrl: string;
+	origin: string;
+	host: string;
+	port: number;
+	isLoopback: boolean;
+};
+
+/** Resolve the endpoint from Pi's provider configuration, with a local default. */
+export function mtplxEndpointFromBaseUrl(baseUrl: unknown): MtplxEndpoint {
+	const fallback = `http://${HOST}:${PORT}/v1`;
+	let url: URL;
+	try {
+		url = new URL(typeof baseUrl === "string" && baseUrl.trim() ? baseUrl.trim() : fallback);
+		if (url.protocol !== "http:" && url.protocol !== "https:") throw new Error("unsupported protocol");
+	} catch {
+		url = new URL(fallback);
+	}
+	// URL.hostname retains brackets around IPv6 literals ("[::1]"). The
+	// MTPLX CLI and node:net both expect the bare address, and ::1 is local.
+	const host = url.hostname.replace(/^\[|\]$/g, "");
+	const port = Number(url.port || (url.protocol === "https:" ? 443 : 80));
+	return {
+		baseUrl: url.toString().replace(/\/$/, ""),
+		origin: url.origin,
+		host,
+		port,
+		isLoopback: host === "127.0.0.1" || host === "::1" || host.toLowerCase() === "localhost",
+	};
+}
+
+export function loadMtplxEndpoint(): MtplxEndpoint {
+	try {
+		const catalog = JSON.parse(readFileSync(PI_MODELS_FILE, "utf8")) as { providers?: { mtplx?: { baseUrl?: unknown } } };
+		return mtplxEndpointFromBaseUrl(catalog.providers?.mtplx?.baseUrl);
+	} catch {
+		return mtplxEndpointFromBaseUrl(undefined);
+	}
+}
 
 // Fan mode ("fan curve") applied at autostart, live-updated from `/mtplx` while the
 // server runs. Persisted to disk so the choice survives Pi restarts: `fanMode` is a
@@ -66,17 +111,188 @@ export function saveSsdSessionCache(enabled: boolean): void {
 	}
 }
 
-// Derive Pi's model id from an artifact ref: "Youssofal/Qwen3.8-27B-MTPLX-Optimized-Quality"
-// → "mtplx-qwen38-27b-optimized-quality" (same scheme as the built-in entry).
-export function modelIdFromRef(ref: string): string {
-	const slug = ref
-		.split("/")
-		.pop() ?? ref
-		.replace(/^mtplx/i, "")
-		.replace(/[^a-zA-Z0-9]+/g, "-")
-		.replace(/-+/g, "-")
-		.replace(/^-|-$/g, "");
-	return `mtplx-${slug || "model"}`.toLowerCase();
+// Whether Pi should stop the server during a normal process shutdown. Kept
+// separate from the server's own settings so users can deliberately leave a
+// loaded model available after Pi exits.
+export const AUTO_SHUTDOWN_FILE = join(homedir(), ".pi", "agent", "mtplx-auto-shutdown.json");
+export const AUTO_START_FILE = join(homedir(), ".pi", "agent", "mtplx-auto-start.json");
+
+export function loadAutoStart(): boolean {
+	try {
+		const parsed = JSON.parse(readFileSync(AUTO_START_FILE, "utf8")) as { enabled?: unknown };
+		if (typeof parsed.enabled === "boolean") return parsed.enabled;
+	} catch {
+		// missing or corrupt file → fall back to the default
+	}
+	return DEFAULT_AUTO_START;
+}
+
+export function saveAutoStart(enabled: boolean): void {
+	try {
+		mkdirSync(join(homedir(), ".pi", "agent"), { recursive: true });
+		writeFileSync(AUTO_START_FILE, JSON.stringify({ enabled }, null, 2) + "\n");
+	} catch (error) {
+		console.error(`MTPLX could not persist auto-start preference: ${error instanceof Error ? error.message : String(error)}`);
+	}
+}
+
+export function loadAutoShutdown(): boolean {
+	try {
+		const parsed = JSON.parse(readFileSync(AUTO_SHUTDOWN_FILE, "utf8")) as { enabled?: unknown };
+		if (typeof parsed.enabled === "boolean") return parsed.enabled;
+	} catch {
+		// missing or corrupt file → default to cleaning up the Pi-managed server
+	}
+	return DEFAULT_AUTO_SHUTDOWN;
+}
+
+export function saveAutoShutdown(enabled: boolean): void {
+	try {
+		mkdirSync(join(homedir(), ".pi", "agent"), { recursive: true });
+		writeFileSync(AUTO_SHUTDOWN_FILE, JSON.stringify({ enabled }, null, 2) + "\n");
+	} catch (error) {
+		console.error(`MTPLX could not persist auto-shutdown preference: ${error instanceof Error ? error.message : String(error)}`);
+	}
+}
+
+/** Read the provider-level key Pi uses for requests to the local MTPLX server. */
+export function mtplxApiKeyFromCatalog(catalog: unknown): string | undefined {
+	if (typeof catalog !== "object" || catalog === null) return undefined;
+	const providers = (catalog as { providers?: unknown }).providers;
+	if (typeof providers !== "object" || providers === null) return undefined;
+	const provider = (providers as { mtplx?: unknown }).mtplx;
+	if (typeof provider !== "object" || provider === null) return undefined;
+	const apiKey = (provider as { apiKey?: unknown }).apiKey;
+	return typeof apiKey === "string" && apiKey.trim() ? apiKey.trim() : undefined;
+}
+
+/**
+ * Lifecycle calls reread this on every request. Pi's inference provider has
+ * its own in-memory catalog, so a key changed here still requires a Pi restart
+ * before model requests use it.
+ */
+export function loadMtplxApiKey(): string | undefined {
+	try {
+		return mtplxApiKeyFromCatalog(JSON.parse(readFileSync(PI_MODELS_FILE, "utf8")) as unknown);
+	} catch {
+		return undefined;
+	}
+}
+
+/** Every Pi-managed MTPLX connection has a deterministic key. */
+export function resolveMtplxApiKey(apiKey: string | undefined): string {
+	return apiKey ?? DEFAULT_MTPLX_API_KEY;
+}
+
+export function loadResolvedMtplxApiKey(): string {
+	return resolveMtplxApiKey(loadMtplxApiKey());
+}
+
+/** A stable, non-secret identifier suitable for the interactive menu. */
+export function maskMtplxApiKey(apiKey: string | undefined): string {
+	if (!apiKey) return `default (${DEFAULT_MTPLX_API_KEY})`;
+	return apiKey.length <= 4 ? "configured" : `••••${apiKey.slice(-4)}`;
+}
+
+/**
+ * Pi gives a stored `auth.json` API-key credential precedence over models.json.
+ * Keep an existing MTPLX credential in sync; do not create one, because the
+ * provider-level key is sufficient when no stored credential exists.
+ */
+function syncStoredMtplxCredential(apiKey: string): void {
+	if (!existsSync(PI_AUTH_FILE)) return;
+	const parsed = JSON.parse(readFileSync(PI_AUTH_FILE, "utf8")) as Record<string, unknown>;
+	// Current Pi stores credentials directly as { "mtplx": { ... } }. Accept
+	// the older nested shape too, without changing either file's structure.
+	const credentials = typeof parsed.auth === "object" && parsed.auth !== null && !Array.isArray(parsed.auth)
+		? parsed.auth as Record<string, unknown>
+		: parsed;
+	const credential = credentials.mtplx;
+	if (credential === undefined) return;
+	if (typeof credential !== "object" || credential === null || Array.isArray(credential) || (credential as { type?: unknown }).type !== "api_key") {
+		throw new Error("auth.json contains an MTPLX credential that is not an API key");
+	}
+	(credential as Record<string, unknown>).key = apiKey;
+	writeFileSync(PI_AUTH_FILE, JSON.stringify(parsed, null, 2) + "\n");
+}
+
+/** Reconcile Pi's stored MTPLX credential with the configured provider key. */
+export function syncMtplxStoredCredential(): void {
+	try {
+		syncStoredMtplxCredential(loadResolvedMtplxApiKey());
+	} catch (error) {
+		console.error(`MTPLX could not synchronize Pi's stored API key: ${error instanceof Error ? error.message : String(error)}`);
+	}
+}
+
+/** Update Pi's MTPLX provider and any stored MTPLX credential to the same key. */
+export function saveMtplxApiKey(apiKey: string): boolean {
+	const normalized = apiKey.trim();
+	const resolved = resolveMtplxApiKey(normalized || undefined);
+	try {
+		let catalog: Record<string, unknown> = {};
+		if (existsSync(PI_MODELS_FILE)) {
+			const parsed = JSON.parse(readFileSync(PI_MODELS_FILE, "utf8")) as unknown;
+			if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) throw new Error("models.json is not an object");
+			catalog = parsed as Record<string, unknown>;
+		}
+		let providers = catalog.providers;
+		if (providers !== undefined && (typeof providers !== "object" || providers === null || Array.isArray(providers))) {
+			throw new Error("models.json providers is not an object");
+		}
+		const providerCatalog = (providers ??= {}) as Record<string, unknown>;
+		const existing = providerCatalog.mtplx;
+		if (existing !== undefined && (typeof existing !== "object" || existing === null || Array.isArray(existing))) {
+			throw new Error("models.json mtplx provider is not an object");
+		}
+		const provider = (existing ?? {
+			api: "openai-completions",
+			authHeader: true,
+			baseUrl: `http://${HOST}:${PORT}/v1`,
+		}) as Record<string, unknown>;
+		if (normalized) provider.apiKey = normalized;
+		else delete provider.apiKey;
+		providerCatalog.mtplx = provider;
+		catalog.providers = providerCatalog;
+		mkdirSync(join(homedir(), ".pi", "agent"), { recursive: true });
+		writeFileSync(PI_MODELS_FILE, JSON.stringify(catalog, null, 2) + "\n");
+		syncStoredMtplxCredential(resolved);
+		return true;
+	} catch (error) {
+		console.error(`MTPLX could not update API key: ${error instanceof Error ? error.message : String(error)}`);
+		return false;
+	}
+}
+
+/** Save the OpenAI-compatible base URL used by Pi and lifecycle health checks. */
+export function saveMtplxEndpoint(baseUrl: string): boolean {
+	const normalized = baseUrl.trim();
+	if (!normalized) return false;
+	let endpoint: MtplxEndpoint;
+	try {
+		const url = new URL(normalized);
+		if (url.protocol !== "http:" && url.protocol !== "https:") return false;
+		endpoint = mtplxEndpointFromBaseUrl(normalized);
+	} catch {
+		return false;
+	}
+	try {
+		const catalog = existsSync(PI_MODELS_FILE)
+			? JSON.parse(readFileSync(PI_MODELS_FILE, "utf8")) as Record<string, unknown>
+			: {};
+		if (typeof catalog !== "object" || catalog === null || Array.isArray(catalog)) throw new Error("models.json is not an object");
+		const providers = (catalog.providers ??= {}) as Record<string, unknown>;
+		const provider = (providers.mtplx ?? { api: "openai-completions", authHeader: true }) as Record<string, unknown>;
+		if (typeof provider !== "object" || provider === null || Array.isArray(provider)) throw new Error("models.json mtplx provider is not an object");
+		provider.baseUrl = endpoint.baseUrl;
+		providers.mtplx = provider;
+		mkdirSync(join(homedir(), ".pi", "agent"), { recursive: true });
+		writeFileSync(PI_MODELS_FILE, JSON.stringify(catalog, null, 2) + "\n");
+		return true;
+	} catch (error) {
+		console.error(`MTPLX could not update endpoint: ${error instanceof Error ? error.message : String(error)}`);
+		return false;
+	}
 }
 
 export function displayNameFromId(id: string): string {
@@ -86,14 +302,6 @@ export function displayNameFromId(id: string): string {
 		.replace(/\b\w/g, (c) => c.toUpperCase());
 }
 
-export function slugFromId(id: string): string {
-	return id
-		.replace(/^mtplx-/, "")
-		.replace(/[^a-zA-Z0-9]+/g, "-")
-		.replace(/-+/g, "-")
-		.replace(/^-|-$/g, "");
-}
-
 export function isMtplxModel(model: { provider: string; id: string } | undefined): boolean {
 	// The provider name is MTPLX's own PI_PROVIDER_ID (`mtplx`), matching the
 	// provider block that `mtplx start pi` writes to models.json.
@@ -111,9 +319,9 @@ export function commandError(error: unknown): string {
 	return stderr || details.message || "unknown command failure";
 }
 
-export function portIsOccupied(): Promise<boolean> {
+export function portIsOccupied(endpoint = loadMtplxEndpoint()): Promise<boolean> {
 	return new Promise((resolve) => {
-		const socket = createConnection({ host: HOST, port: PORT });
+		const socket = createConnection({ host: endpoint.host, port: endpoint.port });
 		const done = (occupied: boolean) => {
 			socket.destroy();
 			resolve(occupied);