Reuse MLX's prompt cache after an edit and survive exact cache hits (stacked on #868)

AGENTS.md

@@ -164,9 +164,11 @@ from Accessibility and Input Monitoring.
 
 Runtime generation is split by responsibility:
 
-- `SuggestionEngineRouter`: selects Apple Intelligence, Open Source, or the configured
+- `SuggestionEngineRouter`: selects Apple Intelligence, Open Source, MLX, or the configured
   OpenAI-compatible endpoint and owns the narrow Apple-to-llama language fallback.
 - `FoundationModelSuggestionEngine`: Apple on-device generation path.
+- `MLXServerManager`: Cotabby's `mlx_lm.server` child process (loopback only), with
+  `MLXToolchainProvisioner` and `MLXModelDownloader` for its install and models.
 - `LlamaSuggestionEngine`: request-to-prompt, llama result handling, and cache reset handoff.
 - `OpenAICompatibleSuggestionEngine`: completion/chat transport, SSE streaming, and Ollama preload
   behavior for the configured endpoint.
@@ -256,6 +258,9 @@ Console.app stream.
 - `~/Library/Logs/Cotabby/llm-io.jsonl` — full LLM prompts and completions, one record per
   generation. Shares `request_id` with the main log so a single suggestion can be joined across
   files.
+- `~/Library/Logs/<app>/mlx-server.log` — everything Cotabby's `mlx_lm.server` wrote to stderr,
+  one header line (pid, model, port) per server start; rotated to `mlx-server.log.1` at 2 MB.
+  Without the flag only the last lines stay in memory for the MLX failure message.
 - `~/Desktop/cotabby-ax-dump.txt` — most recent Chrome AX tree snapshot. Overwritten on each
   Chrome focus change (debounced by focused-element identity).
 - Rotated previous logs: `*.jsonl.1` (one-step rotation when a file exceeds 10 MB).
@@ -295,6 +300,7 @@ jq 'select(.category == "runtime")' ~/Library/Logs/Cotabby/cotabby.jsonl
 - Wrong text inserted → look up the request in `llm-io.jsonl`, then walk `suggestion` for
   acceptance
 - Model won't load / decode fails → `runtime` + `models`
+- MLX failed or keeps restarting → `runtime`, then `mlx-server.log`
 - Permission dialog loop → `app` (permission state transitions)
 - Chrome-specific weirdness → start with `~/Desktop/cotabby-ax-dump.txt`, then `focus`
 - Wrong backend chosen → `suggestion` router selection log (`engine`, `fallback_engine`)

ARCHITECTURE.md

@@ -78,7 +78,7 @@ CGEvent and focus poll
        -> SuggestionRequestFactory
             -> bounded AX / surface / clipboard / visual / user context
        -> SuggestionEngineRouter
-            -> Apple | in-process llama | OpenAI-compatible endpoint
+            -> Apple | in-process llama | MLX (loopback mlx_lm.server) | OpenAI-compatible endpoint
        -> SuggestionTextNormalizer + seam guards
        -> SuggestionInteractionState
        -> SuggestionOverlayPresenter -> OverlayController
@@ -257,6 +257,12 @@ text. Posting events is not treated as proof of success; the later AX snapshot i
   for completion/chat APIs and SSE streams. The default is loopback Ollama at
   http://127.0.0.1:11434/v1; LAN and public HTTPS endpoints are supported, while insecure public HTTP
   is rejected.
+- MLX: `OpenAICompatibleSuggestionEngine.mlx` talks to an `mlx_lm.server` child process that
+  [MLXServerManager.swift](Cotabby/Services/Runtime/MLX/MLXServerManager.swift) runs on 127.0.0.1.
+  `MLXToolchainProvisioner` finds mlx-lm or installs it with a checksum- and signature-pinned uv
+  under Application Support; `MLXModelDownloader` manages MLX models in the Hugging Face cache;
+  `MLXServerDesireResolver` decides when the server runs. MLX is on-device (typing history
+  allowed) but takes no speculative work.
 
 LlamaRuntimeCore is a nonisolated lock/condition-protected native boundary, not a Swift actor. Its
 autocomplete lock serializes cache/decode state, and its lifecycle condition prevents shutdown from

Cotabby/App/Coordinators/SettingsCoordinator.swift

@@ -16,6 +16,9 @@ final class SettingsCoordinator: NSObject, NSWindowDelegate {
     private let permissionGuidanceController: PermissionGuidanceController
     private let suggestionSettings: SuggestionSettingsModel
     private let openAICompatibleConnectionModel: OpenAICompatibleConnectionModel
+    private let mlxToolchainProvisioner: MLXToolchainProvisioner
+    private let mlxServerManager: MLXServerManager
+    private let mlxModelDownloader: MLXModelDownloader
     private let foundationModelAvailabilityService: FoundationModelAvailabilityService
     private let runtimeModel: RuntimeBootstrapModel
     private let modelDownloadManager: ModelDownloadManager
@@ -42,6 +45,9 @@ final class SettingsCoordinator: NSObject, NSWindowDelegate {
         permissionGuidanceController: PermissionGuidanceController,
         suggestionSettings: SuggestionSettingsModel,
         openAICompatibleConnectionModel: OpenAICompatibleConnectionModel,
+        mlxToolchainProvisioner: MLXToolchainProvisioner,
+        mlxServerManager: MLXServerManager,
+        mlxModelDownloader: MLXModelDownloader,
         foundationModelAvailabilityService: FoundationModelAvailabilityService,
         runtimeModel: RuntimeBootstrapModel,
         modelDownloadManager: ModelDownloadManager,
@@ -58,6 +64,9 @@ final class SettingsCoordinator: NSObject, NSWindowDelegate {
         self.permissionGuidanceController = permissionGuidanceController
         self.suggestionSettings = suggestionSettings
         self.openAICompatibleConnectionModel = openAICompatibleConnectionModel
+        self.mlxToolchainProvisioner = mlxToolchainProvisioner
+        self.mlxServerManager = mlxServerManager
+        self.mlxModelDownloader = mlxModelDownloader
         self.foundationModelAvailabilityService = foundationModelAvailabilityService
         self.runtimeModel = runtimeModel
         self.modelDownloadManager = modelDownloadManager
@@ -88,6 +97,9 @@ final class SettingsCoordinator: NSObject, NSWindowDelegate {
                     permissionGuidanceController: permissionGuidanceController,
                     suggestionSettings: suggestionSettings,
                     openAICompatibleConnectionModel: openAICompatibleConnectionModel,
+                    mlxToolchainProvisioner: mlxToolchainProvisioner,
+                    mlxServerManager: mlxServerManager,
+                    mlxModelDownloader: mlxModelDownloader,
                     foundationModelAvailabilityService: foundationModelAvailabilityService,
                     runtimeModel: runtimeModel,
                     modelDownloadManager: modelDownloadManager,

Cotabby/App/Coordinators/Suggestion/SuggestionCoordinator+Continuation.swift

@@ -32,10 +32,12 @@ extension SuggestionCoordinator {
     /// Only terminal word endings need another request. A model-provided phrase stays in its
     /// original session; prefetch never replaces those words with a different model guess.
     func prepareContinuation(after session: ActiveSuggestionSession, rawContext: FocusedInputSnapshot) {
-        // Uncommitted lookahead is new work, so keep it on-device. The coordinator intentionally
-        // has no endpoint URL; allowing all configured endpoints here would also send hypothetical
-        // corrections to hosted servers. Their already-returned phrases still use the same buffer.
-        guard settingsSnapshot.selectedEngine != .openAICompatible,
+        // Uncommitted lookahead is new work the user did not ask for. Only the in-process engines
+        // yield it to the real request at once; on a server (the endpoint or MLX) it shares the
+        // generation work with the suggestion the user waits for until its cancellation arrives,
+        // and the endpoint would also receive hypothetical corrections. Already-returned phrases
+        // still use the same buffer.
+        guard settingsSnapshot.selectedEngine.allowsSpeculativeGeneration,
               !userDefaults.bool(forKey: Self.speculativePrefetchDisabledDefaultsKey),
               currentDisabledReason(focusSnapshot: focusModel.snapshot) == nil else { return }
         let plan: SuggestionContinuationPlan?

Cotabby/App/Coordinators/Suggestion/SuggestionCoordinator+TypingPrediction.swift

@@ -6,10 +6,11 @@ import Foundation
 extension SuggestionCoordinator {
     func beginTypingPrediction(for request: SuggestionRequest) {
         clearTypingPrediction()
-        // Internal partial collection must not add speculative traffic to configured endpoints.
+        // Internal partial collection must not add speculative requests to a server (the endpoint
+        // or MLX), where they would share generation work with the request the user waits for.
         // Post-acceptance speculation has its own publication contract and is kept separate.
         guard settingsSnapshot.predictAheadWhileTyping,
-              settingsSnapshot.selectedEngine != .openAICompatible,
+              settingsSnapshot.selectedEngine.allowsSpeculativeGeneration,
               pendingSpeculativeContext == nil, !request.context.isSecure,
               request.context.selection.length == 0,
               !userDefaults.bool(forKey: Self.speculativePrefetchDisabledDefaultsKey) else { return }

Cotabby/App/Core/AppDelegate.swift

@@ -25,6 +25,8 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
     let suggestionSettings: SuggestionSettingsModel
     let foundationModelAvailabilityService: FoundationModelAvailabilityService
     let powerSourceMonitor: PowerSourceMonitor
+    let mlxServerManager: MLXServerManager
+    let menuResourceMonitor: MenuResourceMonitor
     let suggestionCoordinator: SuggestionCoordinator
     let inlineCommandCoordinator: InlineCommandCoordinator
     let welcomeCoordinator: WelcomeCoordinator
@@ -57,6 +59,8 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
         suggestionSettings = environment.suggestionSettings
         foundationModelAvailabilityService = environment.foundationModelAvailabilityService
         powerSourceMonitor = environment.powerSourceMonitor
+        mlxServerManager = environment.mlxServerManager
+        menuResourceMonitor = environment.menuResourceMonitor
         suggestionCoordinator = environment.suggestionCoordinator
         inlineCommandCoordinator = environment.inlineCommandCoordinator
         welcomeCoordinator = environment.welcomeCoordinator
@@ -164,6 +168,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
         )
         applyLaunchAtLoginDefaultIfNeeded()
         startRuntimeIfPreferredEngineRequiresIt()
+        environment.startMLXEngine()
         focusModel.start()
         inputMonitor.start()
         appUpdateManager.start()
@@ -256,6 +261,9 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
         inputMonitor.stop()
         focusModel.stop()
 
+        // Waits at most 3 s for mlx_lm.server to exit (SIGKILL after that), so quitting Cotabby
+        // never leaves a Python process holding the model's memory.
+        environment.mlxServerManager.shutdownSync(timeoutSeconds: 3)
         runtimeModel.shutdownSync(timeoutSeconds: 1.5)
     }
 
@@ -296,8 +304,9 @@ final class AppDelegate: NSObject, NSApplicationDelegate {
             runtimeModel.startIfNeeded()
         } else {
             // Switching away must release Metal buffers and the mapped GGUF. Otherwise an Ollama
-            // user still pays the duplicate memory cost the external endpoint is meant to avoid, and
-            // an Apple Intelligence user keeps a model they asked not to keep loaded.
+            // user still pays the duplicate memory cost the external endpoint is meant to avoid, an
+            // Apple Intelligence user keeps a model they asked not to keep loaded, and an MLX user
+            // holds a second copy of the weights beside the one mlx_lm.server loaded.
             runtimeModel.stop()
         }
     }

Cotabby/App/Core/CotabbyApp.swift

@@ -27,6 +27,8 @@ struct CotabbyApp: App {
                 suggestionSettings: appDelegate.suggestionSettings,
                 foundationModelAvailabilityService: appDelegate.foundationModelAvailabilityService,
                 powerSourceMonitor: appDelegate.powerSourceMonitor,
+                mlxServerManager: appDelegate.mlxServerManager,
+                menuResourceMonitor: appDelegate.menuResourceMonitor,
                 appUpdateManager: appDelegate.appUpdateManager,
                 onOpenSettings: {
                     appDelegate.settingsCoordinator.showSettings()