diff --git a/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift b/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift index 7e5c24d8..e2709886 100644 --- a/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift +++ b/Cotabby/Services/Runtime/OpenAICompatibleSuggestionEngine.swift @@ -9,6 +9,8 @@ final class OpenAICompatibleSuggestionEngine: SuggestionGenerating { private let client: OpenAICompatibleAPIClient private let configurationProvider: @MainActor () throws -> OpenAICompatibleEndpointConfiguration private let apiKeyProvider: @MainActor () throws -> String? + /// Owned for the engine's lifetime so focus-scoped warmup calls share one in-flight model load. + private let preloadWorkController = OllamaPreloadWorkController() init( client: OpenAICompatibleAPIClient, @@ -132,21 +134,68 @@ final class OpenAICompatibleSuggestionEngine: SuggestionGenerating { 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. + guard configuration.defaultOllamaGenerateURL != nil else { return } + let apiKey = try apiKeyProvider() + await preloadWorkController.run(modelName: configuration.modelName) { [client] in + do { + _ = try await client.preloadDefaultOllamaModel( + configuration: configuration, + apiKey: apiKey + ) + CotabbyLogger.runtime.info( + "Preloaded default Ollama model", + metadata: ["model": .string(configuration.modelName)] + ) + } catch is CancellationError { + // App shutdown may cancel opportunistic work; warmup never becomes user-facing. + } catch { + CotabbyLogger.runtime.warning( + "Default Ollama model preload failed: \(error.localizedDescription)" + ) + } + } } catch { CotabbyLogger.runtime.warning( - "Default Ollama model preload failed: \(error.localizedDescription)" + "Default Ollama preload configuration failed: \(error.localizedDescription)" ) } } } + +/// Coalesces focus-driven Ollama warmups without coupling them to prediction cancellation. +/// +/// `OpenAICompatibleSuggestionEngine` owns one controller for its full lifetime. The first caller +/// for a model creates the load task; later focus changes await that task instead of queuing more +/// `/api/generate` requests. Completed work is removed so a later focus can retry after Ollama has +/// restarted. The flight ID prevents an older waiter from deleting a newer retry for the same model. +@MainActor +final class OllamaPreloadWorkController { + private struct Flight { + let id: UUID + let task: Task + } + + private var flightsByModel: [String: Flight] = [:] + + func run( + modelName: String, + operation: @escaping @MainActor () async -> Void + ) async { + if let existing = flightsByModel[modelName] { + await existing.task.value + return + } + + let id = UUID() + // This unstructured task deliberately outlives cancellation of any one focus caller. Model + // loading is shared infrastructure; tying it to a stale field would recreate the 499 abort. + let task = Task { @MainActor in + await operation() + } + flightsByModel[modelName] = Flight(id: id, task: task) + await task.value + + guard flightsByModel[modelName]?.id == id else { return } + flightsByModel[modelName] = nil + } +} diff --git a/CotabbyTests/OllamaPreloadWorkControllerTests.swift b/CotabbyTests/OllamaPreloadWorkControllerTests.swift new file mode 100644 index 00000000..4a6e6513 --- /dev/null +++ b/CotabbyTests/OllamaPreloadWorkControllerTests.swift @@ -0,0 +1,77 @@ +import XCTest +@testable import Cotabby + +/// Exercises the concurrency boundary independently from URLSession. The endpoint engine owns this +/// controller for its lifetime; each test supplies short-lived operations that stand in for Ollama's +/// native preload request and records how many actually begin. +@MainActor +final class OllamaPreloadWorkControllerTests: XCTestCase { + func test_sameModelConcurrentWarmupsShareOneOperationAndAllowLaterRetry() async { + let controller = OllamaPreloadWorkController() + let started = expectation(description: "First preload started") + var invocationCount = 0 + var finishFirst: CheckedContinuation? + + let first = Task { @MainActor in + await controller.run(modelName: "model-a") { + invocationCount += 1 + started.fulfill() + await withCheckedContinuation { continuation in + finishFirst = continuation + } + } + } + await fulfillment(of: [started], timeout: 1) + + let secondEntered = expectation(description: "Second caller entered the controller") + let second = Task { @MainActor in + secondEntered.fulfill() + await controller.run(modelName: "model-a") { + invocationCount += 1 + } + } + await fulfillment(of: [secondEntered], timeout: 1) + await Task.yield() + XCTAssertEqual(invocationCount, 1) + + finishFirst?.resume() + finishFirst = nil + await first.value + await second.value + + await controller.run(modelName: "model-a") { + invocationCount += 1 + } + XCTAssertEqual(invocationCount, 2, "Completed flights must not block later retry attempts") + } + + func test_differentModelsMayPreloadIndependently() async { + let controller = OllamaPreloadWorkController() + let bothStarted = expectation(description: "Both model preloads started") + bothStarted.expectedFulfillmentCount = 2 + var startedModels = Set() + var continuations: [CheckedContinuation] = [] + + let first = Task { @MainActor in + await controller.run(modelName: "model-a") { + startedModels.insert("model-a") + bothStarted.fulfill() + await withCheckedContinuation { continuations.append($0) } + } + } + let second = Task { @MainActor in + await controller.run(modelName: "model-b") { + startedModels.insert("model-b") + bothStarted.fulfill() + await withCheckedContinuation { continuations.append($0) } + } + } + + await fulfillment(of: [bothStarted], timeout: 1) + XCTAssertEqual(startedModels, Set(["model-a", "model-b"])) + continuations.forEach { $0.resume() } + continuations.removeAll() + await first.value + await second.value + } +}