refactor(editor): explicit state machines for text binding sync and completion sessions (#1616)

* fix(editor): re-highlight the document after a programmatic text replacement (#1612)

* refactor(editor): model text binding sync as an explicit origin state machine

* refactor(editor): model completion as an explicit session with apply-time trigger suppression

Author: datlechin

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Cancelling a SQLite query no longer races a disconnect happening at the same moment. (#1610)
 - Typing in the query editor no longer erases characters or drops focus on each keystroke, a timing-dependent bug most visible on macOS 15. (#1608)
 - The autocomplete popup now filters in place as you type instead of closing and reopening on every keystroke. (#1608)
+- Syntax highlighting no longer disappears after formatting a query. (#1612)
 
 ## [0.49.1] - 2026-06-06
 

LocalPackages/CodeEditSourceEditor/Sources/CodeEditSourceEditor/CodeSuggestion/Model/SuggestionViewModel.swift

@@ -19,6 +19,7 @@ final class SuggestionViewModel: ObservableObject {
 
     var itemsRequestTask: Task<Void, Never>?
     weak var activeTextView: TextViewController?
+    private(set) var isApplyingCompletion = false
 
     weak var delegate: CodeSuggestionDelegate?
 
@@ -83,6 +84,8 @@ final class SuggestionViewModel: ObservableObject {
         isManualTrigger: Bool = false,
         showWindowOnParent: @escaping @MainActor (NSWindow, NSRect) -> Void
     ) {
+        guard !isApplyingCompletion else { return }
+
         self.activeTextView = nil
         self.delegate = nil
         itemsRequestTask?.cancel()
@@ -142,6 +145,8 @@ final class SuggestionViewModel: ObservableObject {
         position: CursorPosition,
         close: () -> Void
     ) {
+        guard !isApplyingCompletion else { return }
+
         if activeTextView !== textView {
             itemsRequestTask?.cancel()
             itemsRequestTask = nil
@@ -173,11 +178,13 @@ final class SuggestionViewModel: ObservableObject {
         guard let activeTextView else {
             return
         }
+        isApplyingCompletion = true
         self.delegate?.completionWindowApplyCompletion(
             item: item,
             textView: activeTextView,
             cursorPosition: activeTextView.cursorPositions.first
         )
+        isApplyingCompletion = false
         onApply?()
     }
 
@@ -187,6 +194,7 @@ final class SuggestionViewModel: ObservableObject {
         items.removeAll()
         selectedIndex = 0
         activeTextView = nil
+        delegate?.completionWindowDidClose()
         delegate = nil
     }
 

LocalPackages/CodeEditSourceEditor/Sources/CodeEditSourceEditor/Controller/TextViewController+Highlighter.swift

@@ -9,6 +9,13 @@ import Foundation
 import SwiftTreeSitter
 
 extension TextViewController {
+    /// Tears down and rebuilds the highlighter for the text view's current storage.
+    ///
+    /// Ends with an explicit `invalidate()` so the rebuilt highlighter queries the
+    /// visible text immediately. A fresh highlighter only highlights in response to
+    /// triggers (an edit, a frame change, or an invalidation); after a mid-session
+    /// storage swap such as `setText`, none of those is guaranteed to fire, and
+    /// without this the document stays unstyled until the next layout change.
     package func setUpHighlighter() {
         if let highlighter {
             textView.removeStorageDelegate(highlighter)
@@ -24,6 +31,7 @@ extension TextViewController {
         )
         textView.addStorageDelegate(highlighter)
         self.highlighter = highlighter
+        highlighter.invalidate()
     }
 
     /// Sets new highlight providers. Recognizes when objects move in the array or are removed or inserted.

LocalPackages/CodeEditSourceEditor/Sources/CodeEditSourceEditor/SourceEditor/RepresentableSyncPhase.swift

@@ -0,0 +1,59 @@
+//
+//  RepresentableSyncPhase.swift
+//  CodeEditSourceEditor
+//
+//  Tracks which side of the SwiftUI representable boundary originated the
+//  change currently being synchronized, replacing the pair of booleans the
+//  coordinator previously juggled. The states are mutually exclusive by
+//  construction: an editor change can never be marked while a representable
+//  value is being applied, and a representable value is never applied while
+//  an editor change is pending.
+//
+
+import Foundation
+
+@MainActor
+final class RepresentableSyncPhase {
+    enum Phase: Equatable {
+        case idle
+        case editorChangePending
+        case applyingRepresentableValue
+    }
+
+    private(set) var phase: Phase = .idle
+
+    var isEditorChangePending: Bool {
+        phase == .editorChangePending
+    }
+
+    var isApplyingRepresentableValue: Bool {
+        phase == .applyingRepresentableValue
+    }
+
+    /// Latches that the editor originated a change. The next representable
+    /// update pass consumes this instead of pushing its own values down.
+    /// Ignored while a representable value is being applied, since editor
+    /// notifications fired during a programmatic application are echoes,
+    /// not user edits.
+    func markEditorChange() {
+        guard phase != .applyingRepresentableValue else { return }
+        phase = .editorChangePending
+    }
+
+    /// Returns whether an editor change was pending and resets to idle.
+    @discardableResult
+    func consumePendingEditorChange() -> Bool {
+        guard phase == .editorChangePending else { return false }
+        phase = .idle
+        return true
+    }
+
+    /// Runs `body` with the phase marked as applying a representable value,
+    /// so editor notifications fired by the application itself are ignored.
+    func applyRepresentableValue(_ body: () -> Void) {
+        let previous = phase
+        phase = .applyingRepresentableValue
+        body()
+        phase = previous
+    }
+}

LocalPackages/CodeEditSourceEditor/Sources/CodeEditSourceEditor/SourceEditor/SourceEditor+Coordinator.swift

@@ -14,23 +14,18 @@ extension SourceEditor {
     @MainActor
     public class Coordinator: NSObject {
         private weak var controller: TextViewController?
-        var isUpdatingFromRepresentable: Bool = false
-        var isUpdateFromTextView: Bool = false
-        var text: TextAPI
-        var lastSyncedText: String?
+        let phase = RepresentableSyncPhase()
+        let textSync: TextBindingSync
         @Binding var editorState: SourceEditorState
 
         private(set) var highlightProviders: [any HighlightProviding]
 
         private var cancellables: Set<AnyCancellable> = []
 
         init(text: TextAPI, editorState: Binding<SourceEditorState>, highlightProviders: [any HighlightProviding]?) {
-            self.text = text
+            self.textSync = TextBindingSync(text: text, phase: phase)
             self._editorState = editorState
             self.highlightProviders = highlightProviders ?? [TreeSitterClient()]
-            if case .binding(let binding) = text {
-                self.lastSyncedText = binding.wrappedValue
-            }
             super.init()
         }
 
@@ -123,57 +118,11 @@ extension SourceEditor {
             self.highlightProviders = highlightProviders
         }
 
-        private var textBindingTask: Task<Void, Never>?
-
         @objc func textViewDidChangeText(_ notification: Notification) {
             guard let textView = notification.object as? TextView else {
                 return
             }
-            guard !isUpdatingFromRepresentable else { return }
-            guard case .binding(let binding) = text else { return }
-
-            // For large documents, debounce the binding writeback to avoid
-            // copying megabytes of text into SwiftUI on every keystroke.
-            let docLength = textView.textStorage.length
-            // Set flag immediately so SwiftUI's updateNSViewController knows
-            // the text view is the source of truth during the debounce window.
-            isUpdateFromTextView = true
-            if docLength > 500_000 {
-                textBindingTask?.cancel()
-                textBindingTask = Task { @MainActor [weak self, weak textView] in
-                    try? await Task.sleep(for: .milliseconds(150))
-                    guard !Task.isCancelled, let self, let textView else { return }
-                    guard case .binding(let currentBinding) = self.text else { return }
-                    let newText = textView.string
-                    self.lastSyncedText = newText
-                    currentBinding.wrappedValue = newText
-                }
-            } else {
-                let newText = textView.string
-                lastSyncedText = newText
-                binding.wrappedValue = newText
-            }
-        }
-
-        /// Pushes an external binding change down into the text view. The text view's
-        /// content wins while one of its own edits is still in flight: `isUpdateFromTextView`
-        /// is set the moment the text view mutates, before SwiftUI re-renders, so a render
-        /// that still carries a stale binding snapshot is skipped entirely and user typing
-        /// is never clobbered by a stale binding.
-        ///
-        /// Uses `setText` rather than `replaceCharacters` on purpose: `replaceCharacters`
-        /// is the user-edit path. It is gated on `isEditable`, runs mutation filters, and
-        /// fires suggestion triggers, none of which should happen for a programmatic
-        /// whole-document replacement. `setText` clearing the undo stack matches the
-        /// new-document semantics of that replacement.
-        func syncBindingText(_ newValue: String, controller: TextViewController) {
-            guard !isUpdateFromTextView else { return }
-            guard newValue != lastSyncedText else { return }
-            textBindingTask?.cancel()
-            isUpdatingFromRepresentable = true
-            controller.textView.setText(newValue)
-            isUpdatingFromRepresentable = false
-            lastSyncedText = newValue
+            textSync.editorTextDidChange(textView)
         }
 
         @objc func textControllerCursorsDidUpdate(_ notification: Notification) {
@@ -218,8 +167,8 @@ extension SourceEditor {
         }
 
         private func updateState(_ modifyCallback: (inout SourceEditorState) -> Void) {
-            guard !isUpdatingFromRepresentable else { return }
-            self.isUpdateFromTextView = true
+            guard !phase.isApplyingRepresentableValue else { return }
+            phase.markEditorChange()
             modifyCallback(&editorState)
         }
 

LocalPackages/CodeEditSourceEditor/Sources/CodeEditSourceEditor/SourceEditor/SourceEditor.swift

@@ -134,18 +134,15 @@ public struct SourceEditor: NSViewControllerRepresentable {
 
         context.coordinator.updateHighlightProviders(highlightProviders)
 
-        context.coordinator.text = text
+        context.coordinator.textSync.text = text
         if case .binding(let binding) = text {
-            context.coordinator.syncBindingText(binding.wrappedValue, controller: controller)
+            context.coordinator.textSync.applyRepresentableText(binding.wrappedValue, controller: controller)
         }
 
-        // Prevent infinite loop of update notifications
-        if context.coordinator.isUpdateFromTextView {
-            context.coordinator.isUpdateFromTextView = false
-        } else {
-            context.coordinator.isUpdatingFromRepresentable = true
-            updateControllerWithState(state, controller: controller)
-            context.coordinator.isUpdatingFromRepresentable = false
+        if !context.coordinator.phase.consumePendingEditorChange() {
+            context.coordinator.phase.applyRepresentableValue {
+                updateControllerWithState(state, controller: controller)
+            }
         }
 
         // Do manual diffing to reduce the amount of reloads.

LocalPackages/CodeEditSourceEditor/Sources/CodeEditSourceEditor/SourceEditor/TextBindingSync.swift

@@ -0,0 +1,88 @@
+//
+//  TextBindingSync.swift
+//  CodeEditSourceEditor
+//
+//  Owns the two-way synchronization between a SwiftUI text binding and the
+//  text view: editor edits are written back to the binding (debounced for
+//  large documents), and external binding changes are pushed down into the
+//  editor. The shared ``RepresentableSyncPhase`` decides which side wins
+//  when both have changes in flight.
+//
+
+import AppKit
+import CodeEditTextView
+import SwiftUI
+
+@MainActor
+final class TextBindingSync {
+    private static let writebackDebounceThreshold = 500_000
+    private static let writebackDebounce: Duration = .milliseconds(150)
+
+    var text: SourceEditor.TextAPI
+    private(set) var lastSyncedText: String?
+
+    private let phase: RepresentableSyncPhase
+    private var writebackTask: Task<Void, Never>?
+
+    init(text: SourceEditor.TextAPI, phase: RepresentableSyncPhase) {
+        self.text = text
+        self.phase = phase
+        if case .binding(let binding) = text {
+            lastSyncedText = binding.wrappedValue
+        }
+    }
+
+    /// Writes an editor-originated text change back to the binding.
+    ///
+    /// Marks the phase before writing so the representable update pass that
+    /// the binding write triggers cannot clobber the editor with a stale
+    /// snapshot. For large documents the writeback is debounced to avoid
+    /// copying megabytes into SwiftUI on every keystroke; the phase is still
+    /// marked immediately so the editor stays the source of truth during the
+    /// debounce window.
+    func editorTextDidChange(_ textView: TextView) {
+        guard !phase.isApplyingRepresentableValue else { return }
+        guard case .binding(let binding) = text else { return }
+
+        phase.markEditorChange()
+
+        guard textView.textStorage.length > Self.writebackDebounceThreshold else {
+            let newText = textView.string
+            lastSyncedText = newText
+            binding.wrappedValue = newText
+            return
+        }
+
+        writebackTask?.cancel()
+        writebackTask = Task { @MainActor [weak self, weak textView] in
+            try? await Task.sleep(for: Self.writebackDebounce)
+            guard !Task.isCancelled, let self, let textView else { return }
+            guard case .binding(let currentBinding) = self.text else { return }
+            let newText = textView.string
+            self.lastSyncedText = newText
+            currentBinding.wrappedValue = newText
+        }
+    }
+
+    /// Pushes an external binding change down into the text view. The editor's
+    /// content wins while one of its own edits is still in flight, so user
+    /// typing is never clobbered by a stale binding snapshot.
+    ///
+    /// Routes through `TextViewController.setText` on purpose. The text-view
+    /// level `replaceCharacters` is the user-edit path: it is gated on
+    /// `isEditable`, runs mutation filters, and fires suggestion triggers,
+    /// none of which should happen for a programmatic whole-document
+    /// replacement. The controller-level call also rebuilds the highlighter
+    /// for the replacement storage; calling the text view directly leaves
+    /// tree-sitter state for the old document and highlighting never recovers.
+    func applyRepresentableText(_ newValue: String, controller: TextViewController) {
+        guard !phase.isEditorChangePending else { return }
+        guard newValue != lastSyncedText else { return }
+
+        writebackTask?.cancel()
+        phase.applyRepresentableValue {
+            controller.setText(newValue)
+        }
+        lastSyncedText = newValue
+    }
+}