Multi-model LLM load balancer (local + frontier + optional Nova)

CHANGELOG.md

@@ -5,6 +5,21 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+- Multi-model load balancing: optionally spread each chat across all installed local models (native MLX + Ollama), all frontier models (OpenRouter), and an optional Nova Gateway — health-gated and load-balanced — instead of a single pinned model
+- Three independent, persisted toggles in **Settings → Balancer**: All local models, All frontier models (OpenRouter), Nova Gateway (with a configurable gateway URL and Keychain-stored OpenRouter key)
+- `ModelRegistry` (model discovery + pool composition), `LoadBalancer` (round-robin / least-busy with health gating), `OpenRouterProvider` / `OpenAICompatibleRequest`, and `KeychainStore` — pure, network-free, unit-tested building blocks ported from AIStudio
+- `LLMBalancer` service wiring discovery → health map → balanced dispatch into `ChatViewModel`
+- Network-free `LoadBalancerTests` suite (24 tests) covering parsing, pool composition, and selection policies
+
+### Changed
+- `ChatViewModel` routes generation through the balancer when any toggle is on; falls back cleanly to the single pinned MLX model when the pool is empty or all toggles are off (existing behavior preserved)
+
+### Notes
+- Nova is never a hard requirement — the balancer works with zero Nova present; the Nova Gateway is one optional, health-probed entry that drops out if unavailable
+
 ## [5.0.0] - 2026-02-19
 
 ### Removed

MLX Code Tests/LoadBalancerTests.swift

@@ -0,0 +1,276 @@
+//
+//  LoadBalancerTests.swift
+//  MLX Code Tests
+//
+//  Deterministic (no-network) tests for model discovery parsing, pool
+//  composition, and the load-balancer selection policies. Ported from AIStudio
+//  and adapted to MLX Code's backend set and MLX discovery source.
+//
+//  Copyright © 2026 Jordan Koch. All rights reserved.
+//
+
+import XCTest
+@testable import MLX_Code
+
+final class LoadBalancerTests: XCTestCase {
+
+    // MARK: - parseOllamaTags
+
+    func testParseOllamaTagsMapsModels() {
+        let json = """
+        {"models": [
+            {"name": "mistral:latest", "size": 123},
+            {"name": "llama3.2:3b"},
+            {"size": 5}
+        ]}
+        """
+        let models = ModelRegistry.parseOllamaTags(Data(json.utf8))
+        XCTAssertEqual(models.count, 2)
+        XCTAssertEqual(models.map { $0.modelName }, ["mistral:latest", "llama3.2:3b"])
+        XCTAssertTrue(models.allSatisfy { $0.backend == .ollama })
+        XCTAssertEqual(models[0].id, "ollama|mistral:latest")
+        XCTAssertEqual(models[0].endpoint, "http://localhost:11434/api/chat")
+    }
+
+    func testParseOllamaTagsEmptyAndGarbage() {
+        XCTAssertTrue(ModelRegistry.parseOllamaTags(Data("nonsense".utf8)).isEmpty)
+        XCTAssertTrue(ModelRegistry.parseOllamaTags(Data("{}".utf8)).isEmpty)
+        XCTAssertTrue(ModelRegistry.parseOllamaTags(Data(#"{"models": []}"#.utf8)).isEmpty)
+        XCTAssertTrue(ModelRegistry.parseOllamaTags(Data()).isEmpty)
+    }
+
+    // MARK: - parseMLXModels (HF hub cache form)
+
+    func testParseMLXModelsFromHubDirs() {
+        let dirs = [
+            "models--mlx-community--Llama-3.2-3B-Instruct-4bit",
+            "models--meta-llama--Llama-3.1-8B",          // not MLX → excluded
+            "models--mlx-community--Qwen2.5-7B-4bit",
+            "blobs",                                       // not a model dir → excluded
+            "models--"                                     // empty repo → excluded
+        ]
+        let models = ModelRegistry.parseMLXModels(hubDirectoryNames: dirs)
+        XCTAssertEqual(models.count, 2)
+        XCTAssertEqual(models.map { $0.modelName },
+                       ["mlx-community/Llama-3.2-3B-Instruct-4bit", "mlx-community/Qwen2.5-7B-4bit"])
+        XCTAssertTrue(models.allSatisfy { $0.backend == .mlx })
+        XCTAssertEqual(models[0].displayName, "Llama-3.2-3B-Instruct-4bit")
+    }
+
+    func testParseMLXModelsEmpty() {
+        XCTAssertTrue(ModelRegistry.parseMLXModels(hubDirectoryNames: []).isEmpty)
+        XCTAssertTrue(ModelRegistry.parseMLXModels(hubDirectoryNames: ["random", "stuff"]).isEmpty)
+    }
+
+    // MARK: - mlxModels(fromLocalNames:) — MLX Code's own discovery source
+
+    func testMLXModelsFromLocalNames() {
+        let models = ModelRegistry.mlxModels(fromLocalNames: [
+            "Qwen2.5-7B-Instruct-4bit",
+            "mlx-community/Llama-3.2-3B-Instruct-4bit",
+            "  ",                                           // blank → excluded
+            ""                                              // empty → excluded
+        ])
+        XCTAssertEqual(models.count, 2)
+        XCTAssertTrue(models.allSatisfy { $0.backend == .mlx })
+        XCTAssertEqual(models[0].modelName, "Qwen2.5-7B-Instruct-4bit")
+        XCTAssertEqual(models[0].id, "mlx|Qwen2.5-7B-Instruct-4bit")
+        // display name is the short (last path component) form
+        XCTAssertEqual(models[1].displayName, "Llama-3.2-3B-Instruct-4bit")
+        XCTAssertTrue(models.allSatisfy { $0.endpoint.isEmpty })
+    }
+
+    // MARK: - frontier + nova mapping
+
+    func testFrontierModelsMapping() {
+        let frontier = ModelRegistry.frontierModels(from: ["openai/gpt-4o", "anthropic/claude-sonnet-4.5", ""])
+        XCTAssertEqual(frontier.count, 2)
+        XCTAssertTrue(frontier.allSatisfy { $0.backend == .openRouter })
+        XCTAssertEqual(frontier[0].endpoint, OpenRouterProvider.chatCompletionsURL)
+    }
+
+    func testNovaGatewayModel() {
+        let nova = ModelRegistry.novaGatewayModel()
+        XCTAssertEqual(nova.backend, .novaGateway)
+        XCTAssertEqual(nova.endpoint, "http://127.0.0.1:18792/v1/chat/completions")
+    }
+
+    // MARK: - Pool composition (toggles)
+
+    private func samplePool() -> (local: [DiscoveredModel], frontier: [DiscoveredModel], nova: DiscoveredModel) {
+        let ollama = [DiscoveredModel(modelName: "mistral:latest", backend: .ollama, endpoint: "e")]
+        let mlx = [DiscoveredModel(modelName: "mlx-community/Qwen", backend: .mlx, endpoint: "")]
+        let frontier = ModelRegistry.frontierModels(from: ["openai/gpt-4o"])
+        let nova = ModelRegistry.novaGatewayModel()
+        return (ollama + mlx, frontier, nova)
+    }
+
+    func testAssemblePoolLocalOnly() {
+        let s = samplePool()
+        let pool = ModelRegistry.assemblePool(
+            ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova,
+            useAllLocalModels: true, enableAllFrontierModels: false, useNovaGateway: false)
+        XCTAssertEqual(pool.count, 2)
+        XCTAssertTrue(pool.allSatisfy { $0.backend == .ollama || $0.backend == .mlx })
+    }
+
+    func testAssemblePoolFrontierOnly() {
+        let s = samplePool()
+        let pool = ModelRegistry.assemblePool(
+            ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova,
+            useAllLocalModels: false, enableAllFrontierModels: true, useNovaGateway: false)
+        XCTAssertEqual(pool.count, 1)
+        XCTAssertEqual(pool[0].backend, .openRouter)
+    }
+
+    func testAssemblePoolBothPlusNova() {
+        let s = samplePool()
+        let pool = ModelRegistry.assemblePool(
+            ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova,
+            useAllLocalModels: true, enableAllFrontierModels: true, useNovaGateway: true)
+        XCTAssertEqual(pool.count, 4)
+        XCTAssertTrue(pool.contains { $0.backend == .novaGateway })
+    }
+
+    func testAssemblePoolNovaAbsentWhenToggleOff() {
+        let s = samplePool()
+        let pool = ModelRegistry.assemblePool(
+            ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova,
+            useAllLocalModels: true, enableAllFrontierModels: false, useNovaGateway: false)
+        XCTAssertFalse(pool.contains { $0.backend == .novaGateway })
+    }
+
+    func testAssemblePoolAllOff() {
+        let s = samplePool()
+        let pool = ModelRegistry.assemblePool(
+            ollama: [s.local[0]], mlx: [s.local[1]], frontier: s.frontier, novaGateway: s.nova,
+            useAllLocalModels: false, enableAllFrontierModels: false, useNovaGateway: false)
+        XCTAssertTrue(pool.isEmpty)
+    }
+
+    func testAssemblePoolDeduplicates() {
+        let dup = DiscoveredModel(modelName: "mistral:latest", backend: .ollama, endpoint: "e")
+        let pool = ModelRegistry.assemblePool(
+            ollama: [dup, dup], mlx: [], frontier: [], novaGateway: nil,
+            useAllLocalModels: true, enableAllFrontierModels: false, useNovaGateway: false)
+        XCTAssertEqual(pool.count, 1)
+    }
+
+    // MARK: - Round-robin policy
+
+    func testRoundRobinCyclesAndWraps() {
+        let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") }
+        let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, true) })
+        let lb = LoadBalancer()
+
+        var picks: [String] = []
+        for _ in 0..<7 {
+            picks.append(lb.next(pool: pool, health: health, policy: .roundRobin)!.modelName)
+        }
+        XCTAssertEqual(picks, ["a", "b", "c", "a", "b", "c", "a"])
+    }
+
+    // MARK: - Least-busy policy
+
+    func testLeastBusyPicksLowestInFlight() {
+        let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") }
+        let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, true) })
+        let lb = LoadBalancer()
+
+        // a: 2 in-flight, b: 0, c: 1 → b is least busy.
+        lb.checkOut(pool[0].id); lb.checkOut(pool[0].id)
+        lb.checkOut(pool[2].id)
+        XCTAssertEqual(lb.next(pool: pool, health: health, policy: .leastBusy)?.modelName, "b")
+    }
+
+    func testLeastBusyTieBreaksByPoolOrder() {
+        let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") }
+        let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, true) })
+        let lb = LoadBalancer()
+        // All zero in-flight → first in pool order wins, deterministically.
+        XCTAssertEqual(lb.next(pool: pool, health: health, policy: .leastBusy)?.modelName, "a")
+    }
+
+    func testCheckInNeverGoesNegative() {
+        let lb = LoadBalancer()
+        lb.checkIn("x")
+        XCTAssertEqual(lb.inFlight["x"], 0)
+        lb.checkOut("x"); lb.checkIn("x"); lb.checkIn("x")
+        XCTAssertEqual(lb.inFlight["x"], 0)
+    }
+
+    // MARK: - Health gating
+
+    func testHealthMapExcludesUnhealthy() {
+        let pool = ["a", "b", "c"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") }
+        let lb = LoadBalancer()
+        // b marked unhealthy; a & c absent-from-map default to healthy.
+        let health = [pool[1].id: false]
+        let picked = (0..<4).map { _ in lb.next(pool: pool, health: health, policy: .roundRobin)!.modelName }
+        XCTAssertFalse(picked.contains("b"))
+        XCTAssertEqual(Set(picked), ["a", "c"])
+    }
+
+    func testAllUnhealthyReturnsNil() {
+        let pool = ["a", "b"].map { DiscoveredModel(modelName: $0, backend: .ollama, endpoint: "e") }
+        let health = Dictionary(uniqueKeysWithValues: pool.map { ($0.id, false) })
+        let lb = LoadBalancer()
+        XCTAssertNil(lb.next(pool: pool, health: health, policy: .roundRobin))
+        XCTAssertNil(lb.next(pool: pool, health: health, policy: .leastBusy))
+    }
+
+    func testEmptyPoolReturnsNil() {
+        let lb = LoadBalancer()
+        XCTAssertNil(lb.next(pool: [], health: [:], policy: .roundRobin))
+    }
+
+    // MARK: - OpenAI-compatible request building (network-free)
+
+    func testOpenAICompatibleRequestBuild() throws {
+        let request = try OpenAICompatibleRequest.build(
+            endpoint: "https://openrouter.ai/api/v1/chat/completions",
+            model: "openai/gpt-4o",
+            messages: [["role": "user", "content": "hi"]],
+            temperature: 0.5,
+            maxTokens: 128,
+            stream: false,
+            headers: ["Authorization": "Bearer test"]
+        )
+        XCTAssertEqual(request.httpMethod, "POST")
+        XCTAssertEqual(request.value(forHTTPHeaderField: "Authorization"), "Bearer test")
+        XCTAssertEqual(request.value(forHTTPHeaderField: "Content-Type"), "application/json")
+        let body = try XCTUnwrap(request.httpBody)
+        let json = try XCTUnwrap(try JSONSerialization.jsonObject(with: body) as? [String: Any])
+        XCTAssertEqual(json["model"] as? String, "openai/gpt-4o")
+        XCTAssertEqual(json["stream"] as? Bool, false)
+    }
+
+    func testOpenAICompatibleRequestInvalidURLThrows() {
+        XCTAssertThrowsError(try OpenAICompatibleRequest.build(
+            endpoint: "",
+            model: "m",
+            messages: [],
+            temperature: 0.5,
+            maxTokens: 10,
+            stream: false
+        ))
+    }
+
+    func testParseOpenRouterModels() {
+        let json = #"{"data":[{"id":"openai/gpt-4o"},{"id":"anthropic/claude-sonnet-4.5"},{"foo":"bar"}]}"#
+        XCTAssertEqual(OpenRouterProvider.parseModels(Data(json.utf8)),
+                       ["openai/gpt-4o", "anthropic/claude-sonnet-4.5"])
+        XCTAssertTrue(OpenRouterProvider.parseModels(Data("garbage".utf8)).isEmpty)
+    }
+
+    // MARK: - Backend type
+
+    func testNovaGatewayBackendType() {
+        XCTAssertEqual(LLMBackendType.novaGateway.rawValue, "novagateway")
+        XCTAssertEqual(LLMBackendType.novaGateway.displayName, "Nova Gateway")
+        XCTAssertEqual(LLMBackendType.novaGateway.defaultURL, "http://127.0.0.1:18792")
+        XCTAssertFalse(LLMBackendType.novaGateway.icon.isEmpty)
+        // MLX Code's balancer backend set: ollama, mlx, openRouter, novaGateway.
+        XCTAssertEqual(LLMBackendType.allCases.count, 4)
+    }
+}

MLX Code.xcodeproj/project.pbxproj

@@ -7,6 +7,11 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		AABB01000000000000000002 /* ModelRegistry.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB01000000000000000001 /* ModelRegistry.swift */; };
+		AABB02000000000000000002 /* OpenRouterProvider.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB02000000000000000001 /* OpenRouterProvider.swift */; };
+		AABB03000000000000000002 /* KeychainStore.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB03000000000000000001 /* KeychainStore.swift */; };
+		AABB04000000000000000002 /* LLMBalancer.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB04000000000000000001 /* LLMBalancer.swift */; };
+		AABB05000000000000000002 /* LoadBalancerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = AABB05000000000000000001 /* LoadBalancerTests.swift */; };
 		043C4E7F3FCA17ECD6EB1A51 /* PromptTemplatesView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 942A3F3878CD789BDF616A78 /* PromptTemplatesView.swift */; };
 		0475B2DD296EECD045126C17 /* PrerequisitesView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 22BF0FCAFA31E5516E7B5EC6 /* PrerequisitesView.swift */; };
 		054C45D18234B5AE6A2A3233 /* GettingStarted.md in Resources */ = {isa = PBXBuildFile; fileRef = 3072558342FD102B940FF3DB /* GettingStarted.md */; };
@@ -178,6 +183,11 @@
 /* End PBXCopyFilesBuildPhase section */
 
 /* Begin PBXFileReference section */
+		AABB01000000000000000001 /* ModelRegistry.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelRegistry.swift; sourceTree = "<group>"; };
+		AABB02000000000000000001 /* OpenRouterProvider.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OpenRouterProvider.swift; sourceTree = "<group>"; };
+		AABB03000000000000000001 /* KeychainStore.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = KeychainStore.swift; sourceTree = "<group>"; };
+		AABB04000000000000000001 /* LLMBalancer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LLMBalancer.swift; sourceTree = "<group>"; };
+		AABB05000000000000000001 /* LoadBalancerTests.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; path = LoadBalancerTests.swift; sourceTree = "<group>"; };
 		0095C89C8F1E4539863AE74D96B702F4 /* MLXService.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MLXService.swift; sourceTree = "<group>"; };
 		01A056387C5D6883F0932FA4 /* GitService.swift */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = sourcecode.swift; name = GitService.swift; path = "MLX Code/Services/GitService.swift"; sourceTree = "<group>"; };
 		0674F2D9E84C4AC3AC8ED170619E99EA /* FileService.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FileService.swift; sourceTree = "<group>"; };
@@ -487,6 +497,10 @@
 		8D7B12981853422D971FCAD3B063E827 /* Services */ = {
 			isa = PBXGroup;
 			children = (
+				AABB01000000000000000001 /* ModelRegistry.swift */,
+				AABB02000000000000000001 /* OpenRouterProvider.swift */,
+				AABB03000000000000000001 /* KeychainStore.swift */,
+				AABB04000000000000000001 /* LLMBalancer.swift */,
 				0095C89C8F1E4539863AE74D96B702F4 /* MLXService.swift */,
 				BB001CDEF123456789ABCDE000022222 /* XcodeActionHandler.swift */,
 				2113E8E02AFC404C917BDECE45AD0C31 /* XcodeService.swift */,
@@ -523,6 +537,7 @@
 		CE911C9235A55156197C996A /* MLX Code Tests */ = {
 			isa = PBXGroup;
 			children = (
+				AABB05000000000000000001 /* LoadBalancerTests.swift */,
 				1B364AD048BC9F2567D45CC7 /* AppSettingsTests.swift */,
 				8C37CB4BC161157E297B1C70 /* CommandValidatorTests.swift */,
 				F975089CB3BD79BDBC9F5CD3 /* ContextManagerTests.swift */,
@@ -859,6 +874,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				AABB05000000000000000002 /* LoadBalancerTests.swift in Sources */,
 				678CCB7627496E9731B7D734 /* AppSettingsTests.swift in Sources */,
 				AA74C582F22930DF98A0CB7B /* CommandValidatorTests.swift in Sources */,
 				2CAAE35717EEF0F0BB6889D1 /* ContextManagerTests.swift in Sources */,
@@ -899,6 +915,10 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				AABB01000000000000000002 /* ModelRegistry.swift in Sources */,
+				AABB02000000000000000002 /* OpenRouterProvider.swift in Sources */,
+				AABB03000000000000000002 /* KeychainStore.swift in Sources */,
+				AABB04000000000000000002 /* LLMBalancer.swift in Sources */,
 				E2FCF04057F99E48BF46CCF3 /* CommandValidator.swift in Sources */,
 				A0013163E03EA758655F244E /* ModelSecurityValidator.swift in Sources */,
 				4A14B6380E9E4A44949059C9DE2D703A /* MLXCodeApp.swift in Sources */,