From 6bb831370689bc6f4b82b48bf3f85fbe603b0e3c Mon Sep 17 00:00:00 2001 From: Automation Date: Fri, 29 May 2026 08:10:16 +0200 Subject: [PATCH] fix: provider model sync pruning and dynamic antigravity MITM proxy mappings --- src/lib/db/models.ts | 25 +++++++ src/lib/providerModels/managedModelImport.ts | 69 ++++++++++++++++++ tests/unit/managed-model-import.test.ts | 73 ++++++++++++++++++++ 3 files changed, 167 insertions(+) diff --git a/src/lib/db/models.ts b/src/lib/db/models.ts index b359358ca013..33e6a59c3494 100644 --- a/src/lib/db/models.ts +++ b/src/lib/db/models.ts @@ -750,6 +750,31 @@ export async function deleteSyncedAvailableModelsForProvider(providerId: string) return Number(result.changes || 0); } +/** + * Prune stale synced available models for a provider, keeping only the specified allowed connection IDs. + * Returns the number of keys deleted. + */ +export async function pruneStaleSyncedAvailableModelsForProvider( + providerId: string, + allowedConnectionIds: string[] +): Promise { + const db = getDbInstance(); + if (allowedConnectionIds.length === 0) { + return deleteSyncedAvailableModelsForProvider(providerId); + } + const placeholders = allowedConnectionIds.map(() => "?").join(","); + const keyPrefix = `${providerId}:`; + const allowedKeys = allowedConnectionIds.map((id) => `${providerId}:${id}`); + const result = db + .prepare( + `DELETE FROM key_value WHERE namespace = 'syncedAvailableModels' AND key LIKE ? AND key NOT IN (${placeholders})` + ) + .run(`${keyPrefix}%`, ...allowedKeys); + backupDbFile("pre-write"); + return Number(result.changes || 0); +} + + export async function updateCustomModel( providerId: string, modelId: string, diff --git a/src/lib/providerModels/managedModelImport.ts b/src/lib/providerModels/managedModelImport.ts index 9e5702a50d76..5610e643ed13 100644 --- a/src/lib/providerModels/managedModelImport.ts +++ b/src/lib/providerModels/managedModelImport.ts @@ -4,14 +4,22 @@ import { mergeModelCompatOverride, replaceCustomModels, replaceSyncedAvailableModelsForConnection, + pruneStaleSyncedAvailableModelsForProvider, + setMitmAliasAll, + getSyncedAvailableModels, type ModelCompatPatch, type SyncedAvailableModel, } from "@/lib/db/models"; +import { getProviderConnections } from "@/lib/db/providers"; import { syncManagedAvailableModelAliases, usesManagedAvailableModels, } from "@/lib/providerModels/managedAvailableModels"; import { normalizeDiscoveredModels } from "@/lib/providerModels/modelDiscovery"; +import { + ANTIGRAVITY_MODEL_ALIASES, + ANTIGRAVITY_REVERSE_MODEL_ALIASES, +} from "@omniroute/open-sse/config/antigravityModelAliases.ts"; type JsonRecord = Record; @@ -238,6 +246,67 @@ export async function importManagedModels({ ); } + // Prune stale/inactive connection caches for this provider + const activeConnections = await getProviderConnections({ provider: providerId, isActive: true }); + const allowedConnectionIds = Array.from( + new Set([...activeConnections.map((c) => String(c.id)), connectionId]) + ); + await pruneStaleSyncedAvailableModelsForProvider(providerId, allowedConnectionIds); + + // If this is the "antigravity" provider, dynamically regenerate and persist the mitmAlias mapping for "antigravity" + if (providerId === "antigravity") { + const allAntigravityModels = await getSyncedAvailableModels("antigravity"); + const syncedIds = new Set(allAntigravityModels.map((m) => m.id).filter(Boolean) as string[]); + + // Transitive/recursive resolution helper + const resolveTransitively = (name: string): string => { + let current = name; + const visited = new Set(); + while (current && !visited.has(current)) { + if (syncedIds.has(current)) { + return current; + } + visited.add(current); + if (ANTIGRAVITY_MODEL_ALIASES && (ANTIGRAVITY_MODEL_ALIASES as any)[current]) { + current = (ANTIGRAVITY_MODEL_ALIASES as any)[current]; + continue; + } + if (ANTIGRAVITY_REVERSE_MODEL_ALIASES && (ANTIGRAVITY_REVERSE_MODEL_ALIASES as any)[current]) { + current = (ANTIGRAVITY_REVERSE_MODEL_ALIASES as any)[current]; + continue; + } + break; + } + return current; + }; + + // Gather all candidate alias names to check + const candidates = new Set(); + for (const id of syncedIds) { + candidates.add(id); + } + for (const [k, v] of Object.entries(ANTIGRAVITY_MODEL_ALIASES || {})) { + candidates.add(k); + candidates.add(v); + } + for (const [k, v] of Object.entries(ANTIGRAVITY_REVERSE_MODEL_ALIASES || {})) { + candidates.add(k); + candidates.add(v); + } + + // Build the dynamic mapping dictionary + const mappings: Record = {}; + for (const alias of candidates) { + const resolvedId = resolveTransitively(alias); + if (syncedIds.has(resolvedId)) { + mappings[alias] = `antigravity/${resolvedId}`; + } + } + + await setMitmAliasAll("antigravity", mappings); + } + + let syncedAliases = 0; if (usesManagedAvailableModels(providerId) && (mode === "merge" || discoveredModels.length > 0)) { const aliasModelIds = mode === "sync" ? syncedAvailableModels : discoveredModels; diff --git a/tests/unit/managed-model-import.test.ts b/tests/unit/managed-model-import.test.ts index 0a09fe12ffe2..a0e9ea5c2240 100644 --- a/tests/unit/managed-model-import.test.ts +++ b/tests/unit/managed-model-import.test.ts @@ -84,3 +84,76 @@ test("provider-level synced model deletion removes only that provider", async () { id: "shared/model-c", name: "Model C", source: "imported" }, ]); }); + +test("pruning stale connection available models during import", async () => { + const db = core.getDbInstance(); + // Insert connections + db.prepare( + "INSERT INTO provider_connections (id, provider, auth_type, name, is_active, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)" + ).run("conn-active", "openrouter", "apikey", "Active Connection", 1, "2026-05-29", "2026-05-29"); + + db.prepare( + "INSERT INTO provider_connections (id, provider, auth_type, name, is_active, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)" + ).run("conn-stale", "openrouter", "apikey", "Stale Connection", 0, "2026-05-29", "2026-05-29"); + + // Create synced available models for both + await modelsDb.replaceSyncedAvailableModelsForConnection("openrouter", "conn-active", [ + { id: "shared/model-active", name: "Model Active", source: "imported" }, + ]); + await modelsDb.replaceSyncedAvailableModelsForConnection("openrouter", "conn-stale", [ + { id: "shared/model-stale", name: "Model Stale", source: "imported" }, + ]); + + // Import on conn-new + await importManagedModels({ + providerId: "openrouter", + connectionId: "conn-new", + mode: "sync", + fetchedModels: [{ id: "shared/model-new", name: "Model New" }], + }); + + // Check models for "openrouter" + const allSyncedModels = await modelsDb.getSyncedAvailableModels("openrouter"); + + // Stale connection should be pruned. Active connection and the new syncing connection should be kept. + const ids = allSyncedModels.map((m) => m.id); + assert.ok(ids.includes("shared/model-active")); + assert.ok(ids.includes("shared/model-new")); + assert.ok(!ids.includes("shared/model-stale")); +}); + +test("antigravity sync dynamically builds and saves mitmAlias mappings", async () => { + const db = core.getDbInstance(); + // Create an antigravity connection + db.prepare( + "INSERT INTO provider_connections (id, provider, auth_type, name, is_active, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)" + ).run("antigravity-conn", "antigravity", "oauth", "Antigravity", 1, "2026-05-29", "2026-05-29"); + + await importManagedModels({ + providerId: "antigravity", + connectionId: "antigravity-conn", + mode: "sync", + fetchedModels: [ + { id: "gemini-3.5-flash", name: "Gemini 3.5 Flash" }, + { id: "custom-antigravity-model", name: "Custom Antigravity Model" }, + ], + }); + + const models = await modelsDb.getSyncedAvailableModels("antigravity"); + console.log("SYNCED MODELS IN TEST:", models); + + const mitmMappings = await modelsDb.getMitmAlias("antigravity"); + console.log("MITM MAPPINGS IN TEST:", mitmMappings); + + // Should contain standard mapping + assert.equal(mitmMappings["gemini-3.5-flash"], "antigravity/gemini-3.5-flash"); + assert.equal(mitmMappings["custom-antigravity-model"], "antigravity/custom-antigravity-model"); + + // Should contain reverse alias mappings (gemini-3.5-flash-preview maps to gemini-3.5-flash) + assert.equal(mitmMappings["gemini-3.5-flash-preview"], "antigravity/gemini-3.5-flash"); + assert.equal(mitmMappings["gemini-3-flash-agent"], "antigravity/gemini-3.5-flash"); + + // Should contain forward alias mappings (gemini-3.5-flash-preview maps to gemini-3.5-flash) + assert.equal(mitmMappings["gemini-3.5-flash-preview"], "antigravity/gemini-3.5-flash"); +}); +