From 5551fb8e65d754d90ab1b415a8390287fa8b5179 Mon Sep 17 00:00:00 2001 From: Jacob Fu <141651335+FuJacob@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:30:44 -0700 Subject: [PATCH] Preload default Ollama model before suggestions --- .../SuggestionCoordinator+Input.swift | 5 ++- .../OpenAICompatibleEndpointModels.swift | 11 ++++++ .../Runtime/OpenAICompatibleAPIClient.swift | 37 ++++++++++++++++++ .../OpenAICompatibleSuggestionEngine.swift | 24 ++++++++++++ .../OpenAICompatibleAPIClientTests.swift | 39 +++++++++++++++++++ 5 files changed, 114 insertions(+), 2 deletions(-) diff --git a/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift b/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift index ce66515f..f0bc4262 100644 --- a/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift +++ b/Cotabby/App/Coordinators/SuggestionCoordinator+Input.swift @@ -118,8 +118,9 @@ extension SuggestionCoordinator { state = .idle // The user is now on a new editable surface and is likely to type soon. Prime the // selected engine in the background so weight loading and instruction tokenization - // happen before the first real `respond` instead of inside its critical path. Llama's - // default `prewarm` is a no-op, so this call is FM-only by design. + // happen before the first real `respond` instead of inside its critical path. The + // external endpoint also uses this hook to cold-load default local Ollama separately + // from the aggressively cancellable autocomplete request. prewarmEngineForCurrentField(rawContext: focusedContext) } diff --git a/Cotabby/Models/OpenAICompatibleEndpointModels.swift b/Cotabby/Models/OpenAICompatibleEndpointModels.swift index af1700ed..3be0ca17 100644 --- a/Cotabby/Models/OpenAICompatibleEndpointModels.swift +++ b/Cotabby/Models/OpenAICompatibleEndpointModels.swift @@ -126,6 +126,17 @@ nonisolated struct OpenAICompatibleEndpointConfiguration: Equatable, Sendable { baseURL.appendingPathComponent(path, isDirectory: false) } + /// Cotabby's default endpoint is specifically the local Ollama server. Ollama's preload API + /// lives outside its OpenAI-compatible `/v1` namespace, so this narrowly scoped URL avoids + /// sending Ollama-only requests to arbitrary LM Studio, vLLM, or hosted endpoints. + var defaultOllamaGenerateURL: URL? { + guard baseURL.absoluteString == Self.defaultBaseURLString, + var components = URLComponents(url: baseURL, resolvingAgainstBaseURL: false) + else { return nil } + components.path = "/api/generate" + return components.url + } + var privacyWarning: String? { switch hostScope { case .loopback: diff --git a/Cotabby/Services/Runtime/OpenAICompatibleAPIClient.swift b/Cotabby/Services/Runtime/OpenAICompatibleAPIClient.swift index c8dec55b..5d3575c0 100644 --- a/Cotabby/Services/Runtime/OpenAICompatibleAPIClient.swift +++ b/Cotabby/Services/Runtime/OpenAICompatibleAPIClient.swift @@ -50,6 +50,29 @@ final class OpenAICompatibleAPIClient { return payload.data.sorted { $0.id.localizedCaseInsensitiveCompare($1.id) == .orderedAscending } } + /// Loads the selected model through Ollama's native API without tying that cold start to a + /// cancellable autocomplete request. Returning `false` means the configured endpoint is not + /// Cotabby's known local Ollama default, so callers can treat warmup as an intentional no-op. + func preloadDefaultOllamaModel( + configuration: OpenAICompatibleEndpointConfiguration, + apiKey: String? + ) async throws -> Bool { + guard !configuration.modelName.isEmpty else { + throw OpenAICompatibleEndpointError.emptyModelName + } + guard let url = configuration.defaultOllamaGenerateURL else { return false } + + var request = URLRequest(url: url, timeoutInterval: 120) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + applyHeaders(to: &request, apiKey: apiKey) + request.httpBody = try encoder.encode(OllamaPreloadRequest(model: configuration.modelName)) + + let (_, response) = try await session.data(for: request) + try Self.validate(response) + return true + } + func generate( configuration: OpenAICompatibleEndpointConfiguration, apiKey: String?, @@ -242,6 +265,20 @@ private nonisolated struct ModelListResponse: Decodable { let data: [OpenAICompatibleModelOption] } +/// Ollama treats an empty, non-streaming generate request as a model load. `keep_alive = -1` +/// keeps the weights resident until Ollama is stopped or explicitly asked to unload them. +private nonisolated struct OllamaPreloadRequest: Encodable { + let model: String + let prompt = "" + let stream = false + let keepAlive = -1 + + private enum CodingKeys: String, CodingKey { + case model, prompt, stream + case keepAlive = "keep_alive" + } +} + private nonisolated struct CompletionRequest: Encodable { let model: String let prompt: String diff --git a/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift b/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift index 33e2f703..7e5c24d8 100644 --- a/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift +++ b/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift @@ -125,4 +125,28 @@ final class OpenAICompatibleSuggestionEngine: SuggestionGenerating { } func resetCachedGenerationContext() async {} + + /// Best-effort Ollama cold-start work runs through a request that is independent from the + /// coordinator's cancellable prediction task. Typing again can therefore discard stale + /// autocomplete work without aborting the model load that future requests need. + func prewarm(for _: SuggestionRequest) async { + do { + let configuration = try configurationProvider() + let didPreload = try await client.preloadDefaultOllamaModel( + configuration: configuration, + apiKey: try apiKeyProvider() + ) + guard didPreload else { return } + CotabbyLogger.runtime.info( + "Preloaded default Ollama model", + metadata: ["model": .string(configuration.modelName)] + ) + } catch is CancellationError { + // App shutdown may cancel this opportunistic work; warmup never becomes user-facing. + } catch { + CotabbyLogger.runtime.warning( + "Default Ollama model preload failed: \(error.localizedDescription)" + ) + } + } } diff --git a/CotabbyTests/OpenAICompatibleAPIClientTests.swift b/CotabbyTests/OpenAICompatibleAPIClientTests.swift index 17d36773..0845a378 100644 --- a/CotabbyTests/OpenAICompatibleAPIClientTests.swift +++ b/CotabbyTests/OpenAICompatibleAPIClientTests.swift @@ -72,11 +72,13 @@ final class OpenAICompatibleAPIClientTests: XCTestCase { let loopback = try configuration(baseURL: " http://127.0.0.1:11434/ ") XCTAssertEqual(loopback.baseURL.absoluteString, "http://127.0.0.1:11434/v1") XCTAssertEqual(loopback.apiURL(path: "models").absoluteString, "http://127.0.0.1:11434/v1/models") + XCTAssertEqual(loopback.defaultOllamaGenerateURL?.absoluteString, "http://127.0.0.1:11434/api/generate") XCTAssertEqual(loopback.hostScope, .loopback) XCTAssertNil(loopback.privacyWarning) let lan = try configuration(baseURL: "http://192.168.1.50:8000/v1/") XCTAssertEqual(lan.baseURL.absoluteString, "http://192.168.1.50:8000/v1") + XCTAssertNil(lan.defaultOllamaGenerateURL) XCTAssertEqual(lan.hostScope, .localNetwork) XCTAssertNotNil(lan.privacyWarning) @@ -186,6 +188,43 @@ final class OpenAICompatibleAPIClientTests: XCTestCase { XCTAssertEqual(partials, [" hel", " hello"]) } + func test_defaultOllamaPreload_usesNativeRouteLongTimeoutAndKeepsModelResident() async throws { + let client = makeClient() + EndpointStubURLProtocol.handler = { request in + XCTAssertEqual(request.url?.absoluteString, "http://127.0.0.1:11434/api/generate") + XCTAssertEqual(request.timeoutInterval, 120) + XCTAssertEqual(request.httpMethod, "POST") + let json = try Self.jsonBody(request) + XCTAssertEqual(json["model"] as? String, "gemma4:12b-mlx") + XCTAssertEqual(json["prompt"] as? String, "") + XCTAssertEqual(json["stream"] as? Bool, false) + XCTAssertEqual(json["keep_alive"] as? Int, -1) + return Self.response(request: request, body: #"{"done":true}"#) + } + + let didPreload = try await client.preloadDefaultOllamaModel( + configuration: configuration(), + apiKey: nil + ) + + XCTAssertTrue(didPreload) + } + + func test_nonDefaultEndpoint_doesNotReceiveOllamaPreloadRequest() async throws { + let client = makeClient() + EndpointStubURLProtocol.handler = { _ in + XCTFail("A generic endpoint must not receive Ollama's native preload request") + throw URLError(.badServerResponse) + } + + let didPreload = try await client.preloadDefaultOllamaModel( + configuration: configuration(baseURL: "https://models.example.com/v1"), + apiKey: nil + ) + + XCTAssertFalse(didPreload) + } + func test_chatGeneration_postsSingleUserMessageAndReadsDeltaContent() async throws { let client = makeClient() EndpointStubURLProtocol.handler = { request in