diff --git a/src/main/java/ee/carlrobert/codegpt/settings/service/FeatureType.java b/src/main/java/ee/carlrobert/codegpt/settings/service/FeatureType.java
index 7893fb84..6f7dff3f 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/service/FeatureType.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/service/FeatureType.java
@@ -1,11 +1,11 @@
 package ee.carlrobert.codegpt.settings.service;
 
 public enum FeatureType {
-    CHAT,
-    CODE_COMPLETION,
-    AUTO_APPLY,
-    COMMIT_MESSAGE,
-    INLINE_EDIT,
-    NEXT_EDIT,
-    LOOKUP
-}
\ No newline at end of file
+  CHAT,
+  CODE_COMPLETION,
+  AUTO_APPLY,
+  COMMIT_MESSAGE,
+  INLINE_EDIT,
+  NEXT_EDIT,
+  LOOKUP
+}
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
index 6ac15c17..b1c283f9 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/completions/CompletionRequestFactory.kt
@@ -1,9 +1,7 @@
 package ee.carlrobert.codegpt.completions
 
 import com.intellij.openapi.components.service
-import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.vfs.LocalFileSystem
-import com.intellij.openapi.vfs.readText
 import ee.carlrobert.codegpt.completions.factory.*
 import ee.carlrobert.codegpt.psistructure.ClassStructureSerializer
 import ee.carlrobert.codegpt.settings.prompts.CoreActionsState
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt
index 9fc07dfa..dcfe1318 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelRegistry.kt
@@ -123,8 +123,8 @@ class ModelRegistry {
             ),
             FeatureType.INLINE_EDIT to ModelSelection(
                 ServiceType.PROXYAI,
-                GPT_5_MINI,
-                "GPT-5 Mini"
+                GEMINI_FLASH_2_5,
+                "Gemini Flash 2.5"
             ),
             FeatureType.LOOKUP to ModelSelection(
                 ServiceType.PROXYAI,
@@ -202,7 +202,11 @@ class ModelRegistry {
             GPT_5_MINI,
             "GPT-5 Mini"
         ),
-        FeatureType.INLINE_EDIT to ModelSelection(ServiceType.PROXYAI, GPT_5_MINI, "GPT-5 Mini"),
+        FeatureType.INLINE_EDIT to ModelSelection(
+            ServiceType.PROXYAI,
+            GEMINI_FLASH_2_5,
+            "Gemini Flash 2.5"
+        ),
         FeatureType.LOOKUP to ModelSelection(ServiceType.PROXYAI, GPT_5_MINI, "GPT-5 Mini"),
         FeatureType.CODE_COMPLETION to ModelSelection(
             ServiceType.PROXYAI,
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 20ec5286..1aa5c5ad 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettings.kt
@@ -60,6 +60,7 @@ class ModelSettings : SimplePersistentStateComponent<ModelSettingsState>(ModelSe
 
         migrateCustomOpenAIModelCodesToIds()
         migrateMissingProviderInformation()
+        migrateEditCodeModel()
         notifyIfChanged(oldState, this.state)
     }
 
@@ -113,7 +114,7 @@ class ModelSettings : SimplePersistentStateComponent<ModelSettingsState>(ModelSe
             val newModel = getModelFromState(newState, featureType)
 
             if (oldModel != newModel) {
-                val service = findServiceTypeForModel(featureType, newModel)
+                val service = getProviderForFeature(featureType) ?: return
                 notifyModelChange(featureType, newModel, service)
             }
         }
@@ -123,10 +124,6 @@ class ModelSettings : SimplePersistentStateComponent<ModelSettingsState>(ModelSe
         return state.getModelSelection(featureType)?.model
     }
 
-    private fun findServiceTypeForModel(featureType: FeatureType, modelId: String?): ServiceType {
-        return ServiceType.CUSTOM_OPENAI
-    }
-
     private fun migrateMissingProviderInformation() {
         FeatureType.entries.forEach { featureType ->
             val modelDetailsState = getModelDetailsState(featureType)
@@ -140,6 +137,13 @@ class ModelSettings : SimplePersistentStateComponent<ModelSettingsState>(ModelSe
         }
     }
 
+    private fun migrateEditCodeModel() {
+        state.modelSelections["EDIT_CODE"]?.let {
+            state.setModelSelection(FeatureType.INLINE_EDIT, it.model, it.provider!!)
+            state.modelSelections.remove("EDIT_CODE")
+        }
+    }
+
     private fun inferProviderFromModelCode(
         featureType: FeatureType,
         modelCode: String
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt
index da35ac48..e789c47a 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsState.kt
@@ -5,7 +5,7 @@ import ee.carlrobert.codegpt.settings.service.FeatureType
 import ee.carlrobert.codegpt.settings.service.ServiceType
 
 class ModelSettingsState : BaseState() {
-    var modelSelections by map<FeatureType, ModelDetailsState>()
+    var modelSelections by map<String, ModelDetailsState>()
 
     init {
         if (modelSelections.isEmpty()) {
@@ -17,18 +17,16 @@ class ModelSettingsState : BaseState() {
         val registry = ModelRegistry.getInstance()
         FeatureType.entries.forEach { featureType ->
             val defaultModel = registry.getDefaultModelForFeature(featureType)
-            if (defaultModel != null) {
-                setModelSelection(featureType, defaultModel.model, defaultModel.provider)
-            }
+            setModelSelection(featureType, defaultModel.model, defaultModel.provider)
         }
     }
 
     fun getModelSelection(featureType: FeatureType): ModelDetailsState? {
-        return modelSelections[featureType]
+        return modelSelections[featureType.name]
     }
 
     fun setModelSelection(featureType: FeatureType, model: String?, provider: ServiceType) {
-        modelSelections[featureType] = ModelDetailsState().apply {
+        modelSelections[featureType.name] = ModelDetailsState().apply {
             this.model = model
             this.provider = provider
         }
diff --git a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt
index c99e9c87..23c3f340 100644
--- a/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt
+++ b/src/test/kotlin/ee/carlrobert/codegpt/settings/models/ModelSettingsTest.kt
@@ -41,7 +41,7 @@ class ModelSettingsTest : IntegrationTest() {
                 lastNotification.set(NotificationData(FeatureType.COMMIT_MESSAGE, newModel, serviceType, "commitMessage"))
             }
             override fun inlineEditModelChanged(newModel: String, serviceType: ServiceType) {
-                lastNotification.set(NotificationData(FeatureType.INLINE_EDIT, newModel, serviceType, "editCode"))
+                lastNotification.set(NotificationData(FeatureType.INLINE_EDIT, newModel, serviceType, "inlineEdit"))
             }
             override fun nextEditModelChanged(newModel: String, serviceType: ServiceType) {
                 lastNotification.set(NotificationData(FeatureType.NEXT_EDIT, newModel, serviceType, "nextEdit"))
@@ -129,7 +129,7 @@ class ModelSettingsTest : IntegrationTest() {
     }
 
     fun `test setModelWithProvider with next edit triggers next edit notification`() {
-        modelSettings.state.modelSelections.remove(FeatureType.NEXT_EDIT)
+        modelSettings.state.modelSelections.remove("NEXT_EDIT")
         lastNotification.set(null)
         
         modelSettings.setModelWithProvider(FeatureType.NEXT_EDIT, "zeta", ServiceType.PROXYAI)
@@ -204,7 +204,7 @@ class ModelSettingsTest : IntegrationTest() {
         val detailsState = ModelDetailsState()
         detailsState.model = "gpt-4o"
         detailsState.provider = null
-        state.modelSelections[FeatureType.CHAT] = detailsState
+        state.modelSelections["CHAT"] = detailsState
 
         modelSettings.loadState(state)
 
@@ -217,7 +217,7 @@ class ModelSettingsTest : IntegrationTest() {
         val detailsState = ModelDetailsState()
         detailsState.model = "unknown-model"
         detailsState.provider = null
-        state.modelSelections[FeatureType.CHAT] = detailsState
+        state.modelSelections["CHAT"] = detailsState
 
         modelSettings.loadState(state)