fix: final review — stop() race, cancellation tracking, xlsx removal, stale callback guard

Author: datlechin

TablePro/Core/MCP/MCPAuthGuard.swift

@@ -138,13 +138,15 @@ actor MCPAuthGuard {
             }
             group.addTask {
                 try await Task.sleep(for: .seconds(30))
+                approvalTask.cancel()
                 throw MCPError.timeout(
                     String(localized: "User approval timed out after 30 seconds")
                 )
             }
             guard let result = try await group.next() else {
                 throw MCPError.internalError("No result from approval prompt")
             }
+            approvalTask.cancel()
             group.cancelAll()
             return result
         }

TablePro/Core/MCP/MCPConnectionBridge.swift

@@ -389,20 +389,19 @@ actor MCPConnectionBridge {
             }
         }
 
+        let (driver, _) = try await resolveDriver(connectionId)
+
         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 limitedTables {
             let columns = try await DatabaseManager.shared.trackOperation(sessionId: connectionId) {

TablePro/Core/MCP/MCPRouter.swift

@@ -268,16 +268,30 @@ final class MCPRouter: Sendable {
             return encodeError(MCPError.internalError("Server not fully initialized"), id: id)
         }
 
+        let session = await server.session(for: sessionId)
+        let toolTask = Task {
+            try await handler(name, arguments, sessionId)
+        }
+        if let session {
+            await session.addRunningTask(id, task: Task { _ = try? await toolTask.value })
+        }
+
         do {
-            let toolResult = try await handler(name, arguments, sessionId)
+            let toolResult = try await toolTask.value
+            if let session { _ = await session.removeRunningTask(id) }
             let resultData = try encoder.encode(toolResult)
             guard let resultValue = try? decoder.decode(JSONValue.self, from: resultData) else {
                 return encodeError(MCPError.internalError("Failed to encode tool result"), id: id)
             }
             return encodeRawResult(resultValue, id: id, sessionId: sessionId)
+        } catch is CancellationError {
+            if let session { _ = await session.removeRunningTask(id) }
+            return encodeError(MCPError.timeout("Request was cancelled"), id: id)
         } catch let mcpError as MCPError {
+            if let session { _ = await session.removeRunningTask(id) }
             return encodeError(mcpError, id: id)
         } catch {
+            if let session { _ = await session.removeRunningTask(id) }
             return encodeError(MCPError.internalError(error.localizedDescription), id: id)
         }
     }
@@ -651,8 +665,8 @@ extension MCPRouter {
                         ]),
                         "format": .object([
                             "type": "string",
-                            "description": "Export format: csv, json, sql, or xlsx",
-                            "enum": .array([.string("csv"), .string("json"), .string("sql"), .string("xlsx")])
+                            "description": "Export format: csv, json, or sql",
+                            "enum": .array([.string("csv"), .string("json"), .string("sql")])
                         ]),
                         "query": .object([
                             "type": "string",

TablePro/Core/MCP/MCPServer.swift

@@ -106,18 +106,15 @@ actor MCPServer {
 
         if let currentListener = listener {
             listener = nil
-            currentListener.cancel()
-            // Wait for the listener to fully release the port
             await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
                 currentListener.stateUpdateHandler = { state in
                     if case .cancelled = state {
                         continuation.resume()
                     }
                 }
+                currentListener.cancel()
             }
         }
-
-        stateCallback(.stopped)
     }
 
     var sessionCount: Int {

TablePro/Core/MCP/MCPServerManager.swift

@@ -24,6 +24,7 @@ final class MCPServerManager {
     private(set) var connectedClients: [MCPServer.SessionSnapshot] = []
     private var server: MCPServer?
     private var clientRefreshTask: Task<Void, Never>?
+    private var serverGeneration: Int = 0
 
     var isRunning: Bool {
         if case .running = state { return true } else { return false }
@@ -43,9 +44,12 @@ final class MCPServerManager {
             await stop()
         }
 
+        serverGeneration += 1
+        let generation = serverGeneration
         let newServer = MCPServer { [weak self] newState in
             Task { @MainActor in
-                self?.state = newState
+                guard let self, self.serverGeneration == generation else { return }
+                self.state = newState
             }
         }
 

TablePro/Core/MCP/MCPToolHandler.swift

@@ -221,13 +221,8 @@ final class MCPToolHandler: Sendable {
         let outputPath = optionalString(args, key: "output_path")
         let maxRows = optionalInt(args, key: "max_rows", default: 50_000, clamp: 1...100_000)
 
-        guard ["csv", "json", "sql", "xlsx"].contains(format) else {
-            throw MCPError.invalidParams("Unsupported format: \(format). Must be csv, json, sql, or xlsx")
-        }
-
-        // XLSX is binary and cannot be returned inline
-        if format == "xlsx" && outputPath == nil {
-            throw MCPError.invalidParams("XLSX export requires output_path since it is a binary format")
+        guard ["csv", "json", "sql"].contains(format) else {
+            throw MCPError.invalidParams("Unsupported format: \(format). Must be csv, json, or sql")
         }
 
         guard query != nil || tables != nil else {
@@ -280,9 +275,6 @@ final class MCPToolHandler: Sendable {
                 formatted = formatJSON(columns: columnNames, rows: rows)
             case "sql":
                 formatted = formatSQL(table: label, columns: columnNames, rows: rows)
-            case "xlsx":
-                // XLSX is handled via file write below; generate CSV as intermediate
-                formatted = formatCSV(columns: columnNames, rows: rows)
             default:
                 formatted = formatCSV(columns: columnNames, rows: rows)
             }