fix(connections): safe mode level no longer resets when opening a table (#1351)

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Safe mode no longer turns off when you open another table; it stays set for the whole connection until you change it (#1351)
 - Reassigning the Execute Query, Execute All Statements, and Cancel Query shortcuts now takes effect, and the Query menu shows the new keys (#1357)
 - Custom shortcuts now require a modifier key, so a plain key like Space is no longer accepted and then silently ignored (#1357)
 - Cancelling a pending connection no longer lets the abandoned attempt overwrite or drop a later successful connection to the same database (#1358)

TablePro/Core/Database/DatabaseManager+Sessions.swift

@@ -362,6 +362,12 @@ extension DatabaseManager {
         setSession(session, for: sessionId)
     }
 
+    func setSafeModeLevel(_ level: SafeModeLevel, for connectionId: UUID) {
+        guard var session = activeSessions[connectionId], session.safeModeLevel != level else { return }
+        session.safeModeLevel = level
+        setSession(session, for: connectionId)
+    }
+
     internal func setSession(_ session: ConnectionSession, for connectionId: UUID) {
         activeSessions[connectionId] = session
         connectionStatusVersions[connectionId, default: 0] &+= 1

TablePro/Core/MCP/MCPAuthPolicy.swift

@@ -298,7 +298,7 @@ public actor MCPAuthPolicy {
                     externalAccess: conn.externalAccess,
                     name: conn.name,
                     databaseType: conn.type.rawValue,
-                    safeModeLevel: conn.safeModeLevel
+                    safeModeLevel: session.safeModeLevel
                 )
             case .stored(let conn):
                 return ConnectionSnapshot(

TablePro/Core/MCP/Protocol/Tools/ToolConnectionMetadata.swift

@@ -11,7 +11,7 @@ struct ToolConnectionMetadata {
             case .live(_, let session):
                 return ToolConnectionMetadata(
                     databaseType: session.connection.type,
-                    safeModeLevel: session.connection.safeModeLevel,
+                    safeModeLevel: session.safeModeLevel,
                     databaseName: session.activeDatabase
                 )
             case .stored(let conn):

TablePro/Core/Services/Infrastructure/MainWindowToolbar.swift

@@ -130,7 +130,8 @@ internal final class MainWindowToolbar: NSObject, NSToolbarDelegate {
                 content: ToolbarPrincipalContent(
                     state: coordinator.toolbarState,
                     onSwitchDatabase: { [weak coordinator] in coordinator?.commandActions?.openDatabaseSwitcher() },
-                    onCancelQuery: { [weak coordinator] in coordinator?.cancelCurrentQuery() }
+                    onCancelQuery: { [weak coordinator] in coordinator?.cancelCurrentQuery() },
+                    onSafeModeChange: { [weak coordinator] level in coordinator?.setSafeModeLevel(level) }
                 )
             )
             item.visibilityPriority = .high

TablePro/Models/Connection/ConnectionSession.swift

@@ -17,6 +17,10 @@ struct ConnectionSession: Identifiable {
     var status: ConnectionStatus = .disconnected
     var lastError: String?
 
+    /// Live write-protection level. Seeded from the saved default; the toolbar
+    /// updates this, so reading `connection.safeModeLevel` is wrong afterwards.
+    var safeModeLevel: SafeModeLevel
+
     // Per-connection state
     var selectedTables: Set<TableInfo> = []
     var pendingTruncates: Set<String> = []
@@ -45,6 +49,7 @@ struct ConnectionSession: Identifiable {
         self.id = connection.id
         self.connection = connection
         self.driver = driver
+        self.safeModeLevel = connection.safeModeLevel
         self.connectedAt = Date()
         self.lastActiveAt = Date()
     }

TablePro/Models/Connection/ConnectionToolbarState.swift

@@ -273,7 +273,6 @@ final class ConnectionToolbarState {
         databaseType = connection.type
         displayColor = connection.displayColor
         tagId = connection.tagId
-        safeModeLevel = connection.safeModeLevel
         databaseGroupingStrategy = PluginManager.shared.databaseGroupingStrategy(for: connection.type)
         syncFromSession(for: connection)
     }
@@ -299,6 +298,12 @@ final class ConnectionToolbarState {
         if currentSchema != resolvedSchema {
             currentSchema = resolvedSchema
         }
+
+        let resolvedSafeMode = DatabaseManager.shared.session(for: connection.id)?.safeModeLevel
+            ?? connection.safeModeLevel
+        if safeModeLevel != resolvedSafeMode {
+            safeModeLevel = resolvedSafeMode
+        }
     }
 
     /// Update connection state from ConnectionStatus

TablePro/ViewModels/AIChatViewModel+ToolApproval.swift

@@ -89,7 +89,7 @@ extension AIChatViewModel {
         // Safe-mode level and "Always Allow" cannot bypass them — the AI must not
         // be able to drop tables, truncate, or alter-drop without an explicit click.
         if toolMode == .agentOnly {
-            if let connection, connection.safeModeLevel.blocksAllWrites {
+            if let connection, liveSafeModeLevel(for: connection).blocksAllWrites {
                 return .denied(reason: String(
                     localized: "Connection is read-only. Destructive operations are not permitted."
                 ))
@@ -101,18 +101,24 @@ extension AIChatViewModel {
             return .approved
         }
         if let connection {
-            if connection.safeModeLevel.blocksAllWrites {
+            let safeModeLevel = liveSafeModeLevel(for: connection)
+            if safeModeLevel.blocksAllWrites {
                 return .denied(reason: String(
                     localized: "Connection is read-only. Set safe mode to Confirm Writes or higher to allow this tool."
                 ))
             }
-            if !connection.safeModeLevel.requiresConfirmation {
+            if !safeModeLevel.requiresConfirmation {
                 return .approved
             }
         }
         return .pending
     }
 
+    @MainActor
+    private func liveSafeModeLevel(for connection: DatabaseConnection) -> SafeModeLevel {
+        DatabaseManager.shared.session(for: connection.id)?.safeModeLevel ?? connection.safeModeLevel
+    }
+
     @MainActor
     func appendPendingToolUseBlocks(_ blocks: [ToolUseBlock], assistantID: UUID) {
         guard let idx = messages.firstIndex(where: { $0.id == assistantID }) else { return }

TablePro/Views/Main/MainContentCommandActions.swift

@@ -240,7 +240,7 @@ final class MainContentCommandActions {
     var isConnected: Bool { coordinator != nil }
     var isQueryExecuting: Bool { coordinator?.toolbarState.isExecuting ?? false }
 
-    var safeModeLevel: SafeModeLevel { connection.safeModeLevel }
+    var safeModeLevel: SafeModeLevel { coordinator?.toolbarState.safeModeLevel ?? connection.safeModeLevel }
 
     var isReadOnly: Bool { safeModeLevel.blocksAllWrites }
 

TablePro/Views/Main/MainContentCoordinator.swift

@@ -91,6 +91,10 @@ final class MainContentCoordinator {
         services.databaseManager.activeDatabaseName(for: connection)
     }
     var safeModeLevel: SafeModeLevel { toolbarState.safeModeLevel }
+    func setSafeModeLevel(_ level: SafeModeLevel) {
+        toolbarState.safeModeLevel = level
+        services.databaseManager.setSafeModeLevel(level, for: connectionId)
+    }
     let selectionState = GridSelectionState()
     let tabManager: QueryTabManager
     let changeManager: DataChangeManager

TablePro/Views/Toolbar/TableProToolbarView.swift

@@ -20,6 +20,7 @@ struct ToolbarPrincipalContent: View {
     var state: ConnectionToolbarState
     var onSwitchDatabase: (() -> Void)?
     var onCancelQuery: (() -> Void)?
+    var onSafeModeChange: ((SafeModeLevel) -> Void)?
 
     var body: some View {
         let tag = state.tagId.flatMap { TagStorage.shared.tag(for: $0) }
@@ -40,7 +41,10 @@ struct ToolbarPrincipalContent: View {
                 onSwitchDatabase: onSwitchDatabase
             )
 
-            SafeModeBadgeView(safeModeLevel: Bindable(state).safeModeLevel)
+            SafeModeBadgeView(safeModeLevel: Binding(
+                get: { state.safeModeLevel },
+                set: { onSafeModeChange?($0) }
+            ))
 
             ExecutionIndicatorView(
                 isExecuting: state.isExecuting,

TableProTests/Models/ConnectionSessionTests.swift

@@ -217,4 +217,22 @@ struct ConnectionSessionStateTests {
         let session = ConnectionSession(connection: connection)
         #expect(session.id == connection.id)
     }
+
+    @Test("seeds safe mode from the connection's saved default")
+    func seedsSafeModeFromConnection() {
+        var connection = TestFixtures.makeConnection()
+        connection.safeModeLevel = .readOnly
+        let session = ConnectionSession(connection: connection)
+        #expect(session.safeModeLevel == .readOnly)
+    }
+
+    @Test("clearCachedData preserves safe mode so reconnect keeps protection")
+    func clearCachedDataPreservesSafeMode() {
+        var connection = TestFixtures.makeConnection()
+        connection.safeModeLevel = .silent
+        var session = ConnectionSession(connection: connection)
+        session.safeModeLevel = .readOnly
+        session.clearCachedData()
+        #expect(session.safeModeLevel == .readOnly)
+    }
 }

TableProTests/Models/ConnectionToolbarStateTests.swift

@@ -93,4 +93,31 @@ struct ConnectionToolbarStateTests {
 
         #expect(state.currentDatabase == "Production")
     }
+
+    // MARK: - safe mode
+
+    @Test("syncFromSession falls back to the connection's saved safe mode when no session exists")
+    func syncFromSessionFallsBackToConnectionSafeMode() {
+        var connection = TestFixtures.makeConnection()
+        connection.safeModeLevel = .readOnly
+        let state = ConnectionToolbarState()
+
+        state.syncFromSession(for: connection)
+
+        #expect(state.safeModeLevel == .readOnly)
+    }
+
+    @Test("A new toolbar state adopts the live session safe mode, not the stale saved default")
+    func newToolbarStateAdoptsLiveSessionSafeMode() {
+        let id = UUID()
+        var connection = TestFixtures.makeConnection(id: id)
+        connection.safeModeLevel = .silent
+        DatabaseManager.shared.injectSession(ConnectionSession(connection: connection), for: id)
+        DatabaseManager.shared.setSafeModeLevel(.readOnly, for: id)
+        defer { DatabaseManager.shared.removeSession(for: id) }
+
+        let state = ConnectionToolbarState(connection: connection)
+
+        #expect(state.safeModeLevel == .readOnly)
+    }
 }

docs/features/safe-mode.mdx

@@ -65,6 +65,8 @@ MongoDB, Redis, and other NoSQL databases can't be parsed for read vs. write ope
 
 The current Safe Mode level appears as a badge in the toolbar (orange for Alert levels, red for Safe Mode/Read Only). Click it to change levels.
 
+Changing the level from the badge applies to the whole connection and stays set as you open other tables and tabs. The saved default in the Customization pane is the level a connection starts at.
+
 Safe Mode gates apply to query execution, saving cell edits, table operations, and sidebar changes.
 
 ## External Clients