fix(mcp): let SQL Server scripts past the execution gate and keep GO lines out of single statements

CHANGELOG.md

@@ -509,6 +509,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - SQL Server scripts refused, or cut to their first result set, over MCP, AppleScript and the AI assistant.
 - SQL import into SQL Server sending `GO` lines to the server and splitting batches at each semicolon.
 - SQL Server dumps failing on their first view, routine or trigger when restored with sqlcmd or imported.
+- A leading `GO` line sent to SQL Server by MCP and AI assistant tools, and a `GO n` count ignored.
 
 ### Security
 

Packages/TableProCore/Sources/TableProSQLGrammar/SQLStatementScanner.swift

@@ -186,14 +186,16 @@ public enum SQLStatementScanner {
         return results
     }
 
-    /// `text` as a driver receives it when it is sent whole: trimmed, and ending where its last statement's executable
-    /// form ends, so a trailing separator comes off and a terminator that belongs to the statement stays. Empty when
-    /// nothing but separators, comments or blanks is left.
+    /// `text` as a driver receives it when it is sent whole: starting where its first statement starts and ending where
+    /// its last statement's executable form ends, so a separator on either side comes off and a terminator that belongs
+    /// to the statement stays. A `GO` line ahead of the first statement is the client's word: SQL Server reads one it
+    /// receives as a call to a procedure named `GO`. Empty when nothing but separators, comments or blanks is left.
     public static func executableText(of text: String, grammar: SQLLexicalGrammar) -> String {
         let trimmed = StatementBlank.trimming(text)
-        guard let last = executableStatements(in: trimmed, grammar: grammar).last else { return "" }
-        let end = last.range.location + last.range.length
-        return StatementBlank.trimming((trimmed as NSString).substring(to: end))
+        let statements = executableStatements(in: trimmed, grammar: grammar)
+        guard let first = statements.first, let last = statements.last else { return "" }
+        let span = NSRange(location: first.range.location, length: last.range.upperBound - first.range.location)
+        return (trimmed as NSString).substring(with: span)
     }
 
     /// The text the driver receives for `located`, or nil when nothing but a separator is left.

Packages/TableProCore/Tests/TableProSQLGrammarTests/SQLBatchSeparatorTests.swift

@@ -26,6 +26,10 @@ struct SQLBatchSeparatorTests {
         SQLStatementScanner.executableStatements(in: text, grammar: grammar).map(\.sql)
     }
 
+    private func executableText(_ text: String) -> String {
+        SQLStatementScanner.executableText(of: text, grammar: Self.sqlServer)
+    }
+
     private func separatorText(_ text: String) -> [String] {
         separators(text).map { (text as NSString).substring(with: $0.range) }
     }
@@ -162,6 +166,16 @@ struct SQLBatchSeparatorTests {
         #expect(statements(text) == ["SELECT 1", "SELECT 2"])
     }
 
+    /// Sent whole, a leading `GO` answers Msg 2812, "Could not find stored procedure 'GO'", and SQL Server still runs
+    /// the statement after it, so a tool reported a failure for a `DROP` that had run (measured on Azure SQL Edge 15).
+    @Test("A GO line before the first statement stays out of the text sent whole")
+    func leadingSeparatorIsNotSent() {
+        #expect(executableText("GO\nDROP TABLE dbo.stale") == "DROP TABLE dbo.stale")
+        #expect(executableText("  go -- lead\n\nDROP TABLE dbo.stale;") == "DROP TABLE dbo.stale")
+        #expect(executableText("GO\nGO 3\n-- keep\nSELECT 1\nGO") == "-- keep\nSELECT 1")
+        #expect(executableText("SELECT 1\nGO\nSELECT 2") == "SELECT 1\nGO\nSELECT 2")
+    }
+
     @Test("The reporter's script has no GO, so it stays five statements in one batch")
     func reporterScriptIsUnchanged() {
         let text = """

TablePro/Core/AI/Chat/Tools/ExecuteQueryChatTool.swift

@@ -53,8 +53,12 @@ struct ExecuteQueryChatTool: ChatTool {
         )
 
         let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
+        let capabilities = ExternalStatementGate.capabilities(
+            context.writeCapabilities,
+            takingScriptsOn: meta.databaseType
+        )
 
-        guard ExternalStatementGate.acceptsScripts(on: meta.databaseType)
+        guard capabilities.contains(.mayRunMultiStatement)
             || !QueryClassifier.isMultiStatement(query, databaseType: meta.databaseType)
         else {
             return ChatToolResult(
@@ -113,7 +117,7 @@ struct ExecuteQueryChatTool: ChatTool {
             sql: query,
             connectionId: connectionId,
             databaseType: meta.databaseType,
-            capabilities: context.writeCapabilities
+            capabilities: capabilities
         )
 
         let services = MCPToolServices(connectionBridge: context.bridge, authPolicy: context.authPolicy)

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

@@ -50,6 +50,9 @@ enum MCPStatementGate {
         if allowsDestructive {
             capabilities.insert(.mayRunDestructive)
         }
+        if allowsMultiStatement {
+            capabilities.insert(.mayRunMultiStatement)
+        }
         capabilities.formUnion(consent.capabilities)
 
         try await services.authPolicy.checkSafeModeDialog(

TablePro/Core/Scripting/ScriptQueryRunner.swift

@@ -34,6 +34,13 @@ internal enum ScriptQueryRunner {
         internal let client: String?
     }
 
+    /// What a script may ask of the execution gate: to write, and to drop once a person confirms it. It is never
+    /// pre-cleared, since no external caller gets to say on its own that a person already agreed. On an engine that
+    /// takes scripts it may also run several statements in one call.
+    internal static func capabilities(on databaseType: DatabaseType) -> CallerCapabilities {
+        ExternalStatementGate.capabilities([.mayWrite, .mayRunDestructive], takingScriptsOn: databaseType)
+    }
+
     internal static func run(
         _ request: Request,
         bridge: DatabaseAccessBridge,
@@ -50,6 +57,8 @@ internal enum ScriptQueryRunner {
             return try ExternalConnectionPolicySnapshot.resolve(connectionId: request.connectionId)
         }
 
+        let capabilities = Self.capabilities(on: snapshot.databaseType)
+
         /// Classified before anything connects, so a statement the connection refuses never opens a
         /// session and never asks the user for a password.
         try ExternalStatementGate.classify(
@@ -60,7 +69,7 @@ internal enum ScriptQueryRunner {
                 externalAccess: snapshot.externalAccess,
                 loadsExtensions: snapshot.loadsExtensions,
                 allowsDestructive: true,
-                allowsMultiStatement: ExternalStatementGate.acceptsScripts(on: snapshot.databaseType)
+                allowsMultiStatement: capabilities.contains(.mayRunMultiStatement)
             )
         )
 
@@ -80,7 +89,7 @@ internal enum ScriptQueryRunner {
             connectionId: request.connectionId,
             databaseType: snapshot.databaseType,
             caller: .appleScript(client: request.client),
-            capabilities: [.mayWrite, .mayRunDestructive],
+            capabilities: capabilities,
             operationDescription: confirmationTitle(client: request.client, connection: snapshot.connectionName)
         )
 

TablePro/Core/Services/Execution/ExternalStatementGate.swift

@@ -119,6 +119,18 @@ internal enum ExternalStatementGate {
         databaseType.lexicalGrammar.contains(.batchSeparatorLines)
     }
 
+    /// `capabilities`, plus leave to run several statements in one call where the engine takes scripts.
+    ///
+    /// The execution gate refuses a text of several statements to a caller without that leave, before Safe Mode is
+    /// asked anything, so a caller that lets a script past `classify` has to claim it there as well.
+    internal static func capabilities(
+        _ capabilities: CallerCapabilities,
+        takingScriptsOn databaseType: DatabaseType
+    ) -> CallerCapabilities {
+        guard acceptsScripts(on: databaseType) else { return capabilities }
+        return capabilities.union(.mayRunMultiStatement)
+    }
+
     /// A loaded SQLite extension can add functions that write files or run a nested statement, and
     /// the classifier reads `SELECT BlobToFile(...)` as a read. So on a connection that loads
     /// extensions, a statement from outside the app may call only what SQLite itself provides. Nil

TablePro/Core/Utilities/SQL/QueryClassifier.swift

@@ -111,10 +111,13 @@ enum QueryClassifier {
         isMultiStatement(sql, databaseType: databaseType, readings: databaseType.lexicalReadings)
     }
 
+    /// Whether `sql` runs more than one statement: it holds several, or a `GO n` line runs its one statement n times.
+    /// A caller that sends one statement once cannot honour the count, and dropping it runs the statement once.
     static func isMultiStatement(_ sql: String, databaseType: DatabaseType, readings: SQLLexicalReadings) -> Bool {
         let model = QueryStatementModel.forDatabaseType(databaseType)
         return readings.distinct(for: sql).contains { grammar in
-            QueryStatementScanner.executableStatements(in: sql, model: model, grammar: grammar).count > 1
+            let batches = QueryBatchPlanner.batches(in: sql, model: model, grammar: grammar)
+            return batches.flatMap(\.statements).count > 1 || batches.contains { $0.repeatCount > 1 }
         }
     }
 

TableProTests/Core/Execution/ExternalStatementGateTests.swift

@@ -85,6 +85,27 @@ struct ExternalStatementGateTests {
         }
     }
 
+    /// `GO 50` runs the batch fifty times in sqlcmd and in the editor. A tool that sends one statement once would run it
+    /// once and report success, so the count is refused rather than dropped.
+    @Test("A statement a GO line repeats is refused unless the caller takes scripts")
+    func repeatedStatementIsRefused() {
+        let repeated = "DELETE TOP (1000) FROM dbo.log WHERE archived = 1\nGO 50"
+        #expect(refusal(statement(repeated, databaseType: .mssql))
+            == .invalidArgument(String(localized: "Send one statement at a time.")))
+        #expect(refusal(statement(repeated, databaseType: .mssql, allowsMultiStatement: true)) == nil)
+        #expect(refusal(statement("DELETE FROM dbo.log WHERE archived = 1\nGO", databaseType: .mssql)) == nil)
+        #expect(refusal(statement("GO 3\nDELETE FROM dbo.log WHERE archived = 1", databaseType: .mssql)) == nil)
+    }
+
+    @Test("A caller that takes scripts claims leave to run several statements only where the engine takes them")
+    func scriptCapabilityFollowsTheEngine() {
+        let base: CallerCapabilities = [.mayWrite, .mayRunDestructive]
+        #expect(ExternalStatementGate.capabilities(base, takingScriptsOn: .mssql) == base.union(.mayRunMultiStatement))
+        for engine: DatabaseType in [.postgresql, .mysql, .sqlite, .oracle] {
+            #expect(ExternalStatementGate.capabilities(base, takingScriptsOn: engine) == base, "\(engine.rawValue)")
+        }
+    }
+
     @Test("A SQL Server script clears the gate for a caller that takes scripts, and is still tiered by its worst statement")
     func sqlServerScriptIsGatedWhole() {
         let script = "DECLARE @sn NVARCHAR(50) = 'x';\nSELECT * FROM a WHERE sn = @sn;\nGO\nSELECT * FROM b"

TableProTests/Core/MCP/MCPScriptResultTests.swift

@@ -126,4 +126,91 @@ struct MCPScriptResultTests {
         #expect(payload["result_sets"]?.arrayValue?.count == 2)
         #expect(payload["columns"]?.arrayValue?.compactMap(\.stringValue) == ["a"])
     }
+
+    // MARK: - Through every gate
+
+    nonisolated private static let scripts = [
+        "SELECT 1 AS a;\nSELECT 2 AS b;",
+        "SELECT 1 AS a\nGO\nSELECT 2 AS b"
+    ]
+
+    /// A SQL Server session at Silent, whose driver answers each batch with one result set per `AS` alias it holds.
+    private func sqlServerSession() -> ScriptAnsweringDriver {
+        var connection = TestFixtures.makeConnection(database: "warehouse", type: .mssql)
+        connection.aiPolicy = .alwaysAllow
+        let driver = ScriptAnsweringDriver(connection: connection) { batch in
+            ScriptAnsweringDriver.batch(
+                ["a", "b"].filter { batch.contains("AS \($0)") }.map { alias in
+                    ScriptAnsweringDriver.resultSet(columns: [alias], rows: [["1"]])
+                }
+            )
+        }
+        var session = ConnectionSession(connection: connection)
+        session.driver = driver
+        DatabaseManager.shared.injectSession(session, for: connection.id)
+        return driver
+    }
+
+    private func authPolicy() -> MCPAuthPolicy {
+        MCPAuthPolicy(connectionResolver: { _ in nil }, connectionIdsProvider: { [] }, historyRecorder: DiscardingHistory())
+    }
+
+    @Test("execute_query runs a SQL Server script of several statements past both gates", arguments: scripts)
+    func executeQueryToolRunsTheScript(script: String) async throws {
+        let driver = sqlServerSession()
+        defer { DatabaseManager.shared.removeSession(for: driver.connection.id) }
+
+        let result = try await ExecuteQueryTool().perform(
+            arguments: .object([
+                "connection_id": .string(driver.connection.id.uuidString),
+                "query": .string(script)
+            ]),
+            context: MCPToolTestHarness.context(),
+            services: MCPToolServices(connectionBridge: MCPConnectionBridge(), authPolicy: authPolicy())
+        )
+
+        let payload = try #require(result.structuredContent)
+        #expect(!result.isError)
+        #expect(payload["result_sets"]?.arrayValue?.map { $0["columns"]?.arrayValue?.compactMap(\.stringValue) }
+            == [["a"], ["b"]])
+        #expect(!driver.sentBatches.isEmpty)
+    }
+
+    @Test("The assistant's execute_query runs a SQL Server script of several statements past both gates", arguments: scripts)
+    func chatToolRunsTheScript(script: String) async throws {
+        let driver = sqlServerSession()
+        defer { DatabaseManager.shared.removeSession(for: driver.connection.id) }
+        let context = ChatToolContext(
+            connectionId: driver.connection.id,
+            bridge: MCPConnectionBridge(),
+            authPolicy: authPolicy()
+        )
+
+        let result = try await ExecuteQueryChatTool().execute(input: .object(["query": .string(script)]), context: context)
+
+        #expect(!result.isError)
+        let payload = try JSONDecoder().decode(JsonValue.self, from: Data(result.content.utf8))
+        #expect(payload["result_sets"]?.arrayValue?.count == 2)
+        #expect(!driver.sentBatches.isEmpty)
+    }
+
+    /// Measured on Azure SQL Edge 15: `GO\nDROP TABLE dbo.stale` sent whole answers Msg 2812, "Could not find stored
+    /// procedure 'GO'", and still drops the table, so the tool reports a failure for a statement that ran.
+    @Test("A GO line ahead of the one statement a tool sends never reaches the driver")
+    func leadingSeparatorIsNotSent() async throws {
+        let driver = sqlServerSession()
+        defer { DatabaseManager.shared.removeSession(for: driver.connection.id) }
+
+        _ = try await ToolQueryExecutor.executeAndLog(
+            services: MCPToolServices(connectionBridge: MCPConnectionBridge(), authPolicy: authPolicy()),
+            query: "GO\nDROP TABLE dbo.stale",
+            scope: DatabaseScope(connectionId: driver.connection.id, database: "warehouse", schema: nil),
+            maxRows: 0,
+            timeoutSeconds: 30,
+            principal: MCPToolTestHarness.principal()
+        )
+
+        #expect(driver.sentStatements == ["DROP TABLE dbo.stale"])
+        #expect(driver.sentBatches.isEmpty)
+    }
 }

TableProTests/Core/MCP/Protocol/Tools/MCPStatementGateTests.swift

@@ -99,6 +99,18 @@ struct MCPStatementGateRefusalTests {
         #expect(error?.code == .invalidArgument)
     }
 
+    /// The execution gate counts statements too, and it counts them before Safe Mode is asked anything, so leave that
+    /// only the first gate was given refused every script at Silent (#3078).
+    @Test("A SQL Server script clears the execution gate as well when the caller takes scripts", arguments: [
+        "DECLARE @sn NVARCHAR(50) = N'x';\nSELECT 1 AS a WHERE @sn = N'x';\nSELECT 2 AS b;",
+        "SELECT 1 AS a\nGO\nSELECT 2 AS b",
+        "SELECT 1 AS a\nGO 3"
+    ])
+    func sqlServerScriptClearsBothGates(script: String) async throws {
+        let error = try await refusal(sql: script, databaseType: .mssql, allowsMultiStatement: true)
+        #expect(error == nil)
+    }
+
     @Test("A destructive statement is refused unless the caller allows destructive work")
     func destructiveIsRefusedWithoutOptIn() async throws {
         let error = try await refusal(sql: "DROP TABLE users")

TableProTests/Core/Scripting/ScriptingPolicyTests.swift

@@ -52,7 +52,7 @@ struct ScriptingPolicyTests {
             connectionId: connectionId,
             databaseType: .postgresql,
             caller: .appleScript(client: "Script Editor"),
-            capabilities: [.mayWrite, .mayRunDestructive],
+            capabilities: ScriptQueryRunner.capabilities(on: .postgresql),
             operationDescription: "Script Editor wants to run a query on \"Production\"",
             gate: gate
         )
@@ -67,24 +67,53 @@ struct ScriptingPolicyTests {
         #expect(!request.capabilities.contains(.cannotPrompt))
     }
 
-    /// A script runs one statement, so `mayRunMultiStatement` is deliberately absent. `classify`
-    /// refuses several statements before this point; the capability is the second line.
-    @Test("A script never carries permission to run several statements at once")
-    func scriptsMayNotRunMultipleStatements() async throws {
-        let gate = RecordingExecutionGate(decision: authorized())
+    /// A script runs one statement, except on SQL Server, where the script is the unit (#3078). `classify` refuses
+    /// several statements everywhere else before this point; the capability is the second line.
+    @Test("A script carries permission to run several statements at once only where the engine takes scripts")
+    func scriptsRunSeveralStatementsOnlyWhereScriptsAreTheUnit() {
+        let sqlServer = ScriptQueryRunner.capabilities(on: .mssql)
+        #expect(sqlServer == [.mayWrite, .mayRunDestructive, .mayRunMultiStatement])
+        for engine: DatabaseType in [.postgresql, .mysql, .sqlite, .oracle] {
+            #expect(!ScriptQueryRunner.capabilities(on: engine).contains(.mayRunMultiStatement), "\(engine.rawValue)")
+        }
+    }
 
-        try await ExternalStatementGate.authorizeExecution(
-            sql: "SELECT 1",
-            connectionId: UUID(),
-            databaseType: .postgresql,
-            caller: .appleScript(client: nil),
-            capabilities: [.mayWrite, .mayRunDestructive],
-            operationDescription: "a query",
-            gate: gate
+    /// The execution gate counts statements before it asks Safe Mode anything, so this holds at Silent too.
+    @Test("A SQL Server script from a script clears the execution gate, and several statements elsewhere do not")
+    @MainActor
+    func sqlServerScriptClearsTheExecutionGate() async throws {
+        let gate = DefaultExecutionGate(
+            confirming: StubConfirming(answer: false),
+            authenticating: StubAuthenticating(answer: false),
+            safeModeLevelResolver: { _ in .silent },
+            forcesWriteResolver: { _ in false }
         )
 
-        let request = try #require(await gate.lastRequest)
-        #expect(!request.capabilities.contains(.mayRunMultiStatement))
+        for script in ["SELECT 1 AS a;\nSELECT 2 AS b;", "SELECT 1 AS a\nGO\nSELECT 2 AS b", "SELECT 1 AS a\nGO 3"] {
+            try await ExternalStatementGate.authorizeExecution(
+                sql: script,
+                connectionId: UUID(),
+                databaseType: .mssql,
+                caller: .appleScript(client: nil),
+                capabilities: ScriptQueryRunner.capabilities(on: .mssql),
+                operationDescription: "a query",
+                gate: gate
+            )
+        }
+
+        await #expect(throws: ExternalStatementGateError.denied(
+            String(localized: "Multiple statements are not permitted for this client")
+        )) {
+            try await ExternalStatementGate.authorizeExecution(
+                sql: "SELECT 1; SELECT 2",
+                connectionId: UUID(),
+                databaseType: .postgresql,
+                caller: .appleScript(client: nil),
+                capabilities: ScriptQueryRunner.capabilities(on: .postgresql),
+                operationDescription: "a query",
+                gate: gate
+            )
+        }
     }
 
     @Test("A denied statement throws the gate's own reason, so the script can read it")