From 1e8a05dc70b092355e3151bf0b12f2e9b3ca2400 Mon Sep 17 00:00:00 2001
From: Mikael Hugo <mikkihugo@users.noreply.github.com>
Date: Tue, 5 May 2026 17:18:56 +0200
Subject: [PATCH] fix: constrain mimo proxy fallbacks

---
 .../core/model-registry-proxy-routing.test.ts | 22 ++++++++++++++++---
 .../src/core/model-registry.ts                | 21 +++++++++++++++---
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/packages/pi-coding-agent/src/core/model-registry-proxy-routing.test.ts b/packages/pi-coding-agent/src/core/model-registry-proxy-routing.test.ts
index 48cf8e050..b484747cb 100644
--- a/packages/pi-coding-agent/src/core/model-registry-proxy-routing.test.ts
+++ b/packages/pi-coding-agent/src/core/model-registry-proxy-routing.test.ts
@@ -253,6 +253,22 @@ describe("ModelRegistry.getModelsForProxy — basic", () => {
 		);
 	});
 
+	it("routes MiMo proxy candidates through direct Xiaomi and subscribed/free relays, never OpenRouter", () => {
+		const registry = createRegistry(() => true);
+		registerNone(registry, "openrouter", "mimo-v2-pro");
+		registerNone(registry, "ollama-cloud", "mimo-v2-pro");
+
+		const result = registry.getModelsForProxy("mimo-v2-pro");
+		assert.deepEqual(
+			result.map((m) => `${m.provider}/${m.id}`),
+			[
+				"xiaomi/mimo-v2-pro",
+				"opencode-go/mimo-v2-pro",
+				"ollama-cloud/mimo-v2-pro",
+			],
+		);
+	});
+
 	it("hides Grok models even when they arrive through aggregators", () => {
 		const registry = createRegistry(() => true);
 		const available = registry.getAvailable();
@@ -439,9 +455,9 @@ describe("ModelRegistry.getModelsForProxy — family priority ordering", () => {
 
 	it("provider not in any family rule falls back to end of list", () => {
 		const registry = createRegistry();
-		registerNone(registry, "zai", "glm-4-air");
-		registerNone(registry, "unknown-aggregator", "glm-4-air");
-		const result = registry.getModelsForProxy("glm-4-air");
+		registerNone(registry, "zai", "unknown-model");
+		registerNone(registry, "unknown-aggregator", "unknown-model");
+		const result = registry.getModelsForProxy("unknown-model");
 		const providers = result.map((m) => m.provider);
 		assert.equal(providers[0], "zai");
 		assert.equal(providers[providers.length - 1], "unknown-aggregator");
diff --git a/packages/pi-coding-agent/src/core/model-registry.ts b/packages/pi-coding-agent/src/core/model-registry.ts
index 295bfcc56..285cc89e8 100644
--- a/packages/pi-coding-agent/src/core/model-registry.ts
+++ b/packages/pi-coding-agent/src/core/model-registry.ts
@@ -87,7 +87,13 @@ export const PROXY_FAMILY_PRIORITY: ReadonlyArray<{
 	// MiMo/Xiaomi — direct API via Xiaomi MiMo Open Platform (api.xiaomimimo.com)
 	// or the Token Plan endpoint (token-plan-sgp.xiaomimimo.com). Both served
 	// under the `xiaomi` provider namespace.
-	{ match: /^mimo-|^XiaomiMiMo\//i, prefix: "mimo-", providers: ["xiaomi"] },
+	{
+		match: /^mimo-|^xiaomi\/mimo-|^XiaomiMiMo\//i,
+		prefix: "mimo-",
+		providers: ["xiaomi"],
+		family_failover: ["opencode-go", "ollama-cloud"],
+		global_fallback: false,
+	},
 	// Gemini/Gemma: route bare model IDs through google-gemini-cli only.
 	// Direct GenAI and Vertex providers stay explicit provider-qualified routes,
 	// but they are hidden from normal SF/TUI selection and default fallback.
@@ -1183,13 +1189,22 @@ export class ModelRegistry {
 		overrides: Record<string, string[]> = {},
 		providerModelAllow?: ProviderModelAllowList,
 	): Model<Api>[] {
+		const order = this.buildCandidateOrder(modelId, overrides);
+		const hasScopedProviderOrder =
+			Object.entries(overrides).some(([k]) => modelId.startsWith(k)) ||
+			PROXY_FAMILY_PRIORITY.some((r) => r.match.test(modelId));
+		const allowedProviders =
+			hasScopedProviderOrder && order.length > 0 ? new Set(order) : undefined;
 		const candidates = this.filterProviderModelAllow(
-			this.models.filter((m) => m.id === modelId),
+			this.getAllWithDiscovered().filter(
+				(m) =>
+					m.id === modelId &&
+					(!allowedProviders || allowedProviders.has(m.provider)),
+			),
 			providerModelAllow,
 		);
 		if (candidates.length === 0) return [];
 
-		const order = this.buildCandidateOrder(modelId, overrides);
 		const sorted = [...candidates].sort((a, b) => {
 			const pa = order.indexOf(a.provider);
 			const pb = order.indexOf(b.provider);