fix: address all code review findings for MCP server

Author: datlechin

TablePro/Core/MCP/MCPAuthGuard.swift

@@ -113,6 +113,9 @@ actor MCPAuthGuard {
     // MARK: - User Approval (askEachTime)
 
     private func promptUserApproval(connectionName: String, databaseType: String) async throws -> Bool {
+        // Use a task group so the actor suspends (freeing it for other requests)
+        // while the approval dialog is shown on the main thread.
+        // Race the dialog against a 30-second timeout.
         let approvalTask = Task { @MainActor in
             NSApp.requestUserAttention(.criticalRequest)
             NSApp.activate(ignoringOtherApps: true)
@@ -129,8 +132,7 @@ actor MCPAuthGuard {
             )
         }
 
-        // Race against a 30-second timeout
-        return try await withThrowingTaskGroup(of: Bool.self) { group in
+        let approved = try await withThrowingTaskGroup(of: Bool.self) { group in
             group.addTask {
                 await approvalTask.value
             }
@@ -146,6 +148,13 @@ actor MCPAuthGuard {
             group.cancelAll()
             return result
         }
+
+        if approved {
+            return true
+        }
+        throw MCPError.forbidden(
+            String(localized: "User denied MCP access to this connection")
+        )
     }
 
     // MARK: - Session Cleanup

TablePro/Core/MCP/MCPConnectionBridge.swift

@@ -45,6 +45,32 @@ actor MCPConnectionBridge {
     func connect(connectionId: UUID) async throws -> JSONValue {
         let connection = try await resolveConnection(connectionId)
 
+        // Check if session already exists and is connected -- reuse without switching UI
+        let existingSession = await MainActor.run {
+            DatabaseManager.shared.activeSessions[connectionId]
+        }
+
+        if let existing = existingSession, existing.driver != nil {
+            // Already connected, return current state without switching the UI's active session
+            let serverVersion = existing.driver?.serverVersion
+            let currentDatabase = existing.activeDatabase
+            let currentSchema = existing.currentSchema
+
+            var result: [String: JSONValue] = [
+                "status": "connected",
+                "current_database": .string(currentDatabase)
+            ]
+            if let version = serverVersion {
+                result["server_version"] = .string(version)
+            }
+            if let schema = currentSchema {
+                result["current_schema"] = .string(schema)
+            }
+            return .object(result)
+        }
+
+        // Not connected yet -- create a new session via DatabaseManager.
+        // connectToSession is @MainActor; Swift hops automatically for async calls.
         try await DatabaseManager.shared.connectToSession(connection)
 
         let (serverVersion, currentDatabase, currentSchema) = await MainActor.run {
@@ -81,43 +107,56 @@ actor MCPConnectionBridge {
     }
 
     func getConnectionStatus(connectionId: UUID) async throws -> JSONValue {
-        let sessionInfo = await MainActor.run {
-            () -> (status: ConnectionStatus, database: String, schema: String?, version: String?, connectedAt: Date)? in
+        let core = await MainActor.run {
+            () -> (status: ConnectionStatus, database: String, schema: String?)? in
             guard let session = DatabaseManager.shared.activeSessions[connectionId] else {
                 return nil
             }
-            return (
-                status: session.status,
-                database: session.activeDatabase,
-                schema: session.currentSchema,
-                version: session.driver?.serverVersion,
-                connectedAt: session.connectedAt
-            )
+            return (session.status, session.activeDatabase, session.currentSchema)
         }
 
-        guard let info = sessionInfo else {
+        guard let core else {
             throw MCPError.notConnected(connectionId)
         }
 
+        let meta = await MainActor.run {
+            () -> (version: String?, connectedAt: Date, lastActiveAt: Date) in
+            let session = DatabaseManager.shared.activeSessions[connectionId]
+            return (
+                session?.driver?.serverVersion,
+                session?.connectedAt ?? Date(),
+                session?.lastActiveAt ?? Date()
+            )
+        }
+
         let statusString: String
-        switch info.status {
+        var errorDetail: JSONValue?
+        switch core.status {
         case .connected: statusString = "connected"
         case .connecting: statusString = "connecting"
         case .disconnected: statusString = "disconnected"
-        case .error(let msg): statusString = "error: \(msg)"
+        case .error(let msg):
+            statusString = "error"
+            errorDetail = .object([
+                "message": .string(msg)
+            ])
         }
 
         var result: [String: JSONValue] = [
             "status": .string(statusString),
-            "current_database": .string(info.database),
-            "connected_at": .string(ISO8601DateFormatter().string(from: info.connectedAt))
+            "current_database": .string(core.database),
+            "connected_at": .string(ISO8601DateFormatter().string(from: meta.connectedAt)),
+            "last_active_at": .string(ISO8601DateFormatter().string(from: meta.lastActiveAt))
         ]
-        if let schema = info.schema {
+        if let schema = core.schema {
             result["current_schema"] = .string(schema)
         }
-        if let version = info.version {
+        if let version = meta.version {
             result["server_version"] = .string(version)
         }
+        if let errorDetail {
+            result["error"] = errorDetail
+        }
 
         return .object(result)
     }
@@ -146,6 +185,8 @@ actor MCPConnectionBridge {
                 }
                 group.addTask {
                     try await Task.sleep(for: .seconds(timeoutSeconds))
+                    // Cancel the driver query before throwing
+                    try? driver.cancelQuery()
                     throw MCPError.timeout("Query timed out after \(timeoutSeconds) seconds")
                 }
                 guard let first = try await group.next() else {
@@ -156,7 +197,7 @@ actor MCPConnectionBridge {
             }
         }
 
-        let executionTimeMs = (CFAbsoluteTimeGetCurrent() - startTime) * 1000
+        let executionTimeMs = (CFAbsoluteTimeGetCurrent() - startTime) * 1_000
         let isTruncated = result.rows.count > maxRows
         let rows = isTruncated ? Array(result.rows.prefix(maxRows)) : result.rows
 
@@ -316,6 +357,7 @@ actor MCPConnectionBridge {
     // MARK: - Database/Schema Switching
 
     func switchDatabase(connectionId: UUID, database: String) async throws -> JSONValue {
+        // switchDatabase is @MainActor; Swift hops automatically for async calls.
         try await DatabaseManager.shared.switchDatabase(to: database, for: connectionId)
         return .object([
             "status": "switched",
@@ -324,6 +366,7 @@ actor MCPConnectionBridge {
     }
 
     func switchSchema(connectionId: UUID, schema: String) async throws -> JSONValue {
+        // switchSchema is @MainActor; Swift hops automatically for async calls.
         try await DatabaseManager.shared.switchSchema(to: schema, for: connectionId)
         return .object([
             "status": "switched",
@@ -334,14 +377,34 @@ actor MCPConnectionBridge {
     // MARK: - Schema Resource (for resources/read)
 
     func fetchSchemaResource(connectionId: UUID) async throws -> JSONValue {
-        let (driver, _) = try await resolveDriver(connectionId)
+        // Check SchemaProviderRegistry cache first
+        let provider = await MainActor.run {
+            SchemaProviderRegistry.shared.provider(for: connectionId)
+        }
+        var cachedTables: [TableInfo] = []
+        if let provider {
+            let cached = await provider.getTables()
+            if !cached.isEmpty {
+                cachedTables = cached
+            }
+        }
 
-        let tables = try await DatabaseManager.shared.trackOperation(sessionId: connectionId) {
-            try await driver.fetchTables()
+        let tables: [TableInfo]
+        if !cachedTables.isEmpty {
+            tables = cachedTables
+        } else {
+            let (driver, _) = try await resolveDriver(connectionId)
+            tables = try await DatabaseManager.shared.trackOperation(sessionId: connectionId) {
+                try await driver.fetchTables()
+            }
         }
 
+        // Limit to first 100 tables to prevent excessive round-trips
+        let limitedTables = Array(tables.prefix(100))
+
+        let (driver, _) = try await resolveDriver(connectionId)
         var tableSchemas: [JSONValue] = []
-        for table in tables {
+        for table in limitedTables {
             let columns = try await DatabaseManager.shared.trackOperation(sessionId: connectionId) {
                 try await driver.fetchColumns(table: table.name)
             }
@@ -362,7 +425,13 @@ actor MCPConnectionBridge {
             ]))
         }
 
-        return .object(["tables": .array(tableSchemas)])
+        var result: [String: JSONValue] = ["tables": .array(tableSchemas)]
+        if tables.count > 100 {
+            result["truncated"] = .bool(true)
+            result["total_tables"] = .int(tables.count)
+        }
+
+        return .object(result)
     }
 
     // MARK: - History Resource
@@ -394,7 +463,7 @@ actor MCPConnectionBridge {
                 "query": .string(entry.query),
                 "database_name": .string(entry.databaseName),
                 "executed_at": .string(ISO8601DateFormatter().string(from: entry.executedAt)),
-                "execution_time_ms": .double(entry.executionTime * 1000),
+                "execution_time_ms": .double(entry.executionTime * 1_000),
                 "row_count": .int(entry.rowCount),
                 "was_successful": .bool(entry.wasSuccessful)
             ]

TablePro/Core/MCP/MCPHTTPParser.swift

@@ -21,6 +21,7 @@ enum HTTPParseError: Error, Sendable {
     case malformedHeaders
     case unsupportedMethod(String)
     case bodyTooLarge
+    case malformedChunkedEncoding
 }
 
 enum MCPHTTPParser {
@@ -90,7 +91,16 @@ enum MCPHTTPParser {
             if let transferEncoding = headers["transfer-encoding"],
                transferEncoding.lowercased().contains("chunked")
             {
-                logger.warning("Chunked transfer encoding not supported, treating body as empty")
+                let bodyData = data[bodyStartIndex...]
+                switch decodeChunkedBody(bodyData) {
+                case .success(let decoded):
+                    if decoded.count > maxBodySize {
+                        return .failure(.bodyTooLarge)
+                    }
+                    body = decoded
+                case .failure(let error):
+                    return .failure(error)
+                }
             } else if let contentLengthStr = headers["content-length"],
                       let contentLength = Int(contentLengthStr)
             {
@@ -115,6 +125,70 @@ enum MCPHTTPParser {
         ))
     }
 
+    // MARK: - Chunked Transfer Encoding
+
+    /// Decode chunked transfer encoding body.
+    /// Format: <chunk-size-hex>\r\n<chunk-data>\r\n ... 0\r\n\r\n
+    private static func decodeChunkedBody(_ data: Data) -> Result<Data, HTTPParseError> {
+        var result = Data()
+        var offset = data.startIndex
+
+        while offset < data.endIndex {
+            // Find the end of the chunk size line
+            guard let lineEnd = findCRLF(in: data, from: offset) else {
+                return .failure(.incomplete)
+            }
+
+            let sizeData = data[offset..<lineEnd]
+            guard let sizeString = String(data: sizeData, encoding: .ascii)?.trimmingCharacters(in: .whitespaces),
+                  let chunkSize = UInt(sizeString, radix: 16)
+            else {
+                return .failure(.malformedChunkedEncoding)
+            }
+
+            // Move past the \r\n after chunk size
+            let chunkDataStart = lineEnd + 2
+
+            // Terminal chunk
+            if chunkSize == 0 {
+                return .success(result)
+            }
+
+            let chunkDataEnd = chunkDataStart + Int(chunkSize)
+
+            // Check we have enough data for the chunk + trailing \r\n
+            guard chunkDataEnd + 2 <= data.endIndex else {
+                return .failure(.incomplete)
+            }
+
+            // Check accumulated size
+            if result.count + Int(chunkSize) > maxBodySize {
+                return .failure(.bodyTooLarge)
+            }
+
+            result.append(data[chunkDataStart..<chunkDataEnd])
+
+            // Skip past the trailing \r\n after chunk data
+            offset = chunkDataEnd + 2
+        }
+
+        return .failure(.incomplete)
+    }
+
+    /// Find \r\n in data starting from given offset
+    private static func findCRLF(in data: Data, from start: Data.Index) -> Data.Index? {
+        var i = start
+        while i < data.endIndex - 1 {
+            if data[i] == 0x0D, data[i + 1] == 0x0A {
+                return i
+            }
+            i += 1
+        }
+        return nil
+    }
+
+    // MARK: - Response Building
+
     static func buildResponse(
         status: Int,
         statusText: String,
@@ -161,6 +235,7 @@ enum MCPHTTPParser {
         case 400: return "Bad Request"
         case 404: return "Not Found"
         case 405: return "Method Not Allowed"
+        case 413: return "Content Too Large"
         case 500: return "Internal Server Error"
         default: return "Unknown"
         }