From 9e61895716dc0b2d67c67c6f7789e22f40fc4663 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nguye=CC=82=CC=83n=20Nam=20Long?=
 <26674308+J2TeamNNL@users.noreply.github.com>
Date: Wed, 21 Oct 2026 15:43:41 +0700
Subject: [PATCH 01/10] feat(ai): give the schema prompt to the tables the
 question is about

---
 .../SchemaRelevanceRanker.swift               | 187 ++++++++++++++++++
 .../SchemaRelevanceRankerTests.swift          |  62 ++++++
 apple/TablePro/Core/AI/AISchemaContext.swift  |  39 +++-
 .../AIChatViewModel+SchemaContext.swift       |  33 +++-
 .../AIChatViewModel+SlashCommands.swift       |   2 +-
 .../AIChatViewModel+Streaming.swift           |   5 +-
 .../TablePro/ViewModels/AIChatViewModel.swift |   2 +
 .../AI/AISchemaContextRelevanceTests.swift    |  82 ++++++++
 8 files changed, 402 insertions(+), 10 deletions(-)
 create mode 100644 apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift
 create mode 100644 apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift
 create mode 100644 apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift

diff --git a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift
new file mode 100644
index 0000000000..76f763f489
--- /dev/null
+++ b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift
@@ -0,0 +1,187 @@
+import Foundation
+
+public struct SchemaRelevanceDocument: Sendable, Equatable {
+    public let name: String
+    public let details: [String]
+
+    public init(name: String, details: [String] = []) {
+        self.name = name
+        self.details = details
+    }
+}
+
+// Orders tables by how well their names and columns match a question: BM25 over identifier
+// words, with trigram matching for words that miss every exact term (typos, inflections).
+public enum SchemaRelevanceRanker {
+    static let nameWeight = 3
+    static let fullNameBoost = 5.0
+    static let fuzzyThreshold = 0.5
+    static let fuzzyMinLength = 4
+    static let k1 = 1.2
+    static let b = 0.75
+
+    public static func rank(_ documents: [SchemaRelevanceDocument], for query: String) -> [Int] {
+        let queryWords = Set(words(in: query))
+        guard !queryWords.isEmpty, !documents.isEmpty else { return [] }
+
+        let index = Index(documents: documents)
+        let matches = queryWords.flatMap { index.matches(for: $0) }
+        let compactQuery = words(in: query).joined()
+
+        var scored: [(offset: Int, score: Double)] = []
+        for (offset, terms) in index.termCounts.enumerated() {
+            var score = 0.0
+            for match in matches {
+                guard let frequency = terms[match.term] else { continue }
+                score += match.weight * index.bm25(frequency: frequency, idfTerm: match.term, length: index.lengths[offset])
+            }
+            let compactName = index.nameWords[offset].joined()
+            if compactName.count >= fuzzyMinLength, compactQuery.contains(compactName) {
+                score += fullNameBoost
+            }
+            if score > 0 {
+                scored.append((offset, score))
+            }
+        }
+
+        return scored
+            .sorted { $0.score != $1.score ? $0.score > $1.score : $0.offset < $1.offset }
+            .map(\.offset)
+    }
+
+    static func words(in text: String) -> [String] {
+        var result: [String] = []
+        var current = ""
+        var previous: Character?
+
+        func flush() {
+            if !current.isEmpty {
+                result.append(stem(current))
+                current = ""
+            }
+        }
+
+        for character in text {
+            guard character.isLetter || character.isNumber else {
+                flush()
+                previous = nil
+                continue
+            }
+            if let previous, startsNewWord(previous: previous, current: character) {
+                flush()
+            }
+            current += fold(character)
+            previous = character
+        }
+        flush()
+        return result
+    }
+
+    // Folding leaves "đ" alone because Unicode treats it as its own letter, not a "d" with a mark.
+    private static func fold(_ character: Character) -> String {
+        if character == "đ" || character == "Đ" { return "d" }
+        return String(character).folding(options: [.diacriticInsensitive, .caseInsensitive, .widthInsensitive], locale: nil)
+    }
+
+    private static func startsNewWord(previous: Character, current: Character) -> Bool {
+        if previous.isLowercase, current.isUppercase { return true }
+        if previous.isLetter, current.isNumber { return true }
+        if previous.isNumber, current.isLetter { return true }
+        return false
+    }
+
+    static func stem(_ word: String) -> String {
+        guard word.count > 3, word.allSatisfy(\.isASCII) else { return word }
+        if word.hasSuffix("ies") { return String(word.dropLast(3)) + "y" }
+        if word.hasSuffix("sses") || word.hasSuffix("xes") || word.hasSuffix("ches") || word.hasSuffix("shes") {
+            return String(word.dropLast(2))
+        }
+        if word.hasSuffix("s"), !word.hasSuffix("ss"), !word.hasSuffix("us"), !word.hasSuffix("is") {
+            return String(word.dropLast())
+        }
+        return word
+    }
+
+    static func trigramSimilarity(_ lhs: String, _ rhs: String) -> Double {
+        let left = trigrams(lhs)
+        let right = trigrams(rhs)
+        guard !left.isEmpty, !right.isEmpty else { return 0 }
+        return Double(left.intersection(right).count) / Double(left.union(right).count)
+    }
+
+    private static func trigrams(_ word: String) -> Set<String> {
+        let padded = Array("  \(word) ")
+        guard padded.count >= 3 else { return [] }
+        return Set((0...(padded.count - 3)).map { String(padded[$0..<($0 + 3)]) })
+    }
+
+    private struct Match {
+        let term: String
+        let weight: Double
+    }
+
+    private struct Index {
+        let termCounts: [[String: Int]]
+        let nameWords: [[String]]
+        let lengths: [Int]
+        let documentFrequency: [String: Int]
+        let averageLength: Double
+
+        init(documents: [SchemaRelevanceDocument]) {
+            var termCounts: [[String: Int]] = []
+            var nameWords: [[String]] = []
+            var lengths: [Int] = []
+            var documentFrequency: [String: Int] = [:]
+            termCounts.reserveCapacity(documents.count)
+
+            for document in documents {
+                let name = SchemaRelevanceRanker.words(in: document.name)
+                var counts: [String: Int] = [:]
+                for word in name {
+                    counts[word, default: 0] += SchemaRelevanceRanker.nameWeight
+                }
+                for detail in document.details {
+                    for word in SchemaRelevanceRanker.words(in: detail) {
+                        counts[word, default: 0] += 1
+                    }
+                }
+                for term in counts.keys {
+                    documentFrequency[term, default: 0] += 1
+                }
+                termCounts.append(counts)
+                nameWords.append(name)
+                lengths.append(counts.values.reduce(0, +))
+            }
+
+            self.termCounts = termCounts
+            self.nameWords = nameWords
+            self.lengths = lengths
+            self.documentFrequency = documentFrequency
+            let total = lengths.reduce(0, +)
+            self.averageLength = documents.isEmpty ? 0 : max(Double(total) / Double(documents.count), 1)
+        }
+
+        func matches(for word: String) -> [Match] {
+            if documentFrequency[word] != nil {
+                return [Match(term: word, weight: 1)]
+            }
+            guard word.count >= SchemaRelevanceRanker.fuzzyMinLength else { return [] }
+            return documentFrequency.keys.compactMap { term in
+                guard term.count >= SchemaRelevanceRanker.fuzzyMinLength else { return nil }
+                let similarity = SchemaRelevanceRanker.trigramSimilarity(word, term)
+                return similarity >= SchemaRelevanceRanker.fuzzyThreshold ? Match(term: term, weight: similarity) : nil
+            }
+        }
+
+        func bm25(frequency: Int, idfTerm: String, length: Int) -> Double {
+            let documentCount = Double(termCounts.count)
+            let containing = Double(documentFrequency[idfTerm] ?? 0)
+            let idf = log(1 + (documentCount - containing + 0.5) / (containing + 0.5))
+            let tf = Double(frequency)
+            let k1 = SchemaRelevanceRanker.k1
+            let b = SchemaRelevanceRanker.b
+            let norm = k1 * (1 - b + b * Double(length) / averageLength)
+            return idf * (tf * (k1 + 1)) / (tf + norm)
+        }
+    }
+}
diff --git a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift
new file mode 100644
index 0000000000..5a76a36a0a
--- /dev/null
+++ b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift
@@ -0,0 +1,62 @@
+import Foundation
+@testable import TableProAutocomplete
+import Testing
+
+struct SchemaRelevanceRankerTests {
+    private let documents = [
+        SchemaRelevanceDocument(name: "audit_log", details: ["id", "created_at", "message"]),
+        SchemaRelevanceDocument(name: "customers", details: ["id", "email", "full_name"]),
+        SchemaRelevanceDocument(name: "order_items", details: ["order_id", "product_id", "quantity"]),
+        SchemaRelevanceDocument(name: "orders", details: ["id", "customer_id", "total_amount", "created_at"]),
+        SchemaRelevanceDocument(name: "products", details: ["id", "title", "price"]),
+        SchemaRelevanceDocument(name: "ProductCategory", details: ["id", "label"])
+    ]
+
+    private func names(_ query: String) -> [String] {
+        SchemaRelevanceRanker.rank(documents, for: query).map { documents[$0].name }
+    }
+
+    @Test("Splits snake case, camel case and digits into folded, stemmed words")
+    func splitsIdentifiers() {
+        #expect(SchemaRelevanceRanker.words(in: "order_items") == ["order", "item"])
+        #expect(SchemaRelevanceRanker.words(in: "ProductCategories") == ["product", "category"])
+        #expect(SchemaRelevanceRanker.words(in: "userId2fa") == ["user", "id", "2", "fa"])
+        #expect(SchemaRelevanceRanker.words(in: "Đơn hàng") == ["don", "hang"])
+    }
+
+    @Test("A table named in the question ranks first")
+    func namedTableRanksFirst() {
+        #expect(names("total amount of orders per customer").first == "orders")
+    }
+
+    @Test("A word that only a column holds still finds its table")
+    func columnWordFindsTable() {
+        #expect(names("which rows have an email").first == "customers")
+    }
+
+    @Test("A misspelled table name matches by trigrams")
+    func typoMatches() {
+        #expect(names("list all custmers").first == "customers")
+    }
+
+    @Test("A camel case table matches its plural words")
+    func camelCaseMatches() {
+        #expect(names("product categories").contains("ProductCategory"))
+    }
+
+    @Test("A question that matches nothing ranks nothing")
+    func noMatchIsEmpty() {
+        #expect(names("hello there").isEmpty)
+        #expect(SchemaRelevanceRanker.rank(documents, for: "").isEmpty)
+        #expect(SchemaRelevanceRanker.rank([], for: "orders").isEmpty)
+    }
+
+    @Test("Equal scores keep the listed order")
+    func tiesKeepListedOrder() {
+        let twins = [
+            SchemaRelevanceDocument(name: "b_events"),
+            SchemaRelevanceDocument(name: "a_events")
+        ]
+        #expect(SchemaRelevanceRanker.rank(twins, for: "events") == [0, 1])
+    }
+}
diff --git a/apple/TablePro/Core/AI/AISchemaContext.swift b/apple/TablePro/Core/AI/AISchemaContext.swift
index b4090c2249..19dadfe5e8 100644
--- a/apple/TablePro/Core/AI/AISchemaContext.swift
+++ b/apple/TablePro/Core/AI/AISchemaContext.swift
@@ -6,6 +6,7 @@
 //
 
 import Foundation
+import TableProAutocomplete
 import TableProPluginKit
 
 /// One table as the schema context describes it. The producer joins each table to its own columns
@@ -156,12 +157,42 @@ struct AISchemaContext {
             }
         }
 
-        if tables.count > maxTables {
-            lines.append(
-                "\n… and \(tables.count - maxTables) more tables (not shown)"
-            )
+        let hiddenTables = tables.dropFirst(maxTables)
+        if !hiddenTables.isEmpty {
+            let listed = hiddenTables.prefix(maxListedHiddenTables).map { entry in
+                SchemaQualifiedName.render(
+                    name: entry.table.name,
+                    schema: entry.table.schema,
+                    implicitSchemaName: defaultSchema,
+                    quote: { "\(q)\($0)\(q)" }
+                )
+            }
+            var line = "\nOther tables (columns not shown): " + listed.joined(separator: ", ")
+            if hiddenTables.count > listed.count {
+                line += ", … and \(hiddenTables.count - listed.count) more"
+            }
+            lines.append(line)
         }
 
         return lines.joined(separator: "\n")
     }
+
+    static let maxListedHiddenTables = 200
+
+    /// Names of the tables that match `question`, best first, so the prompt spends its column budget
+    /// on them. Empty when every table already fits, or when nothing matches.
+    static func relevantTableNames(
+        in tables: [TableInfo],
+        columnsByTable: [String: [ColumnInfo]],
+        question: String,
+        maxTables: Int
+    ) -> [String] {
+        guard tables.count > maxTables else { return [] }
+        let documents = tables.map { table in
+            var details = (columnsByTable[table.name] ?? []).flatMap { [$0.name, $0.comment ?? ""] }
+            if let comment = table.comment { details.append(comment) }
+            return SchemaRelevanceDocument(name: table.name, details: details)
+        }
+        return SchemaRelevanceRanker.rank(documents, for: question).map { tables[$0].name }
+    }
 }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index cddca848f7..97963fe9e3 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -60,10 +60,15 @@ extension AIChatViewModel {
         await task.value
     }
 
-    func ensureSchemaLoaded() async {
+    /// With a `question`, the tables that match it get their columns loaded and listed first.
+    func ensureSchemaLoaded(relevantTo question: String? = nil) async {
+        if let question {
+            prioritizeSchema(for: question)
+        }
         if let inFlight = inFlightSchemaLoad {
             await inFlight.value
-            return
+            // The load in flight fetched for the previous order; a new order may need other tables.
+            guard question != nil else { return }
         }
         let task: Task<Void, Never> = Task { [weak self] in
             guard let self else { return }
@@ -98,7 +103,7 @@ extension AIChatViewModel {
         guard let connection else { return }
         let connId = connection.id
         let settings = services.appSettings.ai
-        let tablesToFetch = Array(tables.prefix(settings.maxSchemaTables))
+        let tablesToFetch = Array(orderedTables().prefix(settings.maxSchemaTables))
         guard !tablesToFetch.isEmpty else { return }
 
         await withTaskGroup(of: (String, [ColumnInfo]).self) { group in
@@ -193,7 +198,7 @@ extension AIChatViewModel {
     /// `tables` lists the one schema the session browses, where a table name is unique, so the
     /// name-keyed column and foreign key maps join to it by name.
     private func schemaTables() -> [AISchemaTable] {
-        tables.map { table in
+        orderedTables().map { table in
             AISchemaTable(
                 table: table,
                 columns: columnsByTable[table.name] ?? [],
@@ -202,6 +207,26 @@ extension AIChatViewModel {
         }
     }
 
+    private func prioritizeSchema(for question: String) {
+        let text = [question, currentQuery ?? ""].joined(separator: "\n")
+        guard !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
+        relevantTableNames = AISchemaContext.relevantTableNames(
+            in: tables,
+            columnsByTable: columnsByTable,
+            question: text,
+            maxTables: services.appSettings.ai.maxSchemaTables
+        )
+    }
+
+    private func orderedTables() -> [TableInfo] {
+        let all = tables
+        guard !relevantTableNames.isEmpty else { return all }
+        let byName = Dictionary(all.map { ($0.name, $0) }, uniquingKeysWith: { first, _ in first })
+        let relevant = relevantTableNames.compactMap { byName[$0] }
+        let picked = Set(relevant.map(\.name))
+        return relevant + all.filter { !picked.contains($0.name) }
+    }
+
     private func defaultSchema(for connection: DatabaseConnection) -> String? {
         connection.type.implicitSchemaName
             ?? services.databaseManager.session(for: connection.id)?.browseSchema
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift b/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
index 36eb178746..0ba44b3db2 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
@@ -39,7 +39,7 @@ extension AIChatViewModel {
         let invocationText = body.isEmpty ? "/\(command.name)" : "/\(command.name) \(body)"
         let needsSchema = command.promptTemplate.contains(CustomSlashCommandVariable.schema.placeholder)
         if needsSchema {
-            await ensureSchemaLoaded()
+            await ensureSchemaLoaded(relevantTo: body)
         }
         let renderingContext = CustomSlashCommandRenderer.Context(
             query: currentQuery,
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
index a091a133cf..e38eb948b9 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
@@ -125,7 +125,10 @@ extension AIChatViewModel {
         prepTask = Task.detached(priority: .userInitiated) { [weak self] in
             guard let self else { return }
             if settings.includeSchema {
-                await self.ensureSchemaLoaded()
+                let question = await MainActor.run {
+                    self.messages.last { $0.role == .user }?.plainText
+                }
+                await self.ensureSchemaLoaded(relevantTo: question)
             }
             if Task.isCancelled { return }
             let priorTurns: [ChatTurn] = await MainActor.run {
diff --git a/apple/TablePro/ViewModels/AIChatViewModel.swift b/apple/TablePro/ViewModels/AIChatViewModel.swift
index 2a118460af..d312b393c2 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel.swift
@@ -104,6 +104,7 @@ final class AIChatViewModel: ObservableObject {
     var queryContextMetadata: any ScopedMetadataProviding
     var inFlightColumnFetches: [String: Task<Void, Never>] = [:]
     var inFlightSchemaLoad: Task<Void, Never>?
+    var relevantTableNames: [String] = []
     nonisolated(unsafe) var streamingTask: Task<Void, Never>?
     var prepTask: Task<Void, Never>?
 
@@ -354,6 +355,7 @@ final class AIChatViewModel: ObservableObject {
         inFlightColumnFetches.removeAll()
         inFlightSchemaLoad?.cancel()
         inFlightSchemaLoad = nil
+        relevantTableNames = []
         currentQuery = nil
         queryResults = nil
         editorTarget = nil
diff --git a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
new file mode 100644
index 0000000000..f9358138f1
--- /dev/null
+++ b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
@@ -0,0 +1,82 @@
+//
+//  AISchemaContextRelevanceTests.swift
+//  TableProTests
+//
+//  A schema with more tables than the prompt's column budget spends that budget on the tables
+//  the question names, and still tells the model what the other tables are called.
+//
+
+import Foundation
+@testable import TablePro
+import TableProPluginKit
+import Testing
+
+struct AISchemaContextRelevanceTests {
+    private static let tables: [TableInfo] = [
+        TestFixtures.makeTableInfo(name: "audit_log"),
+        TestFixtures.makeTableInfo(name: "customers"),
+        TestFixtures.makeTableInfo(name: "invoices"),
+        TableInfo(name: "tbl_x", type: .table, rowCount: 0, comment: "Monthly revenue snapshots"),
+        TestFixtures.makeTableInfo(name: "orders")
+    ]
+
+    @Test("Tables that match the question come first, best first")
+    func matchingTablesComeFirst() {
+        let names = AISchemaContext.relevantTableNames(
+            in: Self.tables,
+            columnsByTable: [:],
+            question: "How many orders did each customer place?",
+            maxTables: 2
+        )
+        #expect(Array(names.prefix(2)).sorted() == ["customers", "orders"])
+    }
+
+    @Test("A table comment and loaded columns count as matches")
+    func commentAndColumnsMatch() {
+        let columns = ["invoices": [TestFixtures.makeColumnInfo(name: "due_date")]]
+        #expect(AISchemaContext.relevantTableNames(
+            in: Self.tables, columnsByTable: columns, question: "revenue by month", maxTables: 2
+        ).first == "tbl_x")
+        #expect(AISchemaContext.relevantTableNames(
+            in: Self.tables, columnsByTable: columns, question: "unpaid past the due date", maxTables: 2
+        ).first == "invoices")
+    }
+
+    @Test("Nothing is reordered when every table fits")
+    func everyTableFits() {
+        let names = AISchemaContext.relevantTableNames(
+            in: Self.tables,
+            columnsByTable: [:],
+            question: "orders",
+            maxTables: Self.tables.count
+        )
+        #expect(names.isEmpty)
+    }
+
+    @Test("Tables past the budget are named without their columns")
+    func hiddenTablesAreNamed() {
+        let entries = Self.tables.map { AISchemaTable(table: $0) }
+        let section = AISchemaContext.buildSchemaSection(
+            tables: entries,
+            defaultSchema: nil,
+            maxTables: 2,
+            identifierQuote: "\""
+        )
+        #expect(section.contains("Other tables (columns not shown): \"invoices\", \"tbl_x\", \"orders\""))
+        #expect(!section.contains("more tables (not shown)"))
+    }
+
+    @Test("A long tail of hidden tables is cut with a count")
+    func longTailIsCounted() {
+        let many = (0..<(AISchemaContext.maxListedHiddenTables + 3)).map {
+            AISchemaTable(table: TestFixtures.makeTableInfo(name: "t\($0)"))
+        }
+        let section = AISchemaContext.buildSchemaSection(
+            tables: many,
+            defaultSchema: nil,
+            maxTables: 1,
+            identifierQuote: "\""
+        )
+        #expect(section.hasSuffix(", … and 2 more"))
+    }
+}

From ca6296fe82d3dba86803b850534eb15df62c1ea4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nguye=CC=82=CC=83n=20Nam=20Long?=
 <26674308+J2TeamNNL@users.noreply.github.com>
Date: Wed, 21 Oct 2026 15:44:06 +0700
Subject: [PATCH 02/10] docs(changelog): note schema relevance in AI chat

---
 apple/changelog.d/3750.changed.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 apple/changelog.d/3750.changed.md

diff --git a/apple/changelog.d/3750.changed.md b/apple/changelog.d/3750.changed.md
new file mode 100644
index 0000000000..e5b535f26f
--- /dev/null
+++ b/apple/changelog.d/3750.changed.md
@@ -0,0 +1 @@
+AI chat sends the columns of the tables your question is about when a database has more tables than the schema limit.

From 8cf9ddfc39c6cf48a104b0f86a4550fe5f63a6bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nguye=CC=82=CC=83n=20Nam=20Long?=
 <26674308+J2TeamNNL@users.noreply.github.com>
Date: Wed, 21 Oct 2026 16:29:02 +0700
Subject: [PATCH 03/10] fix(ai): keep the last ranking, match table names by
 whole words, and clear the schema load only when it is ours

A follow-up that matches no table no longer resets the ranking. The full-name boost needs the name's words in order in the question. Term trigrams are built once per index. The in-flight schema load is cleared only by the task that set it, and the question is read when the turn starts.
---
 .../SchemaRelevanceRanker.swift               | 31 +++++++++++++------
 .../SchemaRelevanceRankerTests.swift          |  9 ++++++
 .../AIChatViewModel+SchemaContext.swift       |  9 ++++--
 .../AIChatViewModel+Streaming.swift           |  4 +--
 4 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift
index 76f763f489..5e19f843e4 100644
--- a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift
+++ b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift
@@ -26,7 +26,7 @@ public enum SchemaRelevanceRanker {
 
         let index = Index(documents: documents)
         let matches = queryWords.flatMap { index.matches(for: $0) }
-        let compactQuery = words(in: query).joined()
+        let orderedQueryWords = words(in: query)
 
         var scored: [(offset: Int, score: Double)] = []
         for (offset, terms) in index.termCounts.enumerated() {
@@ -35,8 +35,8 @@ public enum SchemaRelevanceRanker {
                 guard let frequency = terms[match.term] else { continue }
                 score += match.weight * index.bm25(frequency: frequency, idfTerm: match.term, length: index.lengths[offset])
             }
-            let compactName = index.nameWords[offset].joined()
-            if compactName.count >= fuzzyMinLength, compactQuery.contains(compactName) {
+            let nameWords = index.nameWords[offset]
+            if nameWords.joined().count >= fuzzyMinLength, containsSequence(orderedQueryWords, nameWords) {
                 score += fullNameBoost
             }
             if score > 0 {
@@ -49,6 +49,13 @@ public enum SchemaRelevanceRanker {
             .map(\.offset)
     }
 
+    private static func containsSequence(_ words: [String], _ sequence: [String]) -> Bool {
+        guard !sequence.isEmpty, words.count >= sequence.count else { return false }
+        return (0...(words.count - sequence.count)).contains { start in
+            Array(words[start..<(start + sequence.count)]) == sequence
+        }
+    }
+
     static func words(in text: String) -> [String] {
         var result: [String] = []
         var current = ""
@@ -103,8 +110,10 @@ public enum SchemaRelevanceRanker {
     }
 
     static func trigramSimilarity(_ lhs: String, _ rhs: String) -> Double {
-        let left = trigrams(lhs)
-        let right = trigrams(rhs)
+        trigramSimilarity(trigrams(lhs), trigrams(rhs))
+    }
+
+    private static func trigramSimilarity(_ left: Set<String>, _ right: Set<String>) -> Double {
         guard !left.isEmpty, !right.isEmpty else { return 0 }
         return Double(left.intersection(right).count) / Double(left.union(right).count)
     }
@@ -125,6 +134,7 @@ public enum SchemaRelevanceRanker {
         let nameWords: [[String]]
         let lengths: [Int]
         let documentFrequency: [String: Int]
+        let fuzzyTerms: [(term: String, trigrams: Set<String>)]
         let averageLength: Double
 
         init(documents: [SchemaRelevanceDocument]) {
@@ -157,6 +167,9 @@ public enum SchemaRelevanceRanker {
             self.nameWords = nameWords
             self.lengths = lengths
             self.documentFrequency = documentFrequency
+            self.fuzzyTerms = documentFrequency.keys
+                .filter { $0.count >= SchemaRelevanceRanker.fuzzyMinLength }
+                .map { ($0, SchemaRelevanceRanker.trigrams($0)) }
             let total = lengths.reduce(0, +)
             self.averageLength = documents.isEmpty ? 0 : max(Double(total) / Double(documents.count), 1)
         }
@@ -166,10 +179,10 @@ public enum SchemaRelevanceRanker {
                 return [Match(term: word, weight: 1)]
             }
             guard word.count >= SchemaRelevanceRanker.fuzzyMinLength else { return [] }
-            return documentFrequency.keys.compactMap { term in
-                guard term.count >= SchemaRelevanceRanker.fuzzyMinLength else { return nil }
-                let similarity = SchemaRelevanceRanker.trigramSimilarity(word, term)
-                return similarity >= SchemaRelevanceRanker.fuzzyThreshold ? Match(term: term, weight: similarity) : nil
+            let wordTrigrams = SchemaRelevanceRanker.trigrams(word)
+            return fuzzyTerms.compactMap { entry in
+                let similarity = SchemaRelevanceRanker.trigramSimilarity(wordTrigrams, entry.trigrams)
+                return similarity >= SchemaRelevanceRanker.fuzzyThreshold ? Match(term: entry.term, weight: similarity) : nil
             }
         }
 
diff --git a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift
index 5a76a36a0a..39bf8f6b65 100644
--- a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift
+++ b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift
@@ -51,6 +51,15 @@ struct SchemaRelevanceRankerTests {
         #expect(SchemaRelevanceRanker.rank([], for: "orders").isEmpty)
     }
 
+    @Test("A table name inside a longer word gets no name boost")
+    func nameInsideLongerWordIsNotBoosted() {
+        let documents = [
+            SchemaRelevanceDocument(name: "data"),
+            SchemaRelevanceDocument(name: "events", details: ["metadata"])
+        ]
+        #expect(SchemaRelevanceRanker.rank(documents, for: "show the metadata").first == 1)
+    }
+
     @Test("Equal scores keep the listed order")
     func tiesKeepListedOrder() {
         let twins = [
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index 97963fe9e3..dab2972267 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -76,7 +76,9 @@ extension AIChatViewModel {
         }
         inFlightSchemaLoad = task
         await task.value
-        inFlightSchemaLoad = nil
+        if inFlightSchemaLoad == task {
+            inFlightSchemaLoad = nil
+        }
     }
 
     func ensureSavedQueryLoaded(id: UUID) async {
@@ -210,12 +212,15 @@ extension AIChatViewModel {
     private func prioritizeSchema(for question: String) {
         let text = [question, currentQuery ?? ""].joined(separator: "\n")
         guard !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
-        relevantTableNames = AISchemaContext.relevantTableNames(
+        let ranked = AISchemaContext.relevantTableNames(
             in: tables,
             columnsByTable: columnsByTable,
             question: text,
             maxTables: services.appSettings.ai.maxSchemaTables
         )
+        if !ranked.isEmpty {
+            relevantTableNames = ranked
+        }
     }
 
     private func orderedTables() -> [TableInfo] {
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
index e38eb948b9..354c8e4f1a 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
@@ -122,12 +122,10 @@ extension AIChatViewModel {
         let assistantID = assistantMessage.id
         streamingState = .streaming(assistantID: assistantID)
 
+        let question = messages.last { $0.role == .user }?.plainText
         prepTask = Task.detached(priority: .userInitiated) { [weak self] in
             guard let self else { return }
             if settings.includeSchema {
-                let question = await MainActor.run {
-                    self.messages.last { $0.role == .user }?.plainText
-                }
                 await self.ensureSchemaLoaded(relevantTo: question)
             }
             if Task.isCancelled { return }

From 16cf83812cf2f7cab726a22dc1ee83d5f69aaa0d Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 17:39:32 +0700
Subject: [PATCH 04/10] fix(ai): rank by the editor's table references, read
 columns uniformly, and keep hidden tables a count

---
 apple/TablePro/Core/AI/AISchemaContext.swift  |  74 +++++++----
 .../InlineSuggestion/AIChatInlineSource.swift |  10 +-
 .../AppleIntelligenceInlineSource.swift       |   6 +-
 .../Core/Autocomplete/SQLSchemaProvider.swift |  38 ++++--
 .../Core/Database/DatabaseDriver.swift        |   5 +
 .../Core/Plugins/PluginDriverAdapter.swift    |   2 +
 .../AIChatViewModel+SchemaContext.swift       |  62 +++++++--
 .../TablePro/ViewModels/AIChatViewModel.swift |   2 +
 .../AI/AISchemaContextRelevanceTests.swift    | 122 ++++++++++++------
 .../Core/AI/SchemaContextForAITests.swift     |  34 +++++
 apple/changelog.d/3750.changed.md             |   2 +-
 11 files changed, 270 insertions(+), 87 deletions(-)

diff --git a/apple/TablePro/Core/AI/AISchemaContext.swift b/apple/TablePro/Core/AI/AISchemaContext.swift
index 19dadfe5e8..37b1add061 100644
--- a/apple/TablePro/Core/AI/AISchemaContext.swift
+++ b/apple/TablePro/Core/AI/AISchemaContext.swift
@@ -8,6 +8,7 @@
 import Foundation
 import TableProAutocomplete
 import TableProPluginKit
+import TableProSQLGrammar
 
 /// One table as the schema context describes it. The producer joins each table to its own columns
 /// and foreign keys, because a lookup by name cannot tell two schemas' tables of the same name apart.
@@ -23,11 +24,11 @@ struct AISchemaTable: Sendable {
     }
 }
 
-/// Builds schema context for AI system prompts
 struct AISchemaContext {
+    static let maxQueryLength = 2_000
+
     // MARK: - Public
 
-    /// Build a system prompt including database context
     static func buildSystemPrompt(
         databaseType: DatabaseType,
         databaseName: String,
@@ -67,7 +68,6 @@ struct AISchemaContext {
            let query = currentQuery,
            !query.isEmpty {
             let lang = editorLanguage.codeBlockTag
-            let maxQueryLength = 2_000
             let nsQuery = query as NSString
             let truncated = nsQuery.length > maxQueryLength
                 ? nsQuery.substring(to: maxQueryLength) + "\n-- ... truncated"
@@ -157,42 +157,68 @@ struct AISchemaContext {
             }
         }
 
-        let hiddenTables = tables.dropFirst(maxTables)
-        if !hiddenTables.isEmpty {
-            let listed = hiddenTables.prefix(maxListedHiddenTables).map { entry in
-                SchemaQualifiedName.render(
-                    name: entry.table.name,
-                    schema: entry.table.schema,
-                    implicitSchemaName: defaultSchema,
-                    quote: { "\(q)\($0)\(q)" }
-                )
-            }
-            var line = "\nOther tables (columns not shown): " + listed.joined(separator: ", ")
-            if hiddenTables.count > listed.count {
-                line += ", … and \(hiddenTables.count - listed.count) more"
-            }
-            lines.append(line)
+        if tables.count > maxTables {
+            lines.append(
+                "\n… and \(tables.count - maxTables) more tables (not shown)"
+            )
         }
 
         return lines.joined(separator: "\n")
     }
 
-    static let maxListedHiddenTables = 200
-
-    /// Names of the tables that match `question`, best first, so the prompt spends its column budget
-    /// on them. Empty when every table already fits, or when nothing matches.
+    // The editor's tables come first, since the model is reading that query; the question's matches
+    // follow, best first.
     static func relevantTableNames(
         in tables: [TableInfo],
         columnsByTable: [String: [ColumnInfo]],
         question: String,
+        queryReferences: [TableReference] = [],
         maxTables: Int
     ) -> [String] {
         guard tables.count > maxTables else { return [] }
+        var names = indices(of: queryReferences, in: tables).map { tables[$0].name }
+        // Columns count only once every table has them, or a table loaded on an earlier turn would
+        // outrank a better match whose columns were never read.
+        let columns = tables.allSatisfy { columnsByTable[$0.name] != nil } ? columnsByTable : [:]
         let documents = tables.map { table in
-            var details = (columnsByTable[table.name] ?? []).flatMap { [$0.name, $0.comment ?? ""] }
+            var details = (columns[table.name] ?? []).flatMap { [$0.name, $0.comment ?? ""] }
             if let comment = table.comment { details.append(comment) }
             return SchemaRelevanceDocument(name: table.name, details: details)
         }
-        return SchemaRelevanceRanker.rank(documents, for: question).map { tables[$0].name }
+        var seen = Set(names)
+        for offset in SchemaRelevanceRanker.rank(documents, for: question) where seen.insert(tables[offset].name).inserted {
+            names.append(tables[offset].name)
+        }
+        return names
+    }
+
+    // Only the part of the query the prompt carries, and never a CTE or subquery, which has no
+    // table of its own.
+    static func referencedTables(in query: String, grammar: SQLLexicalGrammar) -> [TableReference] {
+        let nsQuery = query as NSString
+        let text = nsQuery.length > maxQueryLength ? nsQuery.substring(to: maxQueryLength) : query
+        let analyzer = SQLContextAnalyzer()
+        return SQLStatementScanner.allStatements(in: text, grammar: grammar).flatMap { statement in
+            analyzer.analyze(query: statement, cursorPosition: (statement as NSString).length, grammar: grammar)
+                .tableReferences
+                .filter { $0.derivedColumns == nil }
+        }
+    }
+
+    // In the order the query names them, each table once.
+    static func indices(of references: [TableReference], in tables: [TableInfo]) -> [Int] {
+        var result: [Int] = []
+        for reference in references {
+            let match = tables.indices.first { offset in
+                let table = tables[offset]
+                guard table.name.caseInsensitiveCompare(reference.tableName) == .orderedSame else { return false }
+                guard let schema = reference.schema, let tableSchema = table.schema else { return true }
+                return schema.caseInsensitiveCompare(tableSchema) == .orderedSame
+            }
+            if let match, !result.contains(match) {
+                result.append(match)
+            }
+        }
+        return result
     }
 }
diff --git a/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift b/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift
index 0edde1a65f..266516fc4c 100644
--- a/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift
+++ b/apple/TablePro/Core/AI/InlineSuggestion/AIChatInlineSource.swift
@@ -62,7 +62,7 @@ final class AIChatInlineSource: InlineSuggestionSource {
             ChatTurnWire(role: .user, blocks: [.text(userMessage)])
         ]
 
-        let systemPrompt = await buildSystemPrompt(settings: settings)
+        let systemPrompt = await buildSystemPrompt(settings: settings, context: context)
 
         guard accessGate.allowsUnpromptedAccess(to: connectionId) else { return nil }
 
@@ -94,13 +94,17 @@ final class AIChatInlineSource: InlineSuggestionSource {
 
     // MARK: - Private
 
-    private func buildSystemPrompt(settings: AISettings) async -> String {
+    private func buildSystemPrompt(settings: AISettings, context: SuggestionContext) async -> String {
         guard settings.includeSchema,
               let provider = schemaProvider else {
             return AIPromptTemplates.inlineSuggestSystemPrompt()
         }
 
-        let schemaContext = await provider.buildSchemaContextForAI(settings: settings)
+        let schemaContext = await provider.buildSchemaContextForAI(
+            settings: settings,
+            editorText: context.fullText,
+            cursorOffset: context.cursorOffset
+        )
 
         if let schemaContext, !schemaContext.isEmpty {
             return AIPromptTemplates.inlineSuggestSystemPrompt(schemaContext: schemaContext)
diff --git a/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift b/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift
index 945c5d98f3..7b61d44b44 100644
--- a/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift
+++ b/apple/TablePro/Core/AI/InlineSuggestion/AppleIntelligenceInlineSource.swift
@@ -45,7 +45,11 @@ final class AppleIntelligenceInlineSource: InlineSuggestionSource {
 
         var candidates: [OnDeviceRequest] = []
         if settings.includeSchema, let provider = schemaProvider,
-           let schema = await provider.buildSchemaContextForAI(settings: settings), !schema.isEmpty {
+           let schema = await provider.buildSchemaContextForAI(
+               settings: settings,
+               editorText: context.fullText,
+               cursorOffset: context.cursorOffset
+           ), !schema.isEmpty {
             candidates.append(OnDeviceRequest(
                 instructions: AIPromptTemplates.inlineSuggestSystemPrompt(schemaContext: schema),
                 prompt: prompt
diff --git a/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift b/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift
index 5b63da0a12..fbd4acad44 100644
--- a/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift
+++ b/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift
@@ -9,6 +9,7 @@ import Foundation
 import os
 import TableProAutocomplete
 import TableProPluginKit
+import TableProSQLGrammar
 
 actor SQLSchemaProvider {
     private static let logger = Logger(subsystem: "com.TablePro", category: "SQLSchemaProvider")
@@ -430,18 +431,10 @@ actor SQLSchemaProvider {
 
     // MARK: - AI Schema Context
 
-    func buildSchemaContextForAI(settings: AISettings) async -> String? {
+    // The tables the statement at the cursor reads take the column budget first.
+    func buildSchemaContextForAI(settings: AISettings, editorText: String? = nil, cursorOffset: Int = 0) async -> String? {
         guard !tables.isEmpty, let connection = connectionInfo else { return nil }
 
-        let listedTables = tables
-        var schemaTables: [AISchemaTable] = []
-        schemaTables.reserveCapacity(listedTables.count)
-        for table in listedTables.prefix(settings.maxSchemaTables) {
-            let columns = await getColumns(for: table.name, schema: table.schema)
-            schemaTables.append(AISchemaTable(table: table, columns: columns))
-        }
-        schemaTables += listedTables.dropFirst(settings.maxSchemaTables).map { AISchemaTable(table: $0) }
-
         let dbType = connection.type
         let capturedConnection = connection
         let defaultSchema = getDefaultSchema()
@@ -455,6 +448,31 @@ actor SQLSchemaProvider {
             return (resolvedName, quote, lang, langName)
         }
 
+        var listedTables = tables
+        if let editorText, editorLanguage == .sql, listedTables.count > settings.maxSchemaTables {
+            let grammar = dbType.lexicalGrammar
+            let statement = SQLStatementScanner.statementAtCursor(
+                in: editorText,
+                cursorPosition: cursorOffset,
+                grammar: grammar
+            )
+            let first = AISchemaContext.indices(
+                of: AISchemaContext.referencedTables(in: statement, grammar: grammar),
+                in: listedTables
+            )
+            let picked = Set(first)
+            listedTables = first.map { listedTables[$0] }
+                + listedTables.indices.filter { !picked.contains($0) }.map { listedTables[$0] }
+        }
+
+        var schemaTables: [AISchemaTable] = []
+        schemaTables.reserveCapacity(listedTables.count)
+        for table in listedTables.prefix(settings.maxSchemaTables) {
+            let columns = await getColumns(for: table.name, schema: table.schema)
+            schemaTables.append(AISchemaTable(table: table, columns: columns))
+        }
+        schemaTables += listedTables.dropFirst(settings.maxSchemaTables).map { AISchemaTable(table: $0) }
+
         return AISchemaContext.buildSystemPrompt(
             databaseType: dbType,
             databaseName: dbName,
diff --git a/apple/TablePro/Core/Database/DatabaseDriver.swift b/apple/TablePro/Core/Database/DatabaseDriver.swift
index 247b95c9a5..ca902e7dc8 100644
--- a/apple/TablePro/Core/Database/DatabaseDriver.swift
+++ b/apple/TablePro/Core/Database/DatabaseDriver.swift
@@ -152,6 +152,9 @@ protocol DatabaseDriver: AnyObject, Sendable {
     /// Default implementation falls back to per-table fetchColumns.
     func fetchAllColumns() async throws -> [String: [ColumnInfo]]
 
+    // False when `fetchAllColumns` degrades to one round trip per table.
+    var providesBulkColumnFetch: Bool { get }
+
     /// Dotted field paths a document store exposes for a collection, for query authoring.
     /// Default implementation returns nothing, which is correct for every SQL driver.
     func sampleFieldPaths(table: String, limit: Int) async throws -> [PluginFieldPath]
@@ -748,6 +751,8 @@ extension DatabaseDriver {
         return results
     }
 
+    var providesBulkColumnFetch: Bool { false }
+
     var providesBulkForeignKeyFetch: Bool { false }
 
     func fetchAllForeignKeys(schema: String?) async throws -> [String: [ForeignKeyInfo]] {
diff --git a/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift b/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift
index 63b22e4b15..8be9661e5d 100644
--- a/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift
+++ b/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift
@@ -902,6 +902,8 @@ final class PluginDriverAdapter: DatabaseDriver, SchemaSwitchable, DatabaseRepor
         try await pluginDriver.sampleFieldPaths(table: table, schema: pluginDriver.currentSchema, limit: limit)
     }
 
+    var providesBulkColumnFetch: Bool { pluginDriver.providesBulkColumnFetch }
+
     func fetchAllColumns() async throws -> [String: [ColumnInfo]] {
         let pluginResult = try await pluginDriver.fetchAllColumns(schema: pluginDriver.currentSchema)
         return pluginResult.mapValues { $0.map(ColumnInfo.init) }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index dab2972267..603df03a2b 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -5,6 +5,7 @@
 
 import Foundation
 import os
+import TableProAutocomplete
 import TableProPluginKit
 
 extension AIChatViewModel {
@@ -60,11 +61,7 @@ extension AIChatViewModel {
         await task.value
     }
 
-    /// With a `question`, the tables that match it get their columns loaded and listed first.
     func ensureSchemaLoaded(relevantTo question: String? = nil) async {
-        if let question {
-            prioritizeSchema(for: question)
-        }
         if let inFlight = inFlightSchemaLoad {
             await inFlight.value
             // The load in flight fetched for the previous order; a new order may need other tables.
@@ -72,6 +69,11 @@ extension AIChatViewModel {
         }
         let task: Task<Void, Never> = Task { [weak self] in
             guard let self else { return }
+            if let question {
+                await self.loadEveryColumnForRanking()
+                guard !Task.isCancelled else { return }
+                self.prioritizeSchema(for: question)
+            }
             await self.runSchemaLoad()
         }
         inFlightSchemaLoad = task
@@ -123,7 +125,8 @@ extension AIChatViewModel {
                     }
                 }
             }
-            for await (name, cols) in group {
+            // A load cancelled by a session reset still finishes its reads; they belong to the old session.
+            for await (name, cols) in group where !Task.isCancelled {
                 columnsByTable[name] = cols
             }
         }
@@ -210,12 +213,11 @@ extension AIChatViewModel {
     }
 
     private func prioritizeSchema(for question: String) {
-        let text = [question, currentQuery ?? ""].joined(separator: "\n")
-        guard !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
         let ranked = AISchemaContext.relevantTableNames(
             in: tables,
             columnsByTable: columnsByTable,
-            question: text,
+            question: question,
+            queryReferences: editorQueryReferences(),
             maxTables: services.appSettings.ai.maxSchemaTables
         )
         if !ranked.isEmpty {
@@ -223,6 +225,50 @@ extension AIChatViewModel {
         }
     }
 
+    // The editor steers the order only while the model is shown the editor's query.
+    private func editorQueryReferences() -> [TableReference] {
+        guard services.appSettings.ai.includeCurrentQuery,
+              let connection,
+              let query = currentQuery, !query.isEmpty,
+              services.pluginManager.editorLanguage(for: connection.type) == .sql
+        else { return [] }
+        return AISchemaContext.referencedTables(in: query, grammar: connection.type.lexicalGrammar)
+    }
+
+    private func loadEveryColumnForRanking() async {
+        guard let connection else { return }
+        let tableCount = tables.count
+        guard tableCount > services.appSettings.ai.maxSchemaTables,
+              tableCount <= SQLSchemaProvider.maxCachedTables,
+              bulkColumnReadTableCount != tableCount
+        else { return }
+        bulkColumnReadTableCount = tableCount
+        let allColumns: [String: [ColumnInfo]]?
+        do {
+            allColumns = try await DatabaseManager.shared.withBrowseMetadataDriver(
+                connectionId: connection.id,
+                workload: .bulk
+            ) { driver in
+                // A driver without a bulk read would fall back to one query per table.
+                guard driver.providesBulkColumnFetch else { return nil }
+                return try await driver.fetchAllColumns()
+            }
+        } catch {
+            Self.logger.warning("Bulk column fetch failed: \(error.publicLogShape, privacy: .public)")
+            return
+        }
+        guard let allColumns, !Task.isCancelled else { return }
+        let byFoldedName = Dictionary(
+            allColumns.map { ($0.key.lowercased(), $0.value) },
+            uniquingKeysWith: { first, _ in first }
+        )
+        var merged = columnsByTable
+        for table in tables where (merged[table.name] ?? []).isEmpty {
+            merged[table.name] = allColumns[table.name] ?? byFoldedName[table.name.lowercased()] ?? []
+        }
+        columnsByTable = merged
+    }
+
     private func orderedTables() -> [TableInfo] {
         let all = tables
         guard !relevantTableNames.isEmpty else { return all }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel.swift b/apple/TablePro/ViewModels/AIChatViewModel.swift
index bf806c470a..758d51d714 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel.swift
@@ -108,6 +108,7 @@ final class AIChatViewModel: ObservableObject {
     var inFlightColumnFetches: [String: Task<Void, Never>] = [:]
     var inFlightSchemaLoad: Task<Void, Never>?
     var relevantTableNames: [String] = []
+    var bulkColumnReadTableCount: Int?
     nonisolated(unsafe) var streamingTask: Task<Void, Never>?
     var prepTask: Task<Void, Never>?
 
@@ -367,6 +368,7 @@ final class AIChatViewModel: ObservableObject {
         inFlightSchemaLoad?.cancel()
         inFlightSchemaLoad = nil
         relevantTableNames = []
+        bulkColumnReadTableCount = nil
         currentQuery = nil
         queryResults = nil
         editorTarget = nil
diff --git a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
index f9358138f1..6c8d752818 100644
--- a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
+++ b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
@@ -3,12 +3,13 @@
 //  TableProTests
 //
 //  A schema with more tables than the prompt's column budget spends that budget on the tables
-//  the question names, and still tells the model what the other tables are called.
+//  the editor's query reads and the question names.
 //
 
 import Foundation
 @testable import TablePro
 import TableProPluginKit
+import TableProSQLGrammar
 import Testing
 
 struct AISchemaContextRelevanceTests {
@@ -20,63 +21,104 @@ struct AISchemaContextRelevanceTests {
         TestFixtures.makeTableInfo(name: "orders")
     ]
 
+    private static let grammar = DatabaseType.postgresql.lexicalGrammar
+
+    private static func names(
+        question: String,
+        query: String = "",
+        columns: [String: [ColumnInfo]] = [:],
+        maxTables: Int = 2
+    ) -> [String] {
+        AISchemaContext.relevantTableNames(
+            in: tables,
+            columnsByTable: columns,
+            question: question,
+            queryReferences: AISchemaContext.referencedTables(in: query, grammar: grammar),
+            maxTables: maxTables
+        )
+    }
+
+    private static func everyTableColumns(_ extra: [String: [ColumnInfo]]) -> [String: [ColumnInfo]] {
+        var columns = Dictionary(uniqueKeysWithValues: tables.map { ($0.name, [TestFixtures.makeColumnInfo(name: "id")]) })
+        columns.merge(extra) { _, new in new }
+        return columns
+    }
+
     @Test("Tables that match the question come first, best first")
     func matchingTablesComeFirst() {
-        let names = AISchemaContext.relevantTableNames(
-            in: Self.tables,
-            columnsByTable: [:],
-            question: "How many orders did each customer place?",
-            maxTables: 2
-        )
+        let names = Self.names(question: "How many orders did each customer place?")
         #expect(Array(names.prefix(2)).sorted() == ["customers", "orders"])
     }
 
-    @Test("A table comment and loaded columns count as matches")
-    func commentAndColumnsMatch() {
-        let columns = ["invoices": [TestFixtures.makeColumnInfo(name: "due_date")]]
-        #expect(AISchemaContext.relevantTableNames(
-            in: Self.tables, columnsByTable: columns, question: "revenue by month", maxTables: 2
-        ).first == "tbl_x")
-        #expect(AISchemaContext.relevantTableNames(
-            in: Self.tables, columnsByTable: columns, question: "unpaid past the due date", maxTables: 2
-        ).first == "invoices")
+    @Test("A table comment counts as a match")
+    func commentMatches() {
+        #expect(Self.names(question: "revenue by month").first == "tbl_x")
+    }
+
+    @Test("Columns count once every table has them")
+    func columnsMatchWhenEveryTableHasThem() {
+        let columns = Self.everyTableColumns(["invoices": [TestFixtures.makeColumnInfo(name: "due_date")]])
+        #expect(Self.names(question: "unpaid past the due date", columns: columns).first == "invoices")
+    }
+
+    @Test("Columns loaded for some tables only do not count")
+    func partialColumnsAreIgnored() {
+        let columns = ["audit_log": [TestFixtures.makeColumnInfo(name: "customer_note")]]
+        #expect(Self.names(question: "customer", columns: columns).first == "customers")
+        #expect(!Self.names(question: "note", columns: columns).contains("audit_log"))
     }
 
     @Test("Nothing is reordered when every table fits")
     func everyTableFits() {
-        let names = AISchemaContext.relevantTableNames(
-            in: Self.tables,
-            columnsByTable: [:],
-            question: "orders",
-            maxTables: Self.tables.count
+        #expect(Self.names(question: "orders", maxTables: Self.tables.count).isEmpty)
+    }
+
+    @Test("Tables the editor's query reads come before the question's matches")
+    func queryTablesComeFirst() {
+        let names = Self.names(
+            question: "why is this slow for customers",
+            query: "SELECT * FROM invoices i JOIN audit_log a ON a.id = i.id"
         )
-        #expect(names.isEmpty)
+        #expect(Array(names.prefix(3)) == ["invoices", "audit_log", "customers"])
     }
 
-    @Test("Tables past the budget are named without their columns")
-    func hiddenTablesAreNamed() {
-        let entries = Self.tables.map { AISchemaTable(table: $0) }
-        let section = AISchemaContext.buildSchemaSection(
-            tables: entries,
-            defaultSchema: nil,
-            maxTables: 2,
-            identifierQuote: "\""
+    @Test("SQL keywords in the editor's query match no table")
+    func queryKeywordsMatchNothing() {
+        let names = Self.names(
+            question: "",
+            query: "SELECT id FROM customers GROUP BY id ORDER BY id"
         )
-        #expect(section.contains("Other tables (columns not shown): \"invoices\", \"tbl_x\", \"orders\""))
-        #expect(!section.contains("more tables (not shown)"))
+        #expect(names == ["customers"])
     }
 
-    @Test("A long tail of hidden tables is cut with a count")
-    func longTailIsCounted() {
-        let many = (0..<(AISchemaContext.maxListedHiddenTables + 3)).map {
-            AISchemaTable(table: TestFixtures.makeTableInfo(name: "t\($0)"))
-        }
+    @Test("Every statement in the editor counts, and a CTE is not a table")
+    func everyStatementCounts() {
+        let query = """
+        WITH recent AS (SELECT * FROM orders) SELECT * FROM recent;
+        UPDATE customers SET name = 'x';
+        """
+        #expect(Self.names(question: "", query: query) == ["orders", "customers"])
+    }
+
+    @Test("A schema-qualified reference matches only a table in that schema")
+    func schemaQualifiedReference() {
+        let tables = [
+            TestFixtures.makeTableInfo(name: "orders", schema: "archive"),
+            TestFixtures.makeTableInfo(name: "orders", schema: "public")
+        ]
+        let references = AISchemaContext.referencedTables(in: "SELECT * FROM public.orders", grammar: Self.grammar)
+        #expect(AISchemaContext.indices(of: references, in: tables) == [1])
+    }
+
+    @Test("Tables past the budget are counted, not named")
+    func hiddenTablesAreCounted() {
         let section = AISchemaContext.buildSchemaSection(
-            tables: many,
+            tables: Self.tables.map { AISchemaTable(table: $0) },
             defaultSchema: nil,
-            maxTables: 1,
+            maxTables: 2,
             identifierQuote: "\""
         )
-        #expect(section.hasSuffix(", … and 2 more"))
+        #expect(section.hasSuffix("… and 3 more tables (not shown)"))
+        #expect(!section.contains("\"invoices\""))
     }
 }
diff --git a/apple/TableProTests/Core/AI/SchemaContextForAITests.swift b/apple/TableProTests/Core/AI/SchemaContextForAITests.swift
index 4056ac3db7..c6679af295 100644
--- a/apple/TableProTests/Core/AI/SchemaContextForAITests.swift
+++ b/apple/TableProTests/Core/AI/SchemaContextForAITests.swift
@@ -162,6 +162,40 @@ struct SchemaContextForAITests {
         #expect(listed["hr.orders"] == ["salary"])
     }
 
+    @Test("the tables the statement at the cursor reads take the column budget")
+    func cursorStatementTablesTakeTheBudget() async throws {
+        let tables = [
+            TestFixtures.makeTableInfo(name: "audit_log"),
+            TestFixtures.makeTableInfo(name: "orders")
+        ]
+        let provider = Self.provider(
+            tables: tables,
+            columns: [
+                "audit_log": [TestFixtures.makeColumnInfo(name: "message")],
+                "orders": [TestFixtures.makeColumnInfo(name: "total")]
+            ]
+        )
+        await provider.resetForDatabase(
+            "shop",
+            tables: tables,
+            driver: MockDatabaseDriver(),
+            connection: TestFixtures.makeConnection()
+        )
+        var settings = AISettings.default
+        settings.maxSchemaTables = 1
+        let editorText = "SELECT 1;\nSELECT * FROM orders WHERE "
+
+        let context = try #require(await provider.buildSchemaContextForAI(
+            settings: settings,
+            editorText: editorText,
+            cursorOffset: (editorText as NSString).length
+        ))
+        let listed = Self.columnNamesByTableLine(in: context)
+
+        #expect(listed["orders"] == ["total"])
+        #expect(listed["audit_log"] == nil)
+    }
+
     private static func schemaAwareProvider(
         columnsBySchema: [String: [ColumnInfo]],
         bulkColumns: [String: [ColumnInfo]] = [:]
diff --git a/apple/changelog.d/3750.changed.md b/apple/changelog.d/3750.changed.md
index e5b535f26f..938adc0601 100644
--- a/apple/changelog.d/3750.changed.md
+++ b/apple/changelog.d/3750.changed.md
@@ -1 +1 @@
-AI chat sends the columns of the tables your question is about when a database has more tables than the schema limit.
+AI chat and inline suggestions send columns for the tables your question or query uses, not the first tables listed.

From afe04b19f183023e3088401fb24331951cdcdf9f Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 18:00:52 +0700
Subject: [PATCH 05/10] fix(ai): scope the chat column cache, gate slash
 command reads on the AI policy, and read table references past comments

---
 .../SQLContextAnalyzer.swift                  | 18 ++++++++
 ...QLContextAnalyzerTableReferenceTests.swift | 31 +++++++++++++
 apple/TablePro/Core/AI/AISchemaContext.swift  | 44 +++++++++---------
 .../Core/Autocomplete/SQLSchemaProvider.swift |  3 +-
 .../AIChatViewModel+SchemaContext.swift       | 45 ++++++++++++++++---
 .../AIChatViewModel+SlashCommands.swift       | 11 ++++-
 .../TablePro/ViewModels/AIChatViewModel.swift |  2 +
 .../AI/AISchemaContextRelevanceTests.swift    | 15 +++++++
 .../Core/AI/SchemaContextForAITests.swift     | 33 ++++++++++++++
 9 files changed, 174 insertions(+), 28 deletions(-)
 create mode 100644 apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift

diff --git a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
index 6874698878..3feb9a5ce2 100644
--- a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
+++ b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
@@ -288,6 +288,24 @@ public final class SQLContextAnalyzer {
 
     public init() {}
 
+    // `analyze` reads no tables when its cursor sits in a comment or literal; this reads them wherever the statement
+    // ends, so `SELECT * FROM orders -- note` still names `orders`.
+    public func tableReferences(in statement: String, grammar: SQLLexicalGrammar) -> [TableReference] {
+        let nsStatement = statement as NSString
+        let text = Self.lex(nsStatement, cursor: nsStatement.length, grammar: grammar)
+        var references = extractTableReferences(from: text.forReferences)
+        let tokens = Self.tokens(of: text, grammar: grammar)
+        mergeDerivedTables(
+            Self.derivedTableParser.visibleTables(in: tokens, cursor: nsStatement.length),
+            into: &references
+        )
+        if let alterTableName = extractAlterTableName(from: text.forReferences),
+           !references.contains(where: { $0.tableName == alterTableName }) {
+            references.append(TableReference(tableName: alterTableName, alias: nil))
+        }
+        return references
+    }
+
     // MARK: - Main Analysis
 
     public func analyze(query: String, cursorPosition: Int, grammar: SQLLexicalGrammar) -> SQLContext {
diff --git a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
new file mode 100644
index 0000000000..c4d8472f61
--- /dev/null
+++ b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
@@ -0,0 +1,31 @@
+import Foundation
+@testable import TableProAutocomplete
+import TableProSQLGrammar
+import Testing
+
+struct SQLContextAnalyzerTableReferenceTests {
+    private let analyzer = SQLContextAnalyzer()
+
+    private func tables(_ statement: String) -> [String] {
+        analyzer.tableReferences(in: statement, grammar: .ansi).map(\.tableName)
+    }
+
+    @Test("A statement that ends in a comment still names its tables")
+    func trailingComment() {
+        #expect(tables("SELECT * FROM orders -- recent orders") == ["orders"])
+        #expect(tables("SELECT * FROM orders /* recent */") == ["orders"])
+        #expect(analyzer.analyze(query: "SELECT * FROM orders -- recent", cursorPosition: 30).tableReferences.isEmpty)
+    }
+
+    @Test("A statement that ends in a string still names its tables")
+    func trailingLiteral() {
+        #expect(tables("SELECT * FROM orders WHERE note = 'open") == ["orders"])
+    }
+
+    @Test("Joined, updated and altered tables are named")
+    func everyClause() {
+        #expect(tables("SELECT * FROM orders o JOIN customers c ON c.id = o.customer_id") == ["orders", "customers"])
+        #expect(tables("UPDATE customers SET name = 'x'").contains("customers"))
+        #expect(tables("ALTER TABLE invoices ADD COLUMN due DATE") == ["invoices"])
+    }
+}
diff --git a/apple/TablePro/Core/AI/AISchemaContext.swift b/apple/TablePro/Core/AI/AISchemaContext.swift
index 37b1add061..e966c5b691 100644
--- a/apple/TablePro/Core/AI/AISchemaContext.swift
+++ b/apple/TablePro/Core/AI/AISchemaContext.swift
@@ -68,10 +68,8 @@ struct AISchemaContext {
            let query = currentQuery,
            !query.isEmpty {
             let lang = editorLanguage.codeBlockTag
-            let nsQuery = query as NSString
-            let truncated = nsQuery.length > maxQueryLength
-                ? nsQuery.substring(to: maxQueryLength) + "\n-- ... truncated"
-                : query
+            let carried = carriedQuery(query)
+            let truncated = carried.utf16.count < query.utf16.count ? carried + "\n-- ... truncated" : query
             parts.append("\n## Current Query\n```\(lang)\n\(truncated)\n```")
         }
 
@@ -173,10 +171,11 @@ struct AISchemaContext {
         columnsByTable: [String: [ColumnInfo]],
         question: String,
         queryReferences: [TableReference] = [],
+        defaultSchema: String? = nil,
         maxTables: Int
     ) -> [String] {
         guard tables.count > maxTables else { return [] }
-        var names = indices(of: queryReferences, in: tables).map { tables[$0].name }
+        var names = indices(of: queryReferences, in: tables, defaultSchema: defaultSchema).map { tables[$0].name }
         // Columns count only once every table has them, or a table loaded on an earlier turn would
         // outrank a better match whose columns were never read.
         let columns = tables.allSatisfy { columnsByTable[$0.name] != nil } ? columnsByTable : [:]
@@ -192,29 +191,34 @@ struct AISchemaContext {
         return names
     }
 
-    // Only the part of the query the prompt carries, and never a CTE or subquery, which has no
-    // table of its own.
-    static func referencedTables(in query: String, grammar: SQLLexicalGrammar) -> [TableReference] {
+    static func carriedQuery(_ query: String) -> String {
         let nsQuery = query as NSString
-        let text = nsQuery.length > maxQueryLength ? nsQuery.substring(to: maxQueryLength) : query
+        return nsQuery.length > maxQueryLength ? nsQuery.substring(to: maxQueryLength) : query
+    }
+
+    // A CTE or subquery has no table of its own, so it is left out.
+    static func referencedTables(in query: String, grammar: SQLLexicalGrammar) -> [TableReference] {
         let analyzer = SQLContextAnalyzer()
-        return SQLStatementScanner.allStatements(in: text, grammar: grammar).flatMap { statement in
-            analyzer.analyze(query: statement, cursorPosition: (statement as NSString).length, grammar: grammar)
-                .tableReferences
-                .filter { $0.derivedColumns == nil }
+        return SQLStatementScanner.allStatements(in: query, grammar: grammar).flatMap { statement in
+            analyzer.tableReferences(in: statement, grammar: grammar).filter { $0.derivedColumns == nil }
         }
     }
 
-    // In the order the query names them, each table once.
-    static func indices(of references: [TableReference], in tables: [TableInfo]) -> [Int] {
+    // In the order the query names them, each table once. A bare name means the default schema's
+    // table, as it does when the query runs.
+    static func indices(of references: [TableReference], in tables: [TableInfo], defaultSchema: String? = nil) -> [Int] {
         var result: [Int] = []
         for reference in references {
-            let match = tables.indices.first { offset in
-                let table = tables[offset]
-                guard table.name.caseInsensitiveCompare(reference.tableName) == .orderedSame else { return false }
-                guard let schema = reference.schema, let tableSchema = table.schema else { return true }
-                return schema.caseInsensitiveCompare(tableSchema) == .orderedSame
+            let named = tables.indices.filter {
+                tables[$0].name.caseInsensitiveCompare(reference.tableName) == .orderedSame
+            }
+            let schema = reference.schema ?? defaultSchema
+            let inSchema = named.first { offset in
+                guard let schema, let tableSchema = tables[offset].schema else { return false }
+                return tableSchema.caseInsensitiveCompare(schema) == .orderedSame
             }
+            let unscoped = named.first { tables[$0].schema == nil }
+            let match = inSchema ?? unscoped ?? (reference.schema == nil ? named.first : nil)
             if let match, !result.contains(match) {
                 result.append(match)
             }
diff --git a/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift b/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift
index fbd4acad44..6fb7ea2c7a 100644
--- a/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift
+++ b/apple/TablePro/Core/Autocomplete/SQLSchemaProvider.swift
@@ -458,7 +458,8 @@ actor SQLSchemaProvider {
             )
             let first = AISchemaContext.indices(
                 of: AISchemaContext.referencedTables(in: statement, grammar: grammar),
-                in: listedTables
+                in: listedTables,
+                defaultSchema: defaultSchema
             )
             let picked = Set(first)
             listedTables = first.map { listedTables[$0] }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index 603df03a2b..c69a20167a 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -24,6 +24,7 @@ extension AIChatViewModel {
     }
 
     func ensureColumnsLoaded(forTable tableName: String) async {
+        dropSchemaCacheIfScopeChanged()
         if let existing = columnsByTable[tableName], !existing.isEmpty { return }
         if let inFlight = inFlightColumnFetches[tableName] {
             await inFlight.value
@@ -31,6 +32,7 @@ extension AIChatViewModel {
         }
         guard let connection else { return }
         let connId = connection.id
+        let scope = schemaCacheScope
         let task: Task<Void, Never> = Task { [weak self] in
             let columns: [ColumnInfo]
             do {
@@ -50,7 +52,7 @@ extension AIChatViewModel {
                 Self.logger.warning("Foreign key fetch failed for \(tableName, privacy: .private(mask: .hash)): \(error.publicLogShape, privacy: .public)")
                 fkMap = [:]
             }
-            guard !Task.isCancelled, let self else { return }
+            guard !Task.isCancelled, let self, self.schemaCacheScope == scope else { return }
             self.columnsByTable[tableName] = columns
             if let fks = fkMap[tableName] {
                 self.foreignKeysByTable[tableName] = fks
@@ -62,6 +64,7 @@ extension AIChatViewModel {
     }
 
     func ensureSchemaLoaded(relevantTo question: String? = nil) async {
+        dropSchemaCacheIfScopeChanged()
         if let inFlight = inFlightSchemaLoad {
             await inFlight.value
             // The load in flight fetched for the previous order; a new order may need other tables.
@@ -106,6 +109,7 @@ extension AIChatViewModel {
     private func runSchemaLoad() async {
         guard let connection else { return }
         let connId = connection.id
+        let scope = schemaCacheScope
         let settings = services.appSettings.ai
         let tablesToFetch = Array(orderedTables().prefix(settings.maxSchemaTables))
         guard !tablesToFetch.isEmpty else { return }
@@ -125,13 +129,12 @@ extension AIChatViewModel {
                     }
                 }
             }
-            // A load cancelled by a session reset still finishes its reads; they belong to the old session.
-            for await (name, cols) in group where !Task.isCancelled {
+            for await (name, cols) in group where schemaCacheScope == scope {
                 columnsByTable[name] = cols
             }
         }
 
-        guard !Task.isCancelled else { return }
+        guard !Task.isCancelled, schemaCacheScope == scope else { return }
 
         let needsFKFetch = tablesToFetch.contains { foreignKeysByTable[$0.name] == nil }
         guard needsFKFetch else { return }
@@ -139,6 +142,7 @@ extension AIChatViewModel {
             let fkMap = try await DatabaseManager.shared.withBrowseMetadataDriver(connectionId: connId, workload: .bulk) { driver in
                 try await driver.fetchForeignKeys(forTables: tablesToFetch.map(\.name))
             }
+            guard schemaCacheScope == scope else { return }
             for (name, fks) in fkMap {
                 foreignKeysByTable[name] = fks
             }
@@ -218,6 +222,7 @@ extension AIChatViewModel {
             columnsByTable: columnsByTable,
             question: question,
             queryReferences: editorQueryReferences(),
+            defaultSchema: connection.flatMap { defaultSchema(for: $0) },
             maxTables: services.appSettings.ai.maxSchemaTables
         )
         if !ranked.isEmpty {
@@ -232,7 +237,10 @@ extension AIChatViewModel {
               let query = currentQuery, !query.isEmpty,
               services.pluginManager.editorLanguage(for: connection.type) == .sql
         else { return [] }
-        return AISchemaContext.referencedTables(in: query, grammar: connection.type.lexicalGrammar)
+        return AISchemaContext.referencedTables(
+            in: AISchemaContext.carriedQuery(query),
+            grammar: connection.type.lexicalGrammar
+        )
     }
 
     private func loadEveryColumnForRanking() async {
@@ -243,6 +251,7 @@ extension AIChatViewModel {
               bulkColumnReadTableCount != tableCount
         else { return }
         bulkColumnReadTableCount = tableCount
+        let scope = schemaCacheScope
         let allColumns: [String: [ColumnInfo]]?
         do {
             allColumns = try await DatabaseManager.shared.withBrowseMetadataDriver(
@@ -255,9 +264,12 @@ extension AIChatViewModel {
             }
         } catch {
             Self.logger.warning("Bulk column fetch failed: \(error.publicLogShape, privacy: .public)")
+            if schemaCacheScope == scope {
+                bulkColumnReadTableCount = nil
+            }
             return
         }
-        guard let allColumns, !Task.isCancelled else { return }
+        guard let allColumns, !Task.isCancelled, schemaCacheScope == scope else { return }
         let byFoldedName = Dictionary(
             allColumns.map { ($0.key.lowercased(), $0.value) },
             uniquingKeysWith: { first, _ in first }
@@ -269,6 +281,27 @@ extension AIChatViewModel {
         columnsByTable = merged
     }
 
+    // A session outlives a database or schema switch, so its column cache belongs to the scope it
+    // was read in.
+    private func dropSchemaCacheIfScopeChanged() {
+        let scope = connection.map { connection in
+            [
+                connection.id.uuidString,
+                services.databaseManager.browseDatabaseName(for: connection),
+                defaultSchema(for: connection) ?? ""
+            ].joined(separator: "\u{1F}")
+        }
+        guard scope != schemaCacheScope else { return }
+        schemaCacheScope = scope
+        columnsByTable = [:]
+        foreignKeysByTable = [:]
+        relevantTableNames = []
+        bulkColumnReadTableCount = nil
+        inFlightColumnFetches.removeAll()
+        inFlightSchemaLoad?.cancel()
+        inFlightSchemaLoad = nil
+    }
+
     private func orderedTables() -> [TableInfo] {
         let all = tables
         guard !relevantTableNames.isEmpty else { return all }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift b/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
index 0ba44b3db2..b74309ed2a 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
@@ -39,7 +39,16 @@ extension AIChatViewModel {
         let invocationText = body.isEmpty ? "/\(command.name)" : "/\(command.name) \(body)"
         let needsSchema = command.promptTemplate.contains(CustomSlashCommandVariable.schema.placeholder)
         if needsSchema {
-            await ensureSchemaLoaded(relevantTo: body)
+            // The policy check comes later, when the turn is sent, so a read the user has not
+            // allowed stays as small as it was: no ranking, no bulk column read.
+            switch connection == nil ? nil : resolveConnectionPolicy(settings: services.appSettings.ai) {
+            case .never?:
+                break
+            case .askEachTime?:
+                await ensureSchemaLoaded()
+            case .alwaysAllow?, nil:
+                await ensureSchemaLoaded(relevantTo: body)
+            }
         }
         let renderingContext = CustomSlashCommandRenderer.Context(
             query: currentQuery,
diff --git a/apple/TablePro/ViewModels/AIChatViewModel.swift b/apple/TablePro/ViewModels/AIChatViewModel.swift
index 758d51d714..24260bb851 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel.swift
@@ -109,6 +109,7 @@ final class AIChatViewModel: ObservableObject {
     var inFlightSchemaLoad: Task<Void, Never>?
     var relevantTableNames: [String] = []
     var bulkColumnReadTableCount: Int?
+    var schemaCacheScope: String?
     nonisolated(unsafe) var streamingTask: Task<Void, Never>?
     var prepTask: Task<Void, Never>?
 
@@ -369,6 +370,7 @@ final class AIChatViewModel: ObservableObject {
         inFlightSchemaLoad = nil
         relevantTableNames = []
         bulkColumnReadTableCount = nil
+        schemaCacheScope = nil
         currentQuery = nil
         queryResults = nil
         editorTarget = nil
diff --git a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
index 6c8d752818..9c96181dd6 100644
--- a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
+++ b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
@@ -110,6 +110,21 @@ struct AISchemaContextRelevanceTests {
         #expect(AISchemaContext.indices(of: references, in: tables) == [1])
     }
 
+    @Test("A bare name means the default schema's table")
+    func bareNameMeansDefaultSchema() {
+        let tables = [
+            TestFixtures.makeTableInfo(name: "orders", schema: "archive"),
+            TestFixtures.makeTableInfo(name: "orders", schema: "public")
+        ]
+        let references = AISchemaContext.referencedTables(in: "SELECT * FROM orders", grammar: Self.grammar)
+        #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [1])
+    }
+
+    @Test("A query that ends in a comment still puts its tables first")
+    func queryEndingInCommentCounts() {
+        #expect(Self.names(question: "", query: "SELECT * FROM invoices -- unpaid ones") == ["invoices"])
+    }
+
     @Test("Tables past the budget are counted, not named")
     func hiddenTablesAreCounted() {
         let section = AISchemaContext.buildSchemaSection(
diff --git a/apple/TableProTests/Core/AI/SchemaContextForAITests.swift b/apple/TableProTests/Core/AI/SchemaContextForAITests.swift
index c6679af295..838ccd5c89 100644
--- a/apple/TableProTests/Core/AI/SchemaContextForAITests.swift
+++ b/apple/TableProTests/Core/AI/SchemaContextForAITests.swift
@@ -196,6 +196,39 @@ struct SchemaContextForAITests {
         #expect(listed["audit_log"] == nil)
     }
 
+    @Test("a table named past the first 2,000 characters of a long statement still takes the budget")
+    func longStatementTableTakesTheBudget() async throws {
+        let tables = [
+            TestFixtures.makeTableInfo(name: "audit_log"),
+            TestFixtures.makeTableInfo(name: "orders")
+        ]
+        let provider = Self.provider(
+            tables: tables,
+            columns: [
+                "audit_log": [TestFixtures.makeColumnInfo(name: "message")],
+                "orders": [TestFixtures.makeColumnInfo(name: "total")]
+            ]
+        )
+        await provider.resetForDatabase(
+            "shop",
+            tables: tables,
+            driver: MockDatabaseDriver(),
+            connection: TestFixtures.makeConnection()
+        )
+        var settings = AISettings.default
+        settings.maxSchemaTables = 1
+        let columnList = (0..<500).map { "c\($0)" }.joined(separator: ", ")
+        let editorText = "SELECT \(columnList) FROM orders WHERE "
+
+        let context = try #require(await provider.buildSchemaContextForAI(
+            settings: settings,
+            editorText: editorText,
+            cursorOffset: (editorText as NSString).length
+        ))
+
+        #expect(Self.columnNamesByTableLine(in: context)["orders"] == ["total"])
+    }
+
     private static func schemaAwareProvider(
         columnsBySchema: [String: [ColumnInfo]],
         bulkColumns: [String: [ColumnInfo]] = [:]

From 6c879239a8bd839517877d88a2acf1b5fddc6bb5 Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 18:15:15 +0700
Subject: [PATCH 06/10] fix(ai): keep tables a bulk read left out unknown, and
 match the exact table spelling first

---
 apple/TablePro/Core/AI/AISchemaContext.swift      | 15 ++++++++++-----
 .../AIChatViewModel+SchemaContext.swift           |  6 +++++-
 .../Core/AI/AISchemaContextRelevanceTests.swift   | 10 ++++++++++
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/apple/TablePro/Core/AI/AISchemaContext.swift b/apple/TablePro/Core/AI/AISchemaContext.swift
index e966c5b691..5cac8d4eb7 100644
--- a/apple/TablePro/Core/AI/AISchemaContext.swift
+++ b/apple/TablePro/Core/AI/AISchemaContext.swift
@@ -212,13 +212,18 @@ struct AISchemaContext {
             let named = tables.indices.filter {
                 tables[$0].name.caseInsensitiveCompare(reference.tableName) == .orderedSame
             }
+            // Two tables can differ only by case, so the exact spelling wins.
+            let candidates = named.filter { tables[$0].name == reference.tableName }
+                + named.filter { tables[$0].name != reference.tableName }
             let schema = reference.schema ?? defaultSchema
-            let inSchema = named.first { offset in
-                guard let schema, let tableSchema = tables[offset].schema else { return false }
-                return tableSchema.caseInsensitiveCompare(schema) == .orderedSame
+            func inSchema(_ matches: (String, String) -> Bool) -> Int? {
+                guard let schema else { return nil }
+                return candidates.first { tables[$0].schema.map { matches($0, schema) } ?? false }
             }
-            let unscoped = named.first { tables[$0].schema == nil }
-            let match = inSchema ?? unscoped ?? (reference.schema == nil ? named.first : nil)
+            let match = inSchema(==)
+                ?? inSchema { $0.caseInsensitiveCompare($1) == .orderedSame }
+                ?? candidates.first { tables[$0].schema == nil }
+                ?? (reference.schema == nil ? candidates.first : nil)
             if let match, !result.contains(match) {
                 result.append(match)
             }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index c69a20167a..56e7d5d42b 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -275,8 +275,12 @@ extension AIChatViewModel {
             uniquingKeysWith: { first, _ in first }
         )
         var merged = columnsByTable
+        // A table the read left out, such as a SQLite view, stays unknown, so the ranker never
+        // takes this read for a complete one.
         for table in tables where (merged[table.name] ?? []).isEmpty {
-            merged[table.name] = allColumns[table.name] ?? byFoldedName[table.name.lowercased()] ?? []
+            if let columns = allColumns[table.name] ?? byFoldedName[table.name.lowercased()] {
+                merged[table.name] = columns
+            }
         }
         columnsByTable = merged
     }
diff --git a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
index 9c96181dd6..a0a02467b3 100644
--- a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
+++ b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
@@ -120,6 +120,16 @@ struct AISchemaContextRelevanceTests {
         #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [1])
     }
 
+    @Test("Two tables that differ only by case resolve to the exact spelling")
+    func exactSpellingWins() {
+        let tables = [
+            TestFixtures.makeTableInfo(name: "Orders", schema: "public"),
+            TestFixtures.makeTableInfo(name: "orders", schema: "public")
+        ]
+        let references = AISchemaContext.referencedTables(in: "SELECT * FROM \"orders\"", grammar: Self.grammar)
+        #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [1])
+    }
+
     @Test("A query that ends in a comment still puts its tables first")
     func queryEndingInCommentCounts() {
         #expect(Self.names(question: "", query: "SELECT * FROM invoices -- unpaid ones") == ["invoices"])

From 5477979907e23535a6376b90adc11a76de752313 Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 18:20:30 +0700
Subject: [PATCH 07/10] fix(ai): load foreign keys apart from columns, and keep
 a table a CTE shares its name with

---
 .../SQLContextAnalyzer.swift                  |  8 ++---
 ...QLContextAnalyzerTableReferenceTests.swift |  9 +++++
 apple/TablePro/Core/AI/AISchemaContext.swift  |  4 +--
 .../AIChatViewModel+SchemaContext.swift       | 36 +++++++++++--------
 .../AI/AISchemaContextRelevanceTests.swift    |  8 ++++-
 5 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
index 3feb9a5ce2..90609b9e59 100644
--- a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
+++ b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
@@ -289,16 +289,12 @@ public final class SQLContextAnalyzer {
     public init() {}
 
     // `analyze` reads no tables when its cursor sits in a comment or literal; this reads them wherever the statement
-    // ends, so `SELECT * FROM orders -- note` still names `orders`.
+    // ends, so `SELECT * FROM orders -- note` still names `orders`. A CTE's name comes back like a table's, since a
+    // CTE can share a table's name and still read that table in its own body.
     public func tableReferences(in statement: String, grammar: SQLLexicalGrammar) -> [TableReference] {
         let nsStatement = statement as NSString
         let text = Self.lex(nsStatement, cursor: nsStatement.length, grammar: grammar)
         var references = extractTableReferences(from: text.forReferences)
-        let tokens = Self.tokens(of: text, grammar: grammar)
-        mergeDerivedTables(
-            Self.derivedTableParser.visibleTables(in: tokens, cursor: nsStatement.length),
-            into: &references
-        )
         if let alterTableName = extractAlterTableName(from: text.forReferences),
            !references.contains(where: { $0.tableName == alterTableName }) {
             references.append(TableReference(tableName: alterTableName, alias: nil))
diff --git a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
index c4d8472f61..62dd0d98d0 100644
--- a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
+++ b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
@@ -22,6 +22,15 @@ struct SQLContextAnalyzerTableReferenceTests {
         #expect(tables("SELECT * FROM orders WHERE note = 'open") == ["orders"])
     }
 
+    @Test("A CTE that shares a table's name still names the table its body reads")
+    func shadowingCommonTableExpression() {
+        let references = analyzer.tableReferences(
+            in: "WITH orders AS (SELECT * FROM orders WHERE active) SELECT * FROM orders",
+            grammar: .ansi
+        )
+        #expect(references.contains { $0.tableName == "orders" && $0.derivedColumns == nil })
+    }
+
     @Test("Joined, updated and altered tables are named")
     func everyClause() {
         #expect(tables("SELECT * FROM orders o JOIN customers c ON c.id = o.customer_id") == ["orders", "customers"])
diff --git a/apple/TablePro/Core/AI/AISchemaContext.swift b/apple/TablePro/Core/AI/AISchemaContext.swift
index 5cac8d4eb7..1d21cf143a 100644
--- a/apple/TablePro/Core/AI/AISchemaContext.swift
+++ b/apple/TablePro/Core/AI/AISchemaContext.swift
@@ -196,11 +196,11 @@ struct AISchemaContext {
         return nsQuery.length > maxQueryLength ? nsQuery.substring(to: maxQueryLength) : query
     }
 
-    // A CTE or subquery has no table of its own, so it is left out.
+    // A CTE's name matches no table in the catalog, so `indices` drops it.
     static func referencedTables(in query: String, grammar: SQLLexicalGrammar) -> [TableReference] {
         let analyzer = SQLContextAnalyzer()
         return SQLStatementScanner.allStatements(in: query, grammar: grammar).flatMap { statement in
-            analyzer.tableReferences(in: statement, grammar: grammar).filter { $0.derivedColumns == nil }
+            analyzer.tableReferences(in: statement, grammar: grammar)
         }
     }
 
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index 56e7d5d42b..475714a5de 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -25,7 +25,9 @@ extension AIChatViewModel {
 
     func ensureColumnsLoaded(forTable tableName: String) async {
         dropSchemaCacheIfScopeChanged()
-        if let existing = columnsByTable[tableName], !existing.isEmpty { return }
+        // The ranking's bulk read fills columns without foreign keys, so each is checked on its own.
+        let needsColumns = (columnsByTable[tableName] ?? []).isEmpty
+        if !needsColumns, foreignKeysByTable[tableName] != nil { return }
         if let inFlight = inFlightColumnFetches[tableName] {
             await inFlight.value
             return
@@ -34,28 +36,32 @@ extension AIChatViewModel {
         let connId = connection.id
         let scope = schemaCacheScope
         let task: Task<Void, Never> = Task { [weak self] in
-            let columns: [ColumnInfo]
-            do {
-                columns = try await DatabaseManager.shared.withBrowseMetadataDriver(connectionId: connId) { driver in
-                    try await driver.fetchColumns(table: tableName)
+            var columns: [ColumnInfo]?
+            if needsColumns {
+                do {
+                    columns = try await DatabaseManager.shared.withBrowseMetadataDriver(connectionId: connId) { driver in
+                        try await driver.fetchColumns(table: tableName)
+                    }
+                } catch {
+                    Self.logger.warning("Column fetch failed for \(tableName, privacy: .private(mask: .hash)): \(error.publicLogShape, privacy: .public)")
+                    columns = []
                 }
-            } catch {
-                Self.logger.warning("Column fetch failed for \(tableName, privacy: .private(mask: .hash)): \(error.publicLogShape, privacy: .public)")
-                columns = []
             }
-            let fkMap: [String: [ForeignKeyInfo]]
+            let fkMap: [String: [ForeignKeyInfo]]?
             do {
                 fkMap = try await DatabaseManager.shared.withBrowseMetadataDriver(connectionId: connId) { driver in
                     try await driver.fetchForeignKeys(forTables: [tableName])
                 }
             } catch {
                 Self.logger.warning("Foreign key fetch failed for \(tableName, privacy: .private(mask: .hash)): \(error.publicLogShape, privacy: .public)")
-                fkMap = [:]
+                fkMap = nil
             }
             guard !Task.isCancelled, let self, self.schemaCacheScope == scope else { return }
-            self.columnsByTable[tableName] = columns
-            if let fks = fkMap[tableName] {
-                self.foreignKeysByTable[tableName] = fks
+            if let columns {
+                self.columnsByTable[tableName] = columns
+            }
+            if let fkMap {
+                self.foreignKeysByTable[tableName] = fkMap[tableName] ?? []
             }
             self.inFlightColumnFetches[tableName] = nil
         }
@@ -143,8 +149,8 @@ extension AIChatViewModel {
                 try await driver.fetchForeignKeys(forTables: tablesToFetch.map(\.name))
             }
             guard schemaCacheScope == scope else { return }
-            for (name, fks) in fkMap {
-                foreignKeysByTable[name] = fks
+            for table in tablesToFetch {
+                foreignKeysByTable[table.name] = fkMap[table.name] ?? []
             }
         } catch {
             Self.logger.warning("Foreign key bulk fetch failed: \(error.publicLogShape, privacy: .public)")
diff --git a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
index a0a02467b3..878d117d63 100644
--- a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
+++ b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
@@ -91,7 +91,7 @@ struct AISchemaContextRelevanceTests {
         #expect(names == ["customers"])
     }
 
-    @Test("Every statement in the editor counts, and a CTE is not a table")
+    @Test("Every statement in the editor counts, and a CTE name matches no table")
     func everyStatementCounts() {
         let query = """
         WITH recent AS (SELECT * FROM orders) SELECT * FROM recent;
@@ -100,6 +100,12 @@ struct AISchemaContextRelevanceTests {
         #expect(Self.names(question: "", query: query) == ["orders", "customers"])
     }
 
+    @Test("A CTE named like a table puts that table first")
+    func shadowingCommonTableExpression() {
+        let query = "WITH orders AS (SELECT * FROM orders WHERE total > 0) SELECT * FROM orders"
+        #expect(Self.names(question: "", query: query) == ["orders"])
+    }
+
     @Test("A schema-qualified reference matches only a table in that schema")
     func schemaQualifiedReference() {
         let tables = [

From 364f523c2edb29f53a5e55eba30e58130a488ce5 Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 18:27:09 +0700
Subject: [PATCH 08/10] fix(ai): put every table a referenced name could mean
 first, and keep an ALTER target's schema

---
 .../SQLContextAnalyzer.swift                  | 27 ++++++++++------
 ...QLContextAnalyzerTableReferenceTests.swift |  9 ++++++
 apple/TablePro/Core/AI/AISchemaContext.swift  | 32 ++++++++-----------
 .../AI/AISchemaContextRelevanceTests.swift    | 25 +++++++++++----
 4 files changed, 59 insertions(+), 34 deletions(-)

diff --git a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
index 90609b9e59..d2a6349ab4 100644
--- a/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
+++ b/apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift
@@ -278,7 +278,9 @@ public final class SQLContextAnalyzer {
         "|WINDOW|FETCH|FOR)\\b|[;()]|$)"
     )
 
-    private static let alterTableRegex = compileRegex("(?i)\\bALTER\\s+TABLE\\s+(\(quotedSegment))")
+    private static let alterTableRegex = compileRegex(
+        "(?i)\\bALTER\\s+TABLE\\s+(\(quotedSegment)(?:\\.\(quotedSegment))*)"
+    )
 
     private static let derivedTableParser = DerivedTableParser()
 
@@ -295,9 +297,8 @@ public final class SQLContextAnalyzer {
         let nsStatement = statement as NSString
         let text = Self.lex(nsStatement, cursor: nsStatement.length, grammar: grammar)
         var references = extractTableReferences(from: text.forReferences)
-        if let alterTableName = extractAlterTableName(from: text.forReferences),
-           !references.contains(where: { $0.tableName == alterTableName }) {
-            references.append(TableReference(tableName: alterTableName, alias: nil))
+        if let alterRef = extractAlterTableReference(from: text.forReferences), !references.contains(alterRef) {
+            references.append(alterRef)
         }
         return references
     }
@@ -342,8 +343,7 @@ public final class SQLContextAnalyzer {
 
         var seenReferences = Set<TableReference>(tableReferences)
 
-        if let alterTableName = extractAlterTableName(from: text.forReferences) {
-            let alterRef = TableReference(tableName: alterTableName, alias: nil)
+        if let alterRef = extractAlterTableReference(from: text.forReferences) {
             if seenReferences.insert(alterRef).inserted {
                 tableReferences.append(alterRef)
             }
@@ -1017,12 +1017,19 @@ public final class SQLContextAnalyzer {
         )
     }
 
-    private func extractAlterTableName(from text: LexedStatement) -> String? {
+    private func extractAlterTableReference(from text: LexedStatement) -> TableReference? {
         let nsRange = NSRange(location: 0, length: text.code.length)
         guard let match = Self.alterTableRegex.firstMatch(in: text.codeString, range: nsRange) else { return nil }
-        let tableNSRange = match.range(at: 1)
-        guard tableNSRange.location != NSNotFound else { return nil }
-        return Self.unquoted(text.original.substring(with: tableNSRange))
+        let pathRange = match.range(at: 1)
+        guard pathRange.location != NSNotFound else { return nil }
+        let segments = Self.pathSegments(pathRange, in: text)
+        guard let table = segments.last, !table.name.isEmpty else { return nil }
+        return TableReference(
+            tableName: table.name,
+            alias: nil,
+            schema: segments.count >= 2 ? segments[segments.count - 2].name : nil,
+            isIdentifierQuoted: table.isQuoted
+        )
     }
 
     // MARK: - Clause Detection
diff --git a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
index 62dd0d98d0..f45cfea954 100644
--- a/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
+++ b/apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift
@@ -37,4 +37,13 @@ struct SQLContextAnalyzerTableReferenceTests {
         #expect(tables("UPDATE customers SET name = 'x'").contains("customers"))
         #expect(tables("ALTER TABLE invoices ADD COLUMN due DATE") == ["invoices"])
     }
+
+    @Test("A schema-qualified ALTER target keeps its schema")
+    func qualifiedAlterTarget() {
+        let references = analyzer.tableReferences(in: "ALTER TABLE archive.orders ADD COLUMN due DATE", grammar: .ansi)
+        #expect(references.map(\.tableName) == ["orders"])
+        #expect(references.first?.schema == "archive")
+        let context = analyzer.analyze(query: "ALTER TABLE archive.orders DROP COLUMN ", cursorPosition: 39)
+        #expect(context.tableReferences.contains { $0.tableName == "orders" && $0.schema == "archive" })
+    }
 }
diff --git a/apple/TablePro/Core/AI/AISchemaContext.swift b/apple/TablePro/Core/AI/AISchemaContext.swift
index 1d21cf143a..8d09184397 100644
--- a/apple/TablePro/Core/AI/AISchemaContext.swift
+++ b/apple/TablePro/Core/AI/AISchemaContext.swift
@@ -204,28 +204,24 @@ struct AISchemaContext {
         }
     }
 
-    // In the order the query names them, each table once. A bare name means the default schema's
-    // table, as it does when the query runs.
+    // Every table a name could mean, in the order the query names them. How an engine folds an
+    // unquoted name differs, so a name that matches two tables by case or schema puts both first:
+    // a spare slot is cheaper than missing the table the query reads. The default schema's comes
+    // first for a bare name, as it does when the query runs.
     static func indices(of references: [TableReference], in tables: [TableInfo], defaultSchema: String? = nil) -> [Int] {
+        func sameName(_ lhs: String?, _ rhs: String?) -> Bool {
+            guard let lhs, let rhs else { return false }
+            return lhs.caseInsensitiveCompare(rhs) == .orderedSame
+        }
         var result: [Int] = []
         for reference in references {
-            let named = tables.indices.filter {
-                tables[$0].name.caseInsensitiveCompare(reference.tableName) == .orderedSame
-            }
-            // Two tables can differ only by case, so the exact spelling wins.
-            let candidates = named.filter { tables[$0].name == reference.tableName }
-                + named.filter { tables[$0].name != reference.tableName }
+            let named = tables.indices.filter { sameName(tables[$0].name, reference.tableName) }
             let schema = reference.schema ?? defaultSchema
-            func inSchema(_ matches: (String, String) -> Bool) -> Int? {
-                guard let schema else { return nil }
-                return candidates.first { tables[$0].schema.map { matches($0, schema) } ?? false }
-            }
-            let match = inSchema(==)
-                ?? inSchema { $0.caseInsensitiveCompare($1) == .orderedSame }
-                ?? candidates.first { tables[$0].schema == nil }
-                ?? (reference.schema == nil ? candidates.first : nil)
-            if let match, !result.contains(match) {
-                result.append(match)
+            let inSchema = named.filter { sameName(tables[$0].schema, schema) }
+            let unscoped = named.filter { tables[$0].schema == nil }
+            let elsewhere = reference.schema == nil ? named : []
+            for offset in inSchema + unscoped + elsewhere where !result.contains(offset) {
+                result.append(offset)
             }
         }
         return result
diff --git a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
index 878d117d63..8ea0fa4908 100644
--- a/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
+++ b/apple/TableProTests/Core/AI/AISchemaContextRelevanceTests.swift
@@ -116,24 +116,37 @@ struct AISchemaContextRelevanceTests {
         #expect(AISchemaContext.indices(of: references, in: tables) == [1])
     }
 
-    @Test("A bare name means the default schema's table")
+    @Test("A bare name puts the default schema's table first")
     func bareNameMeansDefaultSchema() {
         let tables = [
             TestFixtures.makeTableInfo(name: "orders", schema: "archive"),
             TestFixtures.makeTableInfo(name: "orders", schema: "public")
         ]
         let references = AISchemaContext.referencedTables(in: "SELECT * FROM orders", grammar: Self.grammar)
-        #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [1])
+        #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [1, 0])
     }
 
-    @Test("Two tables that differ only by case resolve to the exact spelling")
-    func exactSpellingWins() {
+    @Test("A name that matches two tables by case puts both first")
+    func caseCollisionKeepsBoth() {
         let tables = [
+            TestFixtures.makeTableInfo(name: "audit_log", schema: "public"),
             TestFixtures.makeTableInfo(name: "Orders", schema: "public"),
             TestFixtures.makeTableInfo(name: "orders", schema: "public")
         ]
-        let references = AISchemaContext.referencedTables(in: "SELECT * FROM \"orders\"", grammar: Self.grammar)
-        #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [1])
+        for query in ["SELECT * FROM Orders", "SELECT * FROM \"orders\""] {
+            let references = AISchemaContext.referencedTables(in: query, grammar: Self.grammar)
+            #expect(Set(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public")) == [1, 2])
+        }
+    }
+
+    @Test("A qualified name leaves another schema's table out")
+    func qualifiedNameStaysInItsSchema() {
+        let tables = [
+            TestFixtures.makeTableInfo(name: "orders", schema: "archive"),
+            TestFixtures.makeTableInfo(name: "orders", schema: "public")
+        ]
+        let references = AISchemaContext.referencedTables(in: "ALTER TABLE archive.orders ADD due DATE", grammar: Self.grammar)
+        #expect(AISchemaContext.indices(of: references, in: tables, defaultSchema: "public") == [0])
     }
 
     @Test("A query that ends in a comment still puts its tables first")

From 525b97d3bbc895b0456017d50e3487334b341a02 Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 18:33:55 +0700
Subject: [PATCH 09/10] fix(ai): rank a custom command by its body, and trust a
 case-folded bulk entry only when it means one table

---
 .../AIChatViewModel+SchemaContext.swift         | 12 +++++++-----
 .../AIChatViewModel+SlashCommands.swift         |  2 +-
 .../ViewModels/AIChatViewModel+Streaming.swift  | 10 +++++++++-
 apple/TablePro/ViewModels/AIChatViewModel.swift |  7 ++++++-
 .../ViewModels/AIChatViewModelSlashTests.swift  | 17 +++++++++++++++++
 5 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
index 475714a5de..4c30f410e6 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SchemaContext.swift
@@ -276,15 +276,17 @@ extension AIChatViewModel {
             return
         }
         guard let allColumns, !Task.isCancelled, schemaCacheScope == scope else { return }
-        let byFoldedName = Dictionary(
-            allColumns.map { ($0.key.lowercased(), $0.value) },
-            uniquingKeysWith: { first, _ in first }
-        )
+        // A name that differs only by case is trusted only when it can mean one table, or a
+        // zero-column `orders` would take a quoted `"Orders"`'s columns.
+        let foldedTableCounts = Dictionary(tables.map { ($0.name.lowercased(), 1) }, uniquingKeysWith: +)
+        let byFoldedName = Dictionary(grouping: allColumns, by: { $0.key.lowercased() })
+            .compactMapValues { $0.count == 1 ? $0[0].value : nil }
         var merged = columnsByTable
         // A table the read left out, such as a SQLite view, stays unknown, so the ranker never
         // takes this read for a complete one.
         for table in tables where (merged[table.name] ?? []).isEmpty {
-            if let columns = allColumns[table.name] ?? byFoldedName[table.name.lowercased()] {
+            let folded = table.name.lowercased()
+            if let columns = allColumns[table.name] ?? (foldedTableCounts[folded] == 1 ? byFoldedName[folded] : nil) {
                 merged[table.name] = columns
             }
         }
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift b/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
index b74309ed2a..700113965c 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+SlashCommands.swift
@@ -58,7 +58,7 @@ extension AIChatViewModel {
         )
         let prompt = CustomSlashCommandRenderer.render(command, context: renderingContext)
         messages.append(ChatTurn(role: .user, blocks: [.text(invocationText)]))
-        sendWithContext(prompt: prompt)
+        sendWithContext(prompt: prompt, schemaQuestion: needsSchema ? body : nil)
     }
 
     private func showHelp(invocationText: String) {
diff --git a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
index ed2d2df00a..8c235da65a 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
@@ -127,7 +127,7 @@ extension AIChatViewModel {
         let assistantID = assistantMessage.id
         streamingState = .streaming(assistantID: assistantID)
 
-        let question = messages.last { $0.role == .user }?.plainText
+        let question = Self.schemaQuestion(in: messages, override: schemaQuestionOverride)
         prepTask = Task.detached(priority: .userInitiated) { [weak self] in
             guard let self else { return }
             if settings.includeSchema {
@@ -514,6 +514,14 @@ extension AIChatViewModel {
         )
     }
 
+    static func schemaQuestion(in messages: [ChatTurn], override: (turnID: UUID, text: String)?) -> String? {
+        let lastUserTurn = messages.last { $0.role == .user }
+        if let override, override.turnID == lastUserTurn?.id {
+            return override.text
+        }
+        return lastUserTurn?.plainText
+    }
+
     nonisolated static func buildSystemPrompt(
         _ promptContext: PromptContext?,
         mode: AIChatMode,
diff --git a/apple/TablePro/ViewModels/AIChatViewModel.swift b/apple/TablePro/ViewModels/AIChatViewModel.swift
index 24260bb851..3df1451e56 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel.swift
@@ -110,6 +110,9 @@ final class AIChatViewModel: ObservableObject {
     var relevantTableNames: [String] = []
     var bulkColumnReadTableCount: Int?
     var schemaCacheScope: String?
+    // A custom command's message carries the schema it rendered, so its tables are ranked by the
+    // user's own words instead.
+    var schemaQuestionOverride: (turnID: UUID, text: String)?
     nonisolated(unsafe) var streamingTask: Task<Void, Never>?
     var prepTask: Task<Void, Never>?
 
@@ -241,8 +244,9 @@ final class AIChatViewModel: ObservableObject {
         return descriptor.supportsImages(fetched: catalog.fetchedInfo(providerID: config.id, modelID: model))
     }
 
-    func sendWithContext(prompt: String) {
+    func sendWithContext(prompt: String, schemaQuestion: String? = nil) {
         let userMessage = ChatTurn(role: .user, blocks: [.text(prompt)])
+        schemaQuestionOverride = schemaQuestion.map { (userMessage.id, $0) }
         messages.append(userMessage)
         trimMessagesIfNeeded()
         clearError()
@@ -371,6 +375,7 @@ final class AIChatViewModel: ObservableObject {
         relevantTableNames = []
         bulkColumnReadTableCount = nil
         schemaCacheScope = nil
+        schemaQuestionOverride = nil
         currentQuery = nil
         queryResults = nil
         editorTarget = nil
diff --git a/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift b/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift
index 346754329c..5818fc5067 100644
--- a/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift
+++ b/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift
@@ -97,4 +97,21 @@ struct AIChatViewModelSlashTests {
         #expect(vm.inputText.isEmpty)
         #expect(vm.errorMessage == nil)
     }
+
+    @Test("A custom command's turn ranks tables by its body, not by the schema it rendered")
+    func customCommandRanksByBody() {
+        let rendered = ChatTurn(role: .user, blocks: [.text("Schema:\n- customers\n  - customer_id\nlist the logs")])
+        let messages = [ChatTurn(role: .user, blocks: [.text("/report logs")]), rendered]
+
+        #expect(AIChatViewModel.schemaQuestion(in: messages, override: (rendered.id, "logs")) == "logs")
+        #expect(AIChatViewModel.schemaQuestion(in: messages, override: nil) == rendered.plainText)
+    }
+
+    @Test("A custom command's body does not rank a later message")
+    func staleOverrideIsIgnored() {
+        let earlier = ChatTurn(role: .user, blocks: [.text("rendered")])
+        let later = ChatTurn(role: .user, blocks: [.text("how many orders")])
+
+        #expect(AIChatViewModel.schemaQuestion(in: [earlier, later], override: (earlier.id, "logs")) == "how many orders")
+    }
 }

From 9bc07fd7609a86fa8d433120ee37f1dc93e34faa Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Thu, 22 Oct 2026 18:39:15 +0700
Subject: [PATCH 10/10] fix(ai): rank by the user's question, not a tool result
 turn

---
 .../ViewModels/AIChatViewModel+Streaming.swift      |  8 +++++++-
 .../ViewModels/AIChatViewModelSlashTests.swift      | 13 +++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
index 8c235da65a..25000d75a0 100644
--- a/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
+++ b/apple/TablePro/ViewModels/AIChatViewModel+Streaming.swift
@@ -515,7 +515,13 @@ extension AIChatViewModel {
     }
 
     static func schemaQuestion(in messages: [ChatTurn], override: (turnID: UUID, text: String)?) -> String? {
-        let lastUserTurn = messages.last { $0.role == .user }
+        // A turn that only returns tool results is the app's, not the user's.
+        let lastUserTurn = messages.last { turn in
+            turn.role == .user && !turn.blocks.allSatisfy { block in
+                if case .toolResult = block.kind { return true }
+                return false
+            }
+        }
         if let override, override.turnID == lastUserTurn?.id {
             return override.text
         }
diff --git a/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift b/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift
index 5818fc5067..f9aa71cb7a 100644
--- a/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift
+++ b/apple/TableProTests/ViewModels/AIChatViewModelSlashTests.swift
@@ -114,4 +114,17 @@ struct AIChatViewModelSlashTests {
 
         #expect(AIChatViewModel.schemaQuestion(in: [earlier, later], override: (earlier.id, "logs")) == "how many orders")
     }
+
+    @Test("A tool result turn does not replace the user's question")
+    func toolResultTurnIsSkipped() {
+        let question = ChatTurn(role: .user, blocks: [.text("how many orders")])
+        let toolResult = ChatTurn(
+            role: .user,
+            blocks: [.toolResult(ToolResultBlock(toolUseId: "call_1", content: "rows", isError: false))]
+        )
+        let messages = [question, ChatTurn(role: .assistant, blocks: []), toolResult]
+
+        #expect(AIChatViewModel.schemaQuestion(in: messages, override: nil) == "how many orders")
+        #expect(AIChatViewModel.schemaQuestion(in: messages, override: (question.id, "orders")) == "orders")
+    }
 }