From 6490b1879e5ab23ecfe186228d13c0b5964afb76 Mon Sep 17 00:00:00 2001
From: Carl-Robert Linnupuu <carlrobertoh@gmail.com>
Date: Wed, 8 Oct 2025 12:13:56 +0100
Subject: [PATCH] fix: Next Edit model migration (ProxyAI)

---
 .../codegpt/settings/models/ModelSettings.kt       | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt
index fbd08608..0d862cb0 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt
@@ -64,7 +64,7 @@ class ModelSettings : SimplePersistentStateComponent<ModelSettingsState>(ModelSe
         migrateMissingProviderInformation()
         migrateEditCodeModel()
         migrateProxyAIApplyModel()
-        migrateProxyAICodeModels()
+        migrateProxyAIAutocompleteAndNextEditModels()
         notifyIfChanged(oldState, this.state)
     }
 
@@ -158,11 +158,13 @@ class ModelSettings : SimplePersistentStateComponent<ModelSettingsState>(ModelSe
         }
     }
 
-    private fun migrateProxyAICodeModels() {
-        val modelSelection =
-            service<ModelSelectionService>().getModelSelectionForFeature(FeatureType.CODE_COMPLETION)
-        if (modelSelection.provider == PROXYAI && modelSelection.model != ModelRegistry.MERCURY_CODER) {
-            setModelWithProvider(FeatureType.CODE_COMPLETION, ModelRegistry.MERCURY_CODER, PROXYAI)
+    private fun migrateProxyAIAutocompleteAndNextEditModels() {
+        val modelService = service<ModelSelectionService>()
+        listOf(FeatureType.CODE_COMPLETION, FeatureType.NEXT_EDIT).forEach {
+            val modelSelection = modelService.getModelSelectionForFeature(it)
+            if (modelSelection.provider == PROXYAI && modelSelection.model != ModelRegistry.MERCURY_CODER) {
+                setModelWithProvider(it, ModelRegistry.MERCURY_CODER, PROXYAI)
+            }
         }
     }