@finos/legend-lego 2.0.202 → 2.0.204

package/lib/legend-ai/stores/LegendAIChatProcessors.js

@@ -13,23 +13,30 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { assertErrorThrown, uuid } from '@finos/legend-shared';
-import { classifyQuestionIntentFast, LegendAIQuestionIntent, LegendAIThinkingStepStatus, LegendAIMessageRole, LegendAIErrorType, LegendAIServiceError, TDSServiceSourceType, buildColumnDefsFromNames, LEGEND_AI_ORCHESTRATOR_FALLBACK_ACTION_ID, getTodayISO, } from '../LegendAITypes.js';
+import { assertErrorThrown, isPlainObject, uuid } from '@finos/legend-shared';
+import { LegendAIQuestionIntent, LegendAIThinkingStepStatus, LegendAIMessageRole, LegendAIErrorType, LegendAIServiceError, LegendAIUnsupportedEngineShapeError, TDSServiceSourceType, buildColumnDefsFromNames, LEGEND_AI_ORCHESTRATOR_FALLBACK_ACTION_ID, getTodayISO, } from '../LegendAITypes.js';
 import { LegendAIJudgeVerdict, } from '../LegendAI_LegendApplicationPlugin_Extension.js';
 import { extractElementNameFromPath, } from '@finos/legend-graph';
+import { buildEnrichedBusinessContext, buildModelContextEnrichmentText, findBestAlternateRoot, splitIdentifierTokens, tokenizeText, } from '../LegendAIDocEnrichment.js';
+import { isFuzzyMatch, preFilterServicesByRelevance, } from '../LegendAIServiceRetrieval.js';
+import { isNumericColumn, isStringColumn, } from '../components/LegendAIChatHelpers.js';
 const MAX_ERROR_MESSAGE_LENGTH = 500;
 const MAX_THINKING_ERROR_PREVIEW_LENGTH = 200;
 const DEFAULT_MAX_JUDGE_ATTEMPTS = 5;
 const DEFAULT_MAX_EXECUTION_RETRIES = 3;
 const ANALYSIS_TIMEOUT_MS = 15_000;
+const ORCHESTRATOR_GENERATION_TIMEOUT_MS = 120_000;
+const EXECUTION_TIMEOUT_MS = 300_000;
 const ANALYSIS_PREVIEW_ROW_LIMIT = 3;
 const ANALYSIS_PREVIEW_VALUE_LIMIT = 40;
 const MAX_NON_SQL_PASS_ATTEMPTS = 2;
-const ALIAS_DOT_COL_PATTERN = /\b(?<tbl>[a-z]\w*)\s*\.\s*"(?<col>[^"]+)"/gi;
 const JOIN_PATTERN = /\bJOIN\b/i;
 const ORDER_BY_SPLIT = /\bORDER\s+BY\b/i;
 const UNION_ALL_PATTERN = /\bUNION\s+ALL\b/i;
 const LITERAL_COL_PATTERN = /,\s*'[^']*'\s+AS\s+(?:"[^"]+"|[a-z]\w*)/gi;
+const SELECT_ALIAS_PATTERN = /\b(?<tbl>[a-z]\w*)\s*\.\s*"(?<col>[^"]+)"\s+AS\s+(?:"(?<qAlias>[^"]+)"|(?<uAlias>\w+))/gi;
+const ALIAS_DOT_COL_PATTERN = /\b(?<tbl>[a-z]\w*)\s*\.\s*"(?<col>[^"]+)"/gi;
+const SERVICE_CALL_PATTERN = /\bservice\s*\([^()]*\)/gi;
 const DEFAULT_SAFETY_LIMIT = 1000;
 const HAS_LIMIT_PATTERN = /\bLIMIT\s+\d+/i;
 const HAS_AGGREGATION_PATTERN = /\bGROUP\s+BY\b|\bCOUNT\s*\(|\bSUM\s*\(|\bAVG\s*\(|\bMIN\s*\(|\bMAX\s*\(/i;
@@ -42,6 +49,7 @@ const SERVICE_PARAM_DATE_LIKE_PATTERNS = [
     /effective|valid|settle|trade|maturity|expir|inception|close|open/i,
     /start|end|from|until|begin|report|cutoff|valuation|pricing/i,
 ];
+const GENERIC_TABLE_PATTERNS = /combined|consolidated|all|master|summary/i;
 function isLikelySqlQuery(text) {
     const trimmed = text.trim().toLowerCase();
     return (trimmed.startsWith('select') ||
@@ -93,6 +101,7 @@ export function createMessagePair(text) {
             isExecuting: false,
             suggestedQueries: [],
             fallbackAction: null,
+            queriedAccessPointGroups: [],
         },
     ];
 }
@@ -130,6 +139,54 @@ export function classifyError(error) {
     }
     return LegendAIErrorType.GENERAL;
 }
+export var ExecutionErrorCategory;
+(function (ExecutionErrorCategory) {
+    ExecutionErrorCategory["INFRASTRUCTURE"] = "INFRASTRUCTURE";
+    ExecutionErrorCategory["SQL_FIXABLE"] = "SQL_FIXABLE";
+    ExecutionErrorCategory["ACCESS"] = "ACCESS";
+    ExecutionErrorCategory["NONE"] = "NONE";
+})(ExecutionErrorCategory || (ExecutionErrorCategory = {}));
+const EXECUTION_ERROR_RULES = [
+    {
+        category: ExecutionErrorCategory.ACCESS,
+        patterns: [
+            /\binsufficient privileges\b/i,
+            /\baccess denied\b/i,
+            /\bpermission denied\b/i,
+            /\bunauthorized\b/i,
+            /\bentitlement\b/i,
+        ],
+    },
+    {
+        category: ExecutionErrorCategory.INFRASTRUCTURE,
+        patterns: [/__lake_action/i, /__lake/i],
+    },
+    {
+        category: ExecutionErrorCategory.SQL_FIXABLE,
+        patterns: [
+            /\binvalid identifier\b/i,
+            /\bambiguous column\b/i,
+            /\bdoes not exist\b/i,
+            /\bsql compilation error\b/i,
+            /\bnot found\b/i,
+        ],
+    },
+];
+export function categorizeExecutionError(errMsg, error) {
+    if (error instanceof LegendAIUnsupportedEngineShapeError) {
+        return ExecutionErrorCategory.SQL_FIXABLE;
+    }
+    if (error instanceof LegendAIServiceError &&
+        error.errorType === LegendAIErrorType.PERMISSION) {
+        return ExecutionErrorCategory.ACCESS;
+    }
+    for (const rule of EXECUTION_ERROR_RULES) {
+        if (rule.patterns.some((pattern) => pattern.test(errMsg))) {
+            return rule.category;
+        }
+    }
+    return ExecutionErrorCategory.NONE;
+}
 export function finishWithThinkingError(setMessages, errorMsg, startTime, errorType) {
     updateLastAssistant(setMessages, (msg) => ({
         thinkingSteps: msg.thinkingSteps.map((s) => s.status === LegendAIThinkingStepStatus.ACTIVE
@@ -148,7 +205,10 @@ function buildTurnFromAssistant(userText, asstMsg) {
             sql: asstMsg.sql,
             intent: LegendAIQuestionIntent.DATA_QUERY,
         };
-        if (asstMsg.gridData) {
+        if (asstMsg.error) {
+            turn.resultSummary = `ERROR: ${asstMsg.error.slice(0, 200)}`;
+        }
+        else if (asstMsg.gridData) {
             turn.rowCount = asstMsg.gridData.rowData.length;
             const colNames = asstMsg.gridData.columnDefs
                 .map((c) => c.headerName ?? c.colId ?? '')
@@ -166,6 +226,14 @@ function buildTurnFromAssistant(userText, asstMsg) {
             intent: LegendAIQuestionIntent.METADATA,
         };
     }
+    if (asstMsg.error && !asstMsg.sql) {
+        return {
+            question: userText,
+            sql: '(generation failed)',
+            resultSummary: `ERROR: ${asstMsg.error.slice(0, 200)}`,
+            intent: LegendAIQuestionIntent.DATA_QUERY,
+        };
+    }
     return undefined;
 }
 export function buildConversationHistory(messages) {
@@ -224,7 +292,21 @@ export function buildGenerationFailureMessage(failure, suggestion, services) {
     return parts.join('');
 }
 function buildFallbackSuggestions(services) {
-    return services.slice(0, 3).map((svc) => `Show 10 records from ${svc.title}`);
+    const result = [];
+    for (const svc of services.slice(0, 3)) {
+        const strCol = svc.columns.find(isStringColumn);
+        const numCol = svc.columns.find(isNumericColumn);
+        if (numCol && strCol) {
+            result.push(`What are the top 10 ${strCol.name} values by ${numCol.name} in ${svc.title}?`);
+        }
+        else if (strCol) {
+            result.push(`Show the breakdown by ${strCol.name} in ${svc.title}`);
+        }
+        else {
+            result.push(`Show 10 records from ${svc.title}`);
+        }
+    }
+    return result;
 }
 function appendFallbackSuggestions(setMessages, services) {
     const suggestions = buildFallbackSuggestions(services);
@@ -234,6 +316,28 @@ function appendFallbackSuggestions(setMessages, services) {
         }));
     }
 }
+function buildSnowflakeSqlError(rawError, services) {
+    const lowerError = rawError.toLowerCase();
+    const coreMatch = /SnowflakeSQLException:\s*(?<core>SQL compilation error[^"]*?)(?:\\n|\n|$)/i.exec(rawError);
+    const coreMsg = coreMatch?.groups?.core?.trim() ?? 'SQL compilation error';
+    if (lowerError.includes('__lake_action') || lowerError.includes('__lake')) {
+        return [
+            `Lakehouse SQL execution failed: ${coreMsg}`,
+            '\nThis is an internal engine error with this access point — it may not be fully supported yet.',
+            '\nTry querying a different access point, or contact the data product team.',
+        ].join('');
+    }
+    if (lowerError.includes('invalid identifier') ||
+        lowerError.includes('does not exist')) {
+        const svcCols = services.map((s) => `${s.title}: ${s.columns.map((c) => c.name).join(', ')}`);
+        return [
+            `Lakehouse SQL execution failed: ${coreMsg}`,
+            `\nAvailable columns:\n${svcCols.join('\n')}`,
+            '\nTry rephrasing your question so the AI can pick the correct columns.',
+        ].join('');
+    }
+    return `Lakehouse SQL execution failed: ${coreMsg}`;
+}
 export function buildExecutionErrorMessage(errStr, services) {
     const errParts = [];
     const errLower = errStr.toLowerCase();
@@ -251,6 +355,23 @@ export function buildExecutionErrorMessage(errStr, services) {
         }
         return errParts.join('');
     }
+    if (errLower.includes('rename(~')) {
+        return [
+            'Cross-access-point JOINs on columns with the same name are not yet supported by the execution engine.',
+            '\nTry querying each access point separately, or ask a metadata question instead.',
+        ].join('');
+    }
+    if (/can't find a match for function '\w+\(Timestamp/i.test(errStr) &&
+        errLower.includes('string')) {
+        return [
+            'Date comparison type mismatch: a Timestamp column was compared with a plain string literal.',
+            '\nPlease rephrase your question so the AI can regenerate the query with the correct date format.',
+        ].join('');
+    }
+    // Snowflake SQL compilation errors — extract the core message from the trace
+    if (errLower.includes('snowflakesqlexception: sql compilation error')) {
+        return buildSnowflakeSqlError(errStr, services);
+    }
     errParts.push(errStr.slice(0, MAX_ERROR_MESSAGE_LENGTH));
     if (errLower.includes('column') &&
         (errLower.includes('not found') ||
@@ -281,11 +402,54 @@ function parseSuggestedQueries(rawAnswer) {
         .slice(0, 3);
     return { answer, suggestedQueries };
 }
+const COVERAGE_STOPWORDS = new Set([
+    'show',
+    'give',
+    'join',
+    'combine',
+    'merge',
+    'columns',
+    'column',
+    'results',
+    'result',
+    'data',
+    'services',
+    'service',
+    'access',
+    'points',
+    'point',
+    'table',
+    'tables',
+    'query',
+    'rows',
+    'records',
+    'values',
+    'list',
+    'display',
+    'fetch',
+    'retrieve',
+    'return',
+    'from',
+    'with',
+    'that',
+    'this',
+    'what',
+    'which',
+    'where',
+    'have',
+    'each',
+    'both',
+    'first',
+    'last',
+    'some',
+    'available',
+    'possible',
+]);
 function tokenizeQuestionForCoverage(question) {
-    return question
-        .toLowerCase()
-        .split(/[^a-z0-9]+/)
-        .filter((token) => token.length >= 4);
+    return tokenizeText(question, {
+        minLength: 4,
+        stopwords: COVERAGE_STOPWORDS,
+    });
 }
 function buildQuestionCoverageNote(question, columns) {
     const questionTokens = tokenizeQuestionForCoverage(question);
@@ -309,13 +473,44 @@ function formatPreviewValue(value) {
     }
     return `${raw.slice(0, ANALYSIS_PREVIEW_VALUE_LIMIT)}...`;
 }
-function buildDeterministicResultSummary(question, _query, columns, rows) {
+function buildJoinDiagnosticNote(query, columns, rows) {
+    if (!/\bJOIN\b/i.test(query) || rows.length === 0 || columns.length < 2) {
+        return '';
+    }
+    const allNullColumns = [];
+    const someDataColumns = [];
+    for (const col of columns) {
+        const allNull = rows.every((row) => {
+            if (isPlainObject(row)) {
+                const val = row[col];
+                return val === null || val === undefined;
+            }
+            return true;
+        });
+        if (allNull) {
+            allNullColumns.push(col);
+        }
+        else {
+            someDataColumns.push(col);
+        }
+    }
+    if (allNullColumns.length > 0 &&
+        someDataColumns.length > 0 &&
+        allNullColumns.length >= 2) {
+        return (`Note: ${allNullColumns.length} columns (${allNullColumns.slice(0, 3).join(', ')}${allNullColumns.length > 3 ? ', ...' : ''}) ` +
+            'returned entirely NULL values. This typically means the joined services do not share ' +
+            'overlapping values for the join key. The services may track different records. ' +
+            'Try querying each service separately to explore their data independently.');
+    }
+    return '';
+}
+function buildDeterministicResultSummary(question, query, columns, rows) {
     const selectedColumns = columns.length === 0 ? 'none' : columns.join(', ');
     const rowCount = rows.length;
     const previewRows = rows
         .slice(0, ANALYSIS_PREVIEW_ROW_LIMIT)
         .map((row, index) => {
-        if (row && typeof row === 'object' && !Array.isArray(row)) {
+        if (isPlainObject(row)) {
             const entries = Object.entries(row).slice(0, 4);
             const formattedEntries = entries.map(([key, value]) => `${key}: ${formatPreviewValue(value)}`);
             return `${index + 1}. ${formattedEntries.join(', ')}`;
@@ -324,11 +519,15 @@ function buildDeterministicResultSummary(question, _query, columns, rows) {
     })
         .join('\n');
     const coverageNote = buildQuestionCoverageNote(question, columns);
+    const joinDiagnostic = buildJoinDiagnosticNote(query, columns, rows);
     const parts = [
         `I retrieved ${rowCount} row${rowCount === 1 ? '' : 's'} for your question using this query.`,
         `Columns returned: ${selectedColumns}.`,
         `Sample rows:\n${previewRows || 'No sample rows available.'}`,
     ];
+    if (joinDiagnostic) {
+        parts.push(joinDiagnostic);
+    }
     if (coverageNote) {
         parts.push(coverageNote);
     }
@@ -363,10 +562,10 @@ export function attachMetadataOverview(setMessages, metadataOverview) {
         };
     });
 }
-export async function handleMetadataQuestion(question, metadata, context, startTime, hasQueryableServices) {
+export async function handleMetadataQuestion(question, metadata, context, startTime, hasQueryableServices, services, modelContextEnrichment) {
     const { config, plugin, history, setMessages } = context;
     addThinkingStep(setMessages, 'Answering from product metadata...');
-    const metadataPromptText = plugin.buildMetadataPrompt(question, metadata, history);
+    const metadataPromptText = plugin.buildMetadataPrompt(question, metadata, history, services, modelContextEnrichment);
     const rawAnswer = await plugin.callLLM(metadataPromptText, config);
     const { answer, suggestedQueries: parsedSuggestions } = parseSuggestedQueries(rawAnswer);
     const suggestedQueries = hasQueryableServices === false && !config.orchestratorUrl
@@ -433,10 +632,10 @@ async function runJudgeLoop(generatedSql, buildJudgePromptFn, context) {
     }
     return null;
 }
-export async function generateAndJudgeSql(question, services, coordinates, context, startTime, metadata) {
+export async function generateAndJudgeSql(question, services, coordinates, context, startTime, metadata, modelContextEnrichment) {
     const { plugin, config, history, setMessages } = context;
     addThinkingStep(setMessages, 'Building context from service schemas...');
-    const prompt = plugin.buildGeneratorPrompt(question, services, coordinates, history, metadata);
+    const prompt = plugin.buildGeneratorPrompt(question, services, coordinates, history, metadata, modelContextEnrichment);
     addThinkingStep(setMessages, 'Generating SQL query...');
     const answerText = await plugin.callLLM(prompt, config);
     const { sql: generatedSql, failure, suggestion, } = plugin.extractSqlFromResponse(answerText);
@@ -457,10 +656,11 @@ export async function generateAndJudgeSql(question, services, coordinates, conte
  * Uses AP-specific prompts that focus on `p()` syntax and omit
  * coordinates, parameters, and service()-related rules.
  */
-export async function generateAndJudgeAccessPointSql(question, accessPoints, context, startTime) {
+export async function generateAndJudgeAccessPointSql(question, accessPoints, context, startTime, modelContextEnrichment) {
     const { plugin, config, history, setMessages } = context;
     addThinkingStep(setMessages, 'Building context from access point schemas...');
-    const prompt = plugin.buildAccessPointGeneratorPrompt(question, accessPoints, history);
+    plugin.preWarmSchemaAnalysis(accessPoints, config);
+    const prompt = plugin.buildAccessPointGeneratorPrompt(question, accessPoints, history, modelContextEnrichment);
     addThinkingStep(setMessages, 'Generating SQL query for access points...');
     const answerText = await plugin.callLLM(prompt, config);
     const { sql: generatedSql, failure, suggestion, } = plugin.extractSqlFromResponse(answerText);
@@ -474,7 +674,7 @@ export async function generateAndJudgeAccessPointSql(question, accessPoints, con
         finishWithThinkingError(setMessages, 'Could not extract SQL from LLM response.\nTry rephrasing your question or ask about a specific access point.', startTime, LegendAIErrorType.GENERATION);
         return null;
     }
-    return runJudgeLoop(generatedSql, (sql) => plugin.buildAccessPointJudgePrompt(sql, question, accessPoints, history), context);
+    return runJudgeLoop(generatedSql, (sql) => plugin.buildAccessPointJudgePrompt(sql, question, accessPoints, history, modelContextEnrichment), context);
 }
 function reportExecutionResult(rawResult, setMessages, execStartTime, startTime) {
     const columns = deduplicateColumns(rawResult.columns);
@@ -494,10 +694,7 @@ function reportExecutionResult(rawResult, setMessages, execStartTime, startTime)
 export async function executeSqlAndReport(sql, services, config, plugin, setMessages, startTime, dataProductCoordinates) {
     const execStartTime = Date.now();
     try {
-        const isAccessPoint = services.some((s) => s.sourceType === TDSServiceSourceType.ACCESS_POINT);
-        const rawResult = isAccessPoint && dataProductCoordinates
-            ? await plugin.executeLakehouseSql(sql, dataProductCoordinates, config)
-            : await plugin.executeSql(sql, config);
+        const rawResult = await executeSqlForServices(sql, services, dataProductCoordinates, plugin, config);
         return reportExecutionResult(rawResult, setMessages, execStartTime, startTime);
     }
     catch (executeError) {
@@ -514,11 +711,28 @@ export async function executeSqlAndReport(sql, services, config, plugin, setMess
         return undefined;
     }
 }
+const HAS_PURE_AGGREGATION_PATTERN = /->groupBy\(|->distinct\(|->count\(\)|->sum\(\)|->average\(\)|->olapGroupBy\(/;
+const HAS_PURE_TAKE_PATTERN = /->take\(\s*\d+\s*\)/;
+const PURE_LIMIT_PATTERN = /->limit\(\s*(?<count>\d+)\s*\)/g;
+export function ensurePureSafetyLimit(pureQuery, limit = DEFAULT_SAFETY_LIMIT) {
+    const normalized = pureQuery.replaceAll(PURE_LIMIT_PATTERN, (_match, n) => `->take(${n})`);
+    if (HAS_PURE_TAKE_PATTERN.test(normalized)) {
+        return normalized;
+    }
+    if (HAS_PURE_AGGREGATION_PATTERN.test(normalized)) {
+        return normalized;
+    }
+    return `${normalized.trimEnd()}->take(${limit})`;
+}
 export async function executePureQueryAndReport(pureQuery, pureExecutionContext, dataProductCoordinates, config, plugin, setMessages, startTime) {
     const execStartTime = Date.now();
+    const safeQuery = ensurePureSafetyLimit(pureQuery);
     try {
         addThinkingStep(setMessages, 'Executing Pure query...');
-        const rawResult = await plugin.executePureQuery(pureQuery, pureExecutionContext, dataProductCoordinates, config);
+        const rawResult = await withTimeout(plugin.executePureQuery(safeQuery, pureExecutionContext, dataProductCoordinates, config), EXECUTION_TIMEOUT_MS);
+        if (!rawResult) {
+            throw new Error('Query execution timed out after 5 minutes. The dataset may be too large — try adding more filters or reducing the result set.');
+        }
         return reportExecutionResult(rawResult, setMessages, execStartTime, startTime);
     }
     catch (executeError) {
@@ -536,7 +750,7 @@ export async function executePureQueryAndReport(pureQuery, pureExecutionContext,
                 : execErrorType,
             thinkingDuration: elapsedSeconds(startTime),
         }));
-        return { columns: [], rows: [] };
+        return undefined;
     }
 }
 export async function analyzeOrchestratorResults(question, query, execResult, metadata, context, startTime) {
@@ -571,43 +785,66 @@ export async function analyzeOrchestratorResults(question, query, execResult, me
         }));
     }
 }
-async function handleEmptyOrchestratorResults(question, legendQuery, orchestratorOptions, metadata, resolvedEntities, context, startTime) {
+async function retryWithAlternateRoot(question, alternateRoot, resolvedEntities, orchestratorOptions, metadata, context, options) {
+    const { startTime, modelContext } = options;
     const { dataProductCoordinates, pureExecutionContext } = orchestratorOptions;
     const { config, plugin, setMessages } = context;
+    addThinkingStep(setMessages, `No results with ${extractElementNameFromPath(resolvedEntities.rootEntity)}, retrying with ${extractElementNameFromPath(alternateRoot)}...`);
+    try {
+        const retryRelated = resolvedEntities.relatedEntities.filter((e) => e !== alternateRoot);
+        const retryEnriched = modelContext
+            ? buildEnrichedBusinessContext(question, alternateRoot, retryRelated, modelContext)
+            : undefined;
+        const retryResponse = await withTimeout(plugin.generateQueryViaOrchestrator({
+            user_question: question,
+            semantic_search_resolution_details: {
+                data_product_coordinates: dataProductCoordinates,
+                root_entity: alternateRoot,
+                related_entities: retryRelated,
+                ...(retryEnriched
+                    ? { enriched_business_context: retryEnriched }
+                    : {}),
+            },
+        }, config), ORCHESTRATOR_GENERATION_TIMEOUT_MS);
+        if (!retryResponse) {
+            addThinkingStep(setMessages, 'Retry query generation timed out');
+            return false;
+        }
+        const retrySafeQuery = ensurePureSafetyLimit(retryResponse.legend_query);
+        updateLastAssistant(setMessages, () => ({
+            sql: retrySafeQuery,
+            sqlGenTime: elapsedSeconds(startTime, 2),
+            isExecuting: true,
+        }));
+        const retryResult = await executePureQueryAndReport(retrySafeQuery, pureExecutionContext, dataProductCoordinates, config, plugin, setMessages, startTime);
+        if (retryResult && retryResult.rows.length > 0) {
+            await analyzeOrchestratorResults(question, retryResponse.legend_query, retryResult, metadata, context, startTime);
+            return true;
+        }
+        return false;
+    }
+    catch (retryError) {
+        assertErrorThrown(retryError);
+        addThinkingStep(setMessages, `Retry with alternate entity failed: ${retryError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
+        return false;
+    }
+}
+async function handleEmptyOrchestratorResults(question, legendQuery, orchestratorOptions, metadata, resolvedEntities, context, options) {
+    const { startTime, modelContext } = options;
+    const { config, plugin, setMessages } = context;
     if (resolvedEntities.relatedEntities.length > 0) {
-        const alternateRoot = resolvedEntities.relatedEntities[0];
+        const alternateRoot = modelContext
+            ? findBestAlternateRoot(resolvedEntities.rootEntity, resolvedEntities.relatedEntities, modelContext)
+            : resolvedEntities.relatedEntities[0];
         if (alternateRoot) {
-            addThinkingStep(setMessages, `No results with ${extractElementNameFromPath(resolvedEntities.rootEntity)}, retrying with ${extractElementNameFromPath(alternateRoot)}...`);
-            try {
-                const retryResponse = await plugin.generateQueryViaOrchestrator({
-                    user_question: question,
-                    semantic_search_resolution_details: {
-                        data_product_coordinates: dataProductCoordinates,
-                        root_entity: alternateRoot,
-                        related_entities: resolvedEntities.relatedEntities.slice(1),
-                    },
-                }, config);
-                updateLastAssistant(setMessages, () => ({
-                    sql: retryResponse.legend_query,
-                    sqlGenTime: elapsedSeconds(startTime, 2),
-                    isExecuting: true,
-                }));
-                const retryResult = await executePureQueryAndReport(retryResponse.legend_query, pureExecutionContext, dataProductCoordinates, config, plugin, setMessages, startTime);
-                if (retryResult.rows.length > 0) {
-                    await analyzeOrchestratorResults(question, retryResponse.legend_query, retryResult, metadata, context, startTime);
-                    return;
-                }
-            }
-            catch (retryError) {
-                assertErrorThrown(retryError);
-                addThinkingStep(setMessages, `Retry with alternate entity failed: ${retryError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
+            const succeeded = await retryWithAlternateRoot(question, alternateRoot, resolvedEntities, orchestratorOptions, metadata, context, { startTime, ...(modelContext === undefined ? {} : { modelContext }) });
+            if (succeeded) {
+                return;
             }
         }
     }
     addThinkingStep(setMessages, 'No results returned \u2014 building contextual guidance...');
-    updateLastAssistant(setMessages, () => ({
-        isProcessing: true,
-    }));
+    updateLastAssistant(setMessages, () => ({ isProcessing: true }));
     const fallback = await withTimeout(plugin.buildNoResultsFallback(question, legendQuery, metadata, config), ANALYSIS_TIMEOUT_MS);
     if (fallback) {
         completeThinkingSteps(setMessages);
@@ -619,37 +856,88 @@ async function handleEmptyOrchestratorResults(question, legendQuery, orchestrato
         }));
     }
 }
-export async function processQuestionViaOrchestrator(question, dataProductCoordinates, metadata, context, pureExecutionContext, preResolvedEntities) {
+async function resolveOrchestrationEntities(question, dataProductCoordinates, context, pureExecutionContext, preResolvedEntities, modelContext) {
     const { config, plugin, setMessages } = context;
-    const startTime = Date.now();
-    try {
-        let resolvedEntities;
-        if (preResolvedEntities) {
-            resolvedEntities = preResolvedEntities;
-            addThinkingStep(setMessages, `Using pre-resolved root entity: ${extractElementNameFromPath(resolvedEntities.rootEntity)}`);
+    if (preResolvedEntities) {
+        addThinkingStep(setMessages, `Using pre-resolved root entity: ${extractElementNameFromPath(preResolvedEntities.rootEntity)}`);
+        return preResolvedEntities;
+    }
+    addThinkingStep(setMessages, 'Resolving entities for your query...');
+    const resolvedEntities = await plugin.resolveEntitiesForQuery(question, dataProductCoordinates, config, pureExecutionContext, modelContext);
+    addThinkingStep(setMessages, `Found root entity: ${extractElementNameFromPath(resolvedEntities.rootEntity)}`);
+    return resolvedEntities;
+}
+function extractUserErrorFromMessage(errorMessage) {
+    const marker = 'with error:';
+    const idx = errorMessage.toLowerCase().indexOf(marker);
+    if (idx !== -1) {
+        const after = errorMessage.slice(idx + marker.length).trimStart();
+        const newlineIdx = after.indexOf('\n');
+        const reason = (newlineIdx === -1 ? after : after.slice(0, newlineIdx)).trim();
+        if (reason.length > 0) {
+            return reason.slice(0, MAX_ERROR_MESSAGE_LENGTH);
         }
-        else {
-            addThinkingStep(setMessages, 'Resolving entities for your query...');
-            resolvedEntities = await plugin.resolveEntitiesForQuery(question, dataProductCoordinates, config, pureExecutionContext);
-            addThinkingStep(setMessages, `Found root entity: ${extractElementNameFromPath(resolvedEntities.rootEntity)}`);
+    }
+    return errorMessage.slice(0, MAX_ERROR_MESSAGE_LENGTH);
+}
+async function handleOrchestratorError(error, startTime, question, metadata, context) {
+    const { config, plugin, setMessages } = context;
+    const orchErrorType = classifyError(error);
+    const genTime = elapsedSeconds(startTime, 2);
+    addThinkingStep(setMessages, `Error: ${error.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
+    try {
+        addThinkingStep(setMessages, 'Building guidance from available metadata...');
+        const fallbackText = await withTimeout(plugin.buildFailureFallback(question, error.message, metadata, config), ANALYSIS_TIMEOUT_MS);
+        if (fallbackText) {
+            completeThinkingSteps(setMessages);
+            updateLastAssistant(setMessages, () => ({
+                dataContext: fallbackText,
+                sqlGenTime: genTime,
+                error: extractUserErrorFromMessage(error.message),
+                errorType: LegendAIErrorType.GENERATION,
+                isProcessing: false,
+                thinkingDuration: elapsedSeconds(startTime),
+            }));
+            return;
         }
+    }
+    catch (fallbackError) {
+        assertErrorThrown(fallbackError);
+        addThinkingStep(setMessages, `Fallback guidance failed: ${fallbackError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
+    }
+    finishWithThinkingError(setMessages, error.message, startTime, orchErrorType);
+}
+export async function processQuestionViaOrchestrator(question, dataProductCoordinates, metadata, context, pureExecutionContext, preResolvedEntities, modelContext) {
+    const { config, plugin, setMessages } = context;
+    const startTime = Date.now();
+    try {
+        const resolvedEntities = await resolveOrchestrationEntities(question, dataProductCoordinates, context, pureExecutionContext, preResolvedEntities, modelContext);
         if (resolvedEntities.relatedEntities.length > 0) {
             addThinkingStep(setMessages, `Found ${resolvedEntities.relatedEntities.length} related entities`);
         }
         addThinkingStep(setMessages, 'Generating Legend query via orchestrator...');
-        const orchestratorResponse = await plugin.generateQueryViaOrchestrator({
+        const enrichedContext = modelContext
+            ? buildEnrichedBusinessContext(question, resolvedEntities.rootEntity, resolvedEntities.relatedEntities, modelContext)
+            : undefined;
+        const orchestratorResponse = await withTimeout(plugin.generateQueryViaOrchestrator({
             user_question: question,
             semantic_search_resolution_details: {
                 data_product_coordinates: dataProductCoordinates,
                 root_entity: resolvedEntities.rootEntity,
                 related_entities: resolvedEntities.relatedEntities,
+                ...(enrichedContext
+                    ? { enriched_business_context: enrichedContext }
+                    : {}),
             },
-        }, config);
-        const queryGenTime = elapsedSeconds(startTime, 2);
+        }, config), ORCHESTRATOR_GENERATION_TIMEOUT_MS);
+        if (!orchestratorResponse) {
+            throw new Error('Query generation timed out. The orchestrator took too long to respond. Try a simpler question.');
+        }
+        const safeQuery = ensurePureSafetyLimit(orchestratorResponse.legend_query);
         completeThinkingSteps(setMessages);
         updateLastAssistant(setMessages, () => ({
-            sql: orchestratorResponse.legend_query,
-            sqlGenTime: queryGenTime,
+            sql: safeQuery,
+            sqlGenTime: elapsedSeconds(startTime, 2),
             isExecuting: true,
             isProcessing: true,
         }));
@@ -663,13 +951,19 @@ export async function processQuestionViaOrchestrator(question, dataProductCoordi
             }));
             return;
         }
-        const execResult = await executePureQueryAndReport(orchestratorResponse.legend_query, pureExecutionContext, dataProductCoordinates, config, plugin, setMessages, startTime);
+        const execResult = await executePureQueryAndReport(safeQuery, pureExecutionContext, dataProductCoordinates, config, plugin, setMessages, startTime);
+        if (!execResult) {
+            return;
+        }
         try {
             if (execResult.rows.length > 0) {
-                await analyzeOrchestratorResults(question, orchestratorResponse.legend_query, execResult, metadata, context, startTime);
+                await analyzeOrchestratorResults(question, safeQuery, execResult, metadata, context, startTime);
             }
             else {
-                await handleEmptyOrchestratorResults(question, orchestratorResponse.legend_query, { dataProductCoordinates, pureExecutionContext }, metadata, resolvedEntities, context, startTime);
+                await handleEmptyOrchestratorResults(question, safeQuery, { dataProductCoordinates, pureExecutionContext }, metadata, resolvedEntities, context, {
+                    startTime,
+                    ...(modelContext === undefined ? {} : { modelContext }),
+                });
             }
         }
         catch (analysisError) {
@@ -686,26 +980,7 @@ export async function processQuestionViaOrchestrator(question, dataProductCoordi
     }
     catch (error) {
         assertErrorThrown(error);
-        const orchErrorType = classifyError(error);
-        addThinkingStep(setMessages, `Error: ${error.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
-        try {
-            addThinkingStep(setMessages, 'Building guidance from available metadata...');
-            const fallbackText = await withTimeout(plugin.buildFailureFallback(question, error.message, metadata, config), ANALYSIS_TIMEOUT_MS);
-            if (fallbackText) {
-                completeThinkingSteps(setMessages);
-                updateLastAssistant(setMessages, () => ({
-                    dataContext: fallbackText,
-                    isProcessing: false,
-                    thinkingDuration: elapsedSeconds(startTime),
-                }));
-                return;
-            }
-        }
-        catch (fallbackError) {
-            assertErrorThrown(fallbackError);
-            addThinkingStep(setMessages, `Fallback guidance failed: ${fallbackError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
-        }
-        finishWithThinkingError(setMessages, error.message, startTime, orchErrorType);
+        await handleOrchestratorError(error, startTime, question, metadata, context);
     }
 }
 export function cleanLlmSqlResponse(raw) {
@@ -732,9 +1007,7 @@ export function sanitizeJoinOrderBy(sql) {
     const beforeOrderBy = parts[0] ?? '';
     const afterOrderBy = parts.slice(1).join('ORDER BY').replace(/^\s+/, '');
     const selectAliases = new Map();
-    const aliasRegex = /\b(?<tbl>[a-z]\w*)\s*\.\s*"(?<col>[^"]+)"\s+AS\s+(?:"(?<qAlias>[^"]+)"|(?<uAlias>\w+))/gi;
-    let m;
-    while ((m = aliasRegex.exec(beforeOrderBy)) !== null) {
+    for (const m of beforeOrderBy.matchAll(SELECT_ALIAS_PATTERN)) {
         const tableAlias = (m.groups?.tbl ?? '').toLowerCase();
         const colName = (m.groups?.col ?? '').toLowerCase();
         const asAlias = m.groups?.qAlias ?? m.groups?.uAlias ?? '';
@@ -743,17 +1016,72 @@ export function sanitizeJoinOrderBy(sql) {
     if (selectAliases.size === 0) {
         return sql;
     }
-    const rewritten = afterOrderBy.replaceAll(ALIAS_DOT_COL_PATTERN, (...args) => {
-        const groups = args[args.length - 1];
-        const key = `${groups.tbl.toLowerCase()}.${groups.col.toLowerCase()}`;
-        const alias = selectAliases.get(key);
-        return alias ? `"${alias}"` : String(args[0]);
-    });
+    // Build a lookup of the original alias.col text → its SELECT alias replacement.
+    // matchAll exposes typed `.groups`, so no cast is needed.
+    const replacements = new Map();
+    for (const m of afterOrderBy.matchAll(ALIAS_DOT_COL_PATTERN)) {
+        const tbl = (m.groups?.tbl ?? '').toLowerCase();
+        const col = (m.groups?.col ?? '').toLowerCase();
+        const alias = selectAliases.get(`${tbl}.${col}`);
+        if (alias) {
+            replacements.set(m[0], `"${alias}"`);
+        }
+    }
+    if (replacements.size === 0) {
+        return sql;
+    }
+    const rewritten = afterOrderBy.replaceAll(ALIAS_DOT_COL_PATTERN, (match) => replacements.get(match) ?? match);
     if (rewritten === afterOrderBy) {
         return sql;
     }
     return `${beforeOrderBy}ORDER BY ${rewritten}`;
 }
+export function sanitizeJoinSameKeyColumns(sql, services) {
+    if (!JOIN_PATTERN.test(sql)) {
+        return sql;
+    }
+    const joinRegex = /\bJOIN\s{1,5}p\(\s{0,5}'(?<pId>[^']{1,200})'\s{0,5}\)\s{1,5}AS\s{1,5}(?<rAlias>[a-z]\w{0,63})\s{1,5}ON\s{1,5}(?<onClause>[^\n]{1,500})/gi;
+    let result = sql;
+    let match;
+    while ((match = joinRegex.exec(sql)) !== null) {
+        const pId = match.groups?.pId ?? '';
+        const rightAlias = match.groups?.rAlias ?? '';
+        const onClause = (match.groups?.onClause ?? '').slice(0, 500);
+        const sameKeyMatch = /(?<lAlias>[a-z]\w{0,63}) {0,5}\. {0,5}"(?<lCol>[^"]{1,200})" {0,5}= {0,5}(?<rAl>[a-z]\w{0,63}) {0,5}\. {0,5}"(?<rCol>[^"]{1,200})"/i.exec(onClause);
+        if (!sameKeyMatch) {
+            continue;
+        }
+        const leftCol = sameKeyMatch.groups?.lCol ?? '';
+        const rightCol = sameKeyMatch.groups?.rCol ?? '';
+        const matchedRightAlias = sameKeyMatch.groups?.rAl ?? '';
+        if (leftCol.toLowerCase() !== rightCol.toLowerCase() ||
+            matchedRightAlias.toLowerCase() !== rightAlias.toLowerCase()) {
+            continue;
+        }
+        const renamedKey = `${rightAlias}_${rightCol}`;
+        const columnList = buildSubqueryColumnList(pId, rightCol, renamedKey, services);
+        const originalFragment = `p('${pId}') AS ${rightAlias}`;
+        const subqueryFragment = `(SELECT ${columnList} FROM p('${pId}')) AS ${rightAlias}`;
+        const oldOnRef = `${rightAlias}."${rightCol}"`;
+        const newOnRef = `${rightAlias}."${renamedKey}"`;
+        result = result.replace(originalFragment, subqueryFragment);
+        result = result.replaceAll(oldOnRef, newOnRef);
+    }
+    return result;
+}
+function buildSubqueryColumnList(pId, originalKey, renamedKey, services) {
+    const svc = services?.find((s) => s.dataProductPath &&
+        s.pattern &&
+        pId === `${s.dataProductPath}.${s.pattern.replace(/^\//, '')}`);
+    if (!svc || svc.columns.length === 0) {
+        return `"${originalKey}" AS "${renamedKey}", *`;
+    }
+    return svc.columns
+        .map((c) => c.name.toLowerCase() === originalKey.toLowerCase()
+        ? `"${c.name}" AS "${renamedKey}"`
+        : `"${c.name}"`)
+        .join(', ');
+}
 export function sanitizeLiteralColumns(sql) {
     if (!UNION_ALL_PATTERN.test(sql)) {
         return sql;
@@ -817,8 +1145,7 @@ export function ensureDateParameters(sql, services) {
         return sql;
     }
     const today = getTodayISO();
-    const serviceCallPattern = /\bservice\s*\([^()]*\)/gi;
-    return sql.replaceAll(serviceCallPattern, (match) => {
+    return sql.replaceAll(SERVICE_CALL_PATTERN, (match) => {
         let patched = match;
         for (const param of dateParams) {
             if (new RegExp(String.raw `\b${param}\s*=>`, 'i').test(patched)) {
@@ -897,6 +1224,84 @@ export function buildMissingParamsWarning(missingParams) {
         `Please provide ${missingParams.length > 1 ? 'values' : 'a value'} in your question, e.g. "show data where ${exampleValues}".`,
     ].join('\n');
 }
+const NESTED_P_IN_CLAUSE = /\bIN\s*\(\s*SELECT\b[^)]*\bFROM\s+p\s*\(/i;
+const NESTED_P_CROSS_JOIN = /\bCROSS\s+JOIN\s*\(\s*SELECT\b[^)]*\bFROM\s+p\s*\(/i;
+const NESTED_P_SCALAR = /\(\s*SELECT\s+(?:COUNT|SUM|AVG|MIN|MAX)\s*\([^)]*\)\s+(?:AS\s+\w+\s+)?FROM\s+p\s*\(/i;
+export function hasNestedPCalls(sql) {
+    return (NESTED_P_IN_CLAUSE.test(sql) ||
+        NESTED_P_CROSS_JOIN.test(sql) ||
+        NESTED_P_SCALAR.test(sql));
+}
+const NESTED_AGGREGATE_PATTERN = /\b(?:SUM|AVG|MIN|MAX|COUNT)\s*\(\s*(?:SUM|AVG|MIN|MAX|COUNT)\s*\(/i;
+const OVER_CLAUSE_PATTERN = /\bOVER\s*\(/i;
+const AGGREGATE_IN_WINDOW_ARGS_PATTERN = /\bOVER\s*\((?:[^()]|\([^()]*\)){0,300}\b(?:SUM|AVG|MIN|MAX|COUNT)\s*\(/i;
+const NON_WINDOW_SCALAR_FUNCTIONS = new Set([
+    'ROUND',
+    'CAST',
+    'COALESCE',
+    'CEIL',
+    'CEILING',
+    'FLOOR',
+    'ABS',
+    'CONCAT',
+    'NULLIF',
+    'IFNULL',
+    'GREATEST',
+    'LEAST',
+    'TRUNC',
+    'TRUNCATE',
+    'MOD',
+    'POWER',
+    'SQRT',
+    'LN',
+    'LOG',
+    'EXP',
+    'SUBSTR',
+    'SUBSTRING',
+    'TO_CHAR',
+    'TO_NUMBER',
+    'TO_DATE',
+]);
+const FUNCTION_CALL_WITH_OVER_PATTERN = /\b(?<fn>[A-Z_]{1,32})\s*\((?:[^()]|\([^()]*\)){0,200}\bOVER\s*\(/gi;
+function hasWindowInsideNonWindowFunctionCall(sql) {
+    FUNCTION_CALL_WITH_OVER_PATTERN.lastIndex = 0;
+    let match;
+    while ((match = FUNCTION_CALL_WITH_OVER_PATTERN.exec(sql)) !== null) {
+        const fn = match.groups?.fn?.toUpperCase();
+        if (fn && NON_WINDOW_SCALAR_FUNCTIONS.has(fn)) {
+            return true;
+        }
+    }
+    return false;
+}
+const LATERAL_SUBQUERY_PATTERN = /\b(?:CROSS|INNER|LEFT|RIGHT)?\s*JOIN\s+LATERAL\b|\bLATERAL\s*\(/i;
+export function detectUnsupportedEnginePattern(sql) {
+    if (NESTED_AGGREGATE_PATTERN.test(sql) && OVER_CLAUSE_PATTERN.test(sql)) {
+        return {
+            kind: 'NESTED_AGGREGATE_IN_WINDOW',
+            hint: "The engine's SQL→Pure translator cannot combine an aggregate function with a window function in the same SELECT (e.g. SUM(COUNT(*)) OVER ()). Move the aggregation into a CTE first, then apply the window over the CTE.",
+        };
+    }
+    if (AGGREGATE_IN_WINDOW_ARGS_PATTERN.test(sql)) {
+        return {
+            kind: 'AGGREGATE_IN_WINDOW_ARGS',
+            hint: "The engine's SQL→Pure translator cannot evaluate an aggregate function inside an OVER clause's PARTITION BY or ORDER BY at the same level as GROUP BY (e.g. ROW_NUMBER() OVER (ORDER BY COUNT(*) DESC)). Materialize the aggregate into a CTE column first, then reference that column from the window — for example: SELECT col, cnt, ROW_NUMBER() OVER (ORDER BY cnt DESC) AS rn FROM (SELECT col, COUNT(*) AS cnt FROM p('…') GROUP BY col) agg.",
+        };
+    }
+    if (hasWindowInsideNonWindowFunctionCall(sql)) {
+        return {
+            kind: 'WINDOW_INSIDE_FUNCTION_CALL',
+            hint: "The engine's SQL→Pure translator cannot evaluate a window function inside another function call (e.g. ROUND(SUM(x) OVER (), 2)). Materialize the window expression into its own column in a subquery or CTE, then apply the wrapping function in the outer SELECT — for example: SELECT col, ROUND(cnt * 100.0 / total, 2) FROM (SELECT col, cnt, SUM(cnt) OVER () AS total FROM (SELECT col, COUNT(*) AS cnt FROM p('…') GROUP BY col) agg) windowed.",
+        };
+    }
+    if (LATERAL_SUBQUERY_PATTERN.test(sql)) {
+        return {
+            kind: 'LATERAL_SUBQUERY',
+            hint: "The engine's SQL parser does not support LATERAL subqueries (CROSS JOIN LATERAL (...)). Rewrite using ROW_NUMBER() OVER (PARTITION BY ...) in a CTE and filter on the rank, or use a correlated subquery.",
+        };
+    }
+    return undefined;
+}
 /**
  * Appends a safety LIMIT to queries that lack one, preventing unbounded
  * result sets on large services. Skips aggregation queries since those
@@ -908,8 +1313,23 @@ export function ensureSafeLimit(sql, limit = DEFAULT_SAFETY_LIMIT) {
     }
     return `${sql.trimEnd()}\nLIMIT ${limit}`;
 }
+function prepareSafeSql(sql, services) {
+    const safeSql = ensureSafeLimit(ensureDateParameters(sanitizeLiteralColumns(sanitizeJoinSameKeyColumns(sanitizeJoinOrderBy(sql), services)), services));
+    const unsupported = detectUnsupportedEnginePattern(safeSql);
+    if (unsupported) {
+        throw new LegendAIUnsupportedEngineShapeError(unsupported.hint);
+    }
+    return safeSql;
+}
+/**
+ * The single execution chokepoint for SQL in this module. Every code
+ * path that needs to run SQL — primary orchestrator, retry-with-fix,
+ * exported `executeSqlAndReport`, internal helpers — calls through
+ * here. No other code in this file may call `plugin.executeLakehouseSql`
+ * or `plugin.executeSql` directly.
+ */
 async function executeSqlForServices(sql, services, dataProductCoordinates, plugin, config) {
-    const safeSql = ensureSafeLimit(ensureDateParameters(sanitizeLiteralColumns(sanitizeJoinOrderBy(sql)), services));
+    const safeSql = prepareSafeSql(sql, services);
     const isAccessPoint = services.some((s) => s.sourceType === TDSServiceSourceType.ACCESS_POINT);
     if (isAccessPoint && dataProductCoordinates) {
         return plugin.executeLakehouseSql(safeSql, dataProductCoordinates, config);
@@ -1106,61 +1526,162 @@ async function reportQueryResults(report, metadata, context, startTime, hasOrche
         }));
     }
 }
+export function supplementMissingCoverage(question, selected, allServices, maxTotal = 4) {
+    if (selected.length >= maxTotal || allServices.length <= selected.length) {
+        return selected;
+    }
+    const questionTokens = tokenizeText(question);
+    if (questionTokens.length === 0) {
+        return selected;
+    }
+    const coveredTokens = new Set();
+    for (const svc of selected) {
+        for (const col of svc.columns) {
+            for (const token of splitIdentifierTokens(col.name)) {
+                coveredTokens.add(token);
+            }
+            coveredTokens.add(col.name.toLowerCase());
+        }
+        for (const token of splitIdentifierTokens(svc.title)) {
+            coveredTokens.add(token);
+        }
+    }
+    const coveredArray = Array.from(coveredTokens);
+    const uncovered = questionTokens.filter((t) => !coveredTokens.has(t) &&
+        !(t.length >= 4 && coveredArray.some((ct) => isFuzzyMatch(t, ct))));
+    if (uncovered.length === 0) {
+        return selected;
+    }
+    const selectedSet = new Set(selected.map((s) => s.pattern));
+    const supplementCandidates = allServices
+        .filter((svc) => !selectedSet.has(svc.pattern))
+        .map((svc) => {
+        const svcTokens = new Set();
+        for (const col of svc.columns) {
+            svcTokens.add(col.name.toLowerCase());
+            for (const token of splitIdentifierTokens(col.name)) {
+                svcTokens.add(token);
+            }
+        }
+        for (const token of splitIdentifierTokens(svc.title)) {
+            svcTokens.add(token);
+        }
+        const coverCount = uncovered.filter((t) => svcTokens.has(t)).length;
+        const isGeneric = GENERIC_TABLE_PATTERNS.test(svc.title) ? 1 : 0;
+        return { svc, coverCount, isGeneric };
+    })
+        .filter((c) => c.coverCount > 0)
+        .sort((a, b) => b.coverCount - a.coverCount || a.isGeneric - b.isGeneric);
+    const result = [...selected];
+    for (const candidate of supplementCandidates) {
+        if (result.length >= maxTotal) {
+            break;
+        }
+        result.push(candidate.svc);
+    }
+    return result;
+}
+// ─── Question Normalization ──────────────────────────────────────────────────
+/** Delegates to the plugin's LLM to fix natural-language typos in the user's question before AP selection; falls back to the original on failure. */
+export async function normalizeQuestion(question, plugin, config) {
+    try {
+        return await plugin.normalizeQuestion(question, config);
+    }
+    catch (error) {
+        assertErrorThrown(error);
+        return question;
+    }
+}
+// ─── Multi-Turn AP Bias ──────────────────────────────────────────────────────
 /**
- * Scores each service against the user question by counting keyword
- * overlap between the question tokens and the service title, description,
- * column names, and parameter names. Returns services sorted by
- * descending relevance score.
+ * Extracts AP patterns that were successfully used in previous conversation
+ * turns. When the user asks a follow-up question, these APs should be
+ * preferred since they're likely still relevant.
  */
-export function preFilterServicesByRelevance(question, services, limit) {
-    const queryTokens = question
-        .toLowerCase()
-        .split(/[\s,.:;!?'"()\-/]+/)
-        .filter((t) => t.length > 2);
-    if (queryTokens.length === 0) {
-        return services.slice(0, limit);
-    }
-    const scored = services.map((svc) => {
-        const haystack = [
-            svc.title,
-            svc.description ?? '',
-            svc.pattern,
-            ...svc.columns.slice(0, 30).map((c) => c.name),
-            ...svc.parameters,
-            ...(svc.preFilters ?? []).flatMap((pf) => {
-                const parts = [pf.property, pf.operator];
-                if (pf.value !== undefined) {
-                    parts.push(String(pf.value));
+function extractPreviousTurnAPPatterns(history) {
+    const patterns = new Set();
+    for (const turn of history) {
+        const pCallPattern = /p\(\s*'(?<pId>[^']+)'\s*\)/g;
+        let pCall;
+        while ((pCall = pCallPattern.exec(turn.sql)) !== null) {
+            const pattern = pCall.groups?.pId;
+            if (pattern) {
+                const lastDot = pattern.lastIndexOf('.');
+                if (lastDot >= 0) {
+                    patterns.add(pattern.slice(lastDot + 1).toLowerCase());
                 }
-                return parts;
-            }),
-        ]
-            .join(' ')
-            .toLowerCase();
-        let score = 0;
-        for (const token of queryTokens) {
-            if (haystack.includes(token)) {
-                score += 1;
+                patterns.add(pattern.toLowerCase());
             }
         }
-        return { svc, score };
+    }
+    return patterns;
+}
+/**
+ * Applies a scoring boost to services whose patterns appeared in previous
+ * conversation turns, biasing the pre-filter toward continuity when the
+ * user asks follow-up questions about the same data.
+ */
+export function applyMultiTurnBias(services, history) {
+    if (history.length === 0) {
+        return services;
+    }
+    const previousAPs = extractPreviousTurnAPPatterns(history);
+    if (previousAPs.size === 0) {
+        return services;
+    }
+    // Sort so that previously-used APs come first, preserving relative order
+    const biased = [...services].sort((a, b) => {
+        const aUsed = previousAPs.has(a.pattern.replace(/^\//, '').toLowerCase()) ||
+            previousAPs.has(a.title.toLowerCase())
+            ? 1
+            : 0;
+        const bUsed = previousAPs.has(b.pattern.replace(/^\//, '').toLowerCase()) ||
+            previousAPs.has(b.title.toLowerCase())
+            ? 1
+            : 0;
+        return bUsed - aUsed;
     });
-    scored.sort((a, b) => b.score - a.score);
-    return scored.slice(0, limit).map((s) => s.svc);
+    return biased;
 }
 async function selectBestServices(question, services, context) {
-    const { plugin, config, setMessages } = context;
+    const { plugin, config, setMessages, history } = context;
     if (services.length <= 1) {
         return services;
     }
-    let candidates = services;
-    if (services.length > MAX_SERVICES_FOR_LLM_SELECTION) {
-        candidates = preFilterServicesByRelevance(question, services, MAX_SERVICES_FOR_LLM_SELECTION);
-        addThinkingStep(setMessages, `Pre-filtered ${services.length} services to ${candidates.length} candidates`);
+    // ── Step 0: Normalize the question (expand abbreviations, fix typos) ──
+    addThinkingStep(setMessages, 'Normalizing question...');
+    const normalizedQuestion = await normalizeQuestion(question, plugin, config);
+    if (normalizedQuestion !== question) {
+        addThinkingStep(setMessages, `Normalized: "${normalizedQuestion}"`);
+    }
+    // ── Step 0b: Apply multi-turn bias if this is a follow-up ──
+    let inputServices = services;
+    if (history.length > 0) {
+        inputServices = applyMultiTurnBias(services, history);
+    }
+    let candidates = inputServices;
+    if (inputServices.length > MAX_SERVICES_FOR_LLM_SELECTION) {
+        candidates = preFilterServicesByRelevance(normalizedQuestion, inputServices, MAX_SERVICES_FOR_LLM_SELECTION);
+        const topNames = candidates
+            .slice(0, 5)
+            .map((s) => `${s.title} (${s.columns.length} cols)`)
+            .join(', ');
+        addThinkingStep(setMessages, `Pre-filtered ${services.length} services to ${candidates.length} candidates. Top: ${topNames}`);
     }
     try {
         addThinkingStep(setMessages, 'Selecting best service for your query...');
-        return await plugin.selectRelevantServices(question, candidates, config);
+        const selected = await plugin.selectRelevantServices(normalizedQuestion, candidates, config);
+        const selectedNames = selected.map((s) => s.title).join(', ');
+        addThinkingStep(setMessages, `LLM selected: ${selectedNames}`);
+        const validated = supplementMissingCoverage(normalizedQuestion, selected, candidates);
+        if (validated.length > selected.length) {
+            const supplemented = validated
+                .slice(selected.length)
+                .map((s) => s.title)
+                .join(', ');
+            addThinkingStep(setMessages, `Supplemented with ${supplemented} for column coverage`);
+        }
+        return validated;
     }
     catch (selectionError) {
         assertErrorThrown(selectionError);
@@ -1197,64 +1718,131 @@ async function tryRecoverZeroRows(currentSql, sqlResult, question, selectedServi
     }
     return { sql: recoveredSql, result: recoveredResult };
 }
+async function resolveNestedPCalls(sql, question, selectedAPs, context, modelContextEnrichment) {
+    if (!hasNestedPCalls(sql)) {
+        return sql;
+    }
+    const { plugin, config, setMessages } = context;
+    addThinkingStep(setMessages, 'Detected nested p() call — rewriting query...');
+    const fixPrompt = plugin.buildAccessPointErrorCorrectionPrompt(sql, question, 'The SQL contains p() inside a subquery (IN clause, CROSS JOIN, or scalar subquery). p() can ONLY appear directly after FROM or JOIN.', selectedAPs, context.history, modelContextEnrichment);
+    try {
+        const fixResponse = await plugin.callLLM(fixPrompt, config);
+        const fixResult = plugin.extractJudgeResult(fixResponse);
+        const corrected = fixResult.correctedSql;
+        if (corrected && !hasNestedPCalls(corrected)) {
+            return corrected;
+        }
+    }
+    catch (fixError) {
+        assertErrorThrown(fixError);
+        addThinkingStep(setMessages, `Could not rewrite query: ${fixError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}. Proceeding with original SQL.`);
+    }
+    return sql;
+}
+async function retryAccessPointSqlFix(options, context, metadata) {
+    const { failedSql, question, errMsg, selectedAPs, dataProductCoordinates, modelContextEnrichment, startTime, } = options;
+    const { plugin, config, setMessages } = context;
+    addThinkingStep(setMessages, 'SQL error detected — asking LLM to fix the query on the same access points...');
+    const fixPrompt = plugin.buildAccessPointErrorCorrectionPrompt(failedSql, question, errMsg.slice(0, MAX_ERROR_MESSAGE_LENGTH), selectedAPs, context.history, modelContextEnrichment);
+    try {
+        const fixResponse = await plugin.callLLM(fixPrompt, config);
+        const fixResult = plugin.extractJudgeResult(fixResponse);
+        const correctedSql = fixResult.correctedSql;
+        if (correctedSql) {
+            addThinkingStep(setMessages, 'Re-executing corrected SQL...');
+            updateLastAssistant(setMessages, () => ({
+                sql: correctedSql,
+                isExecuting: true,
+            }));
+            const retryResult = await executeSqlForServices(correctedSql, selectedAPs, dataProductCoordinates, plugin, config);
+            await reportQueryResults({
+                currentSql: correctedSql,
+                sqlResult: retryResult,
+                question,
+                services: selectedAPs,
+            }, metadata, context, startTime, false);
+            return true;
+        }
+    }
+    catch (retryError) {
+        assertErrorThrown(retryError);
+        addThinkingStep(setMessages, `SQL correction also failed: ${retryError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
+    }
+    return false;
+}
 /**
  * Dedicated query flow for data product access points.
  * Unlike processDataQuery this:
  * - Uses AP-specific generator/judge prompts (p() syntax, no coordinates/params)
  * - Skips parameter detection (APs have no parameters)
  */
-async function processAccessPointQuery(question, accessPoints, metadata, context, startTime, dataProductCoordinates) {
+async function processAccessPointQuery(question, accessPoints, metadata, context, startTime, dataProductCoordinates, modelContextEnrichment) {
     const { config, plugin, setMessages } = context;
-    const hasOrchestratorFallback = Boolean(config.orchestratorUrl);
     addThinkingStep(setMessages, 'Found relevant access points to query');
     const selectedAPs = await selectBestServices(question, accessPoints, context);
-    const judgedSql = await generateAndJudgeAccessPointSql(question, selectedAPs, context, startTime);
+    const judgedSql = await generateAndJudgeAccessPointSql(question, selectedAPs, context, startTime, modelContextEnrichment);
     if (!judgedSql) {
         addThinkingStep(setMessages, 'SQL generation could not produce a valid query.');
-        handleSqlGenerationFailure(setMessages, startTime, hasOrchestratorFallback, SQL_GENERATION_FAILURE_WITH_ORCHESTRATOR, SQL_GENERATION_FAILURE_NO_ORCHESTRATOR, LegendAIErrorType.GENERATION, buildFallbackSuggestions(selectedAPs));
+        handleSqlGenerationFailure(setMessages, startTime, false, SQL_GENERATION_FAILURE_WITH_ORCHESTRATOR, SQL_GENERATION_FAILURE_NO_ORCHESTRATOR, LegendAIErrorType.GENERATION, buildFallbackSuggestions(selectedAPs));
         return;
     }
     const sqlGenTimeValue = elapsedSeconds(startTime, 2);
     completeThinkingSteps(setMessages);
+    const finalSql = await resolveNestedPCalls(judgedSql, question, selectedAPs, context, modelContextEnrichment);
     updateLastAssistant(setMessages, () => ({
-        sql: judgedSql,
+        sql: finalSql,
         sqlGenTime: sqlGenTimeValue,
         isExecuting: true,
     }));
     const execStartTime = Date.now();
     try {
-        const safeSql = ensureSafeLimit(judgedSql);
-        const rawResult = await plugin.executeLakehouseSql(safeSql, dataProductCoordinates, config);
+        const rawResult = await executeSqlForServices(finalSql, selectedAPs, dataProductCoordinates, plugin, config);
         await reportQueryResults({
-            currentSql: judgedSql,
+            currentSql: finalSql,
             sqlResult: rawResult,
             question,
             services: selectedAPs,
-        }, metadata, context, startTime, hasOrchestratorFallback);
+        }, metadata, context, startTime, false);
     }
     catch (executeError) {
         assertErrorThrown(executeError);
         const execErrorType = classifyError(executeError);
-        addThinkingStep(setMessages, `Execution failed: ${executeError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
-        finishWithThinkingError(setMessages, buildExecutionErrorMessage(executeError.message, selectedAPs), startTime, execErrorType === LegendAIErrorType.GENERAL
+        const errMsg = executeError.message;
+        const errorCategory = categorizeExecutionError(errMsg, executeError);
+        addThinkingStep(setMessages, `Execution failed: ${errMsg.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
+        if (errorCategory === ExecutionErrorCategory.SQL_FIXABLE) {
+            const retried = await retryAccessPointSqlFix({
+                failedSql: finalSql,
+                question,
+                errMsg,
+                selectedAPs,
+                dataProductCoordinates,
+                ...(modelContextEnrichment === undefined
+                    ? {}
+                    : { modelContextEnrichment }),
+                startTime,
+            }, context, metadata);
+            if (retried) {
+                return;
+            }
+        }
+        finishWithThinkingError(setMessages, buildExecutionErrorMessage(errMsg, selectedAPs), startTime, execErrorType === LegendAIErrorType.GENERAL
             ? LegendAIErrorType.EXECUTION
             : execErrorType);
+        const queriedGroups = [
+            ...new Set(selectedAPs
+                .map((ap) => ap.accessPointGroupTitle)
+                .filter((t) => t !== undefined)),
+        ];
         updateLastAssistant(setMessages, () => ({
             execTime: elapsedSeconds(execStartTime, 2),
             isExecuting: false,
             suggestedQueries: buildFallbackSuggestions(selectedAPs),
-            ...(hasOrchestratorFallback
-                ? {
-                    fallbackAction: {
-                        label: ORCHESTRATOR_FALLBACK_LABEL,
-                        actionId: LEGEND_AI_ORCHESTRATOR_FALLBACK_ACTION_ID,
-                    },
-                }
-                : {}),
+            queriedAccessPointGroups: queriedGroups,
         }));
     }
 }
-async function processDataQuery(question, services, coordinates, metadata, context, startTime, orchestratorOptions) {
+async function processDataQuery(question, services, coordinates, metadata, context, startTime, orchestratorOptions, modelContextEnrichment) {
     const { config, setMessages } = context;
     const dataProductCoordinates = orchestratorOptions?.dataProductCoordinates;
     const hasOrchestratorFallback = Boolean(config.orchestratorUrl && dataProductCoordinates);
@@ -1264,7 +1852,7 @@ async function processDataQuery(question, services, coordinates, metadata, conte
     }
     addThinkingStep(setMessages, 'Found relevant services to query');
     const selectedServices = await selectBestServices(question, services, context);
-    const judgedSql = await generateAndJudgeSql(question, selectedServices, coordinates, context, startTime, metadata);
+    const judgedSql = await generateAndJudgeSql(question, selectedServices, coordinates, context, startTime, metadata, modelContextEnrichment);
     if (!judgedSql) {
         addThinkingStep(setMessages, 'SQL generation could not produce a valid query.');
         handleSqlGenerationFailure(setMessages, startTime, hasOrchestratorFallback, SQL_GENERATION_FAILURE_WITH_ORCHESTRATOR, SQL_GENERATION_FAILURE_NO_ORCHESTRATOR, LegendAIErrorType.GENERATION, buildFallbackSuggestions(selectedServices));
@@ -1338,45 +1926,24 @@ function splitServicesByType(services) {
         accessPoints: services.filter((s) => s.sourceType === TDSServiceSourceType.ACCESS_POINT),
     };
 }
-async function routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOpts) {
+async function routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOpts, modelContextEnrichment) {
     const dataProductCoordinates = orchestratorOpts?.dataProductCoordinates;
     const { tdsServices, accessPoints } = splitServicesByType(services);
     if (tdsServices.length === 0 &&
         accessPoints.length > 0 &&
         dataProductCoordinates) {
-        await processAccessPointQuery(question, accessPoints, metadata, context, startTime, dataProductCoordinates);
+        await processAccessPointQuery(question, accessPoints, metadata, context, startTime, dataProductCoordinates, modelContextEnrichment);
     }
     else {
-        await processDataQuery(question, tdsServices, coordinates, metadata, context, startTime, orchestratorOpts);
+        await processDataQuery(question, tdsServices, coordinates, metadata, context, startTime, orchestratorOpts, modelContextEnrichment);
     }
 }
-async function handleAmbiguousIntent(question, services, coordinates, metadata, context, startTime, orchestratorOpts) {
-    const { setMessages } = context;
-    addThinkingStep(setMessages, 'Intent is ambiguous, providing metadata context and querying data...');
-    let metadataOverview;
-    try {
-        addThinkingStep(setMessages, 'Building metadata context...');
-        metadataOverview = await buildMetadataOverview(question, metadata, context);
-    }
-    catch (metadataError) {
-        assertErrorThrown(metadataError);
-        addThinkingStep(setMessages, `Metadata context failed: ${metadataError.message.slice(0, MAX_THINKING_ERROR_PREVIEW_LENGTH)}`);
-    }
-    try {
-        await routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOpts);
-        if (metadataOverview) {
-            attachMetadataOverview(setMessages, metadataOverview);
-        }
-    }
-    catch (queryError) {
-        assertErrorThrown(queryError);
-        addThinkingStep(setMessages, 'Query failed, answering from product metadata...');
-        await handleMetadataQuestion(question, metadata, context, startTime, true);
-    }
-}
-export async function processQuestion(question, services, coordinates, metadata, context, dataProductCoordinates, pureExecutionContext) {
+export async function processQuestion(question, services, coordinates, metadata, context, dataProductCoordinates, pureExecutionContext, modelContext) {
     const { config, plugin, setMessages } = context;
     const startTime = Date.now();
+    const modelContextEnrichment = modelContext
+        ? buildModelContextEnrichmentText(modelContext, services)
+        : undefined;
     try {
         addThinkingStep(setMessages, 'Analyzing your question...');
         const orchestratorOpts = dataProductCoordinates
@@ -1388,7 +1955,7 @@ export async function processQuestion(question, services, coordinates, metadata,
             }
             : undefined;
         if (services.length === 0) {
-            await handleMetadataQuestion(question, metadata, context, startTime, false);
+            await handleMetadataQuestion(question, metadata, context, startTime, false, services, modelContextEnrichment);
             if (config.orchestratorUrl && dataProductCoordinates) {
                 updateLastAssistant(setMessages, () => ({
                     fallbackAction: {
@@ -1399,29 +1966,19 @@ export async function processQuestion(question, services, coordinates, metadata,
             }
             return;
         }
-        const fastIntent = classifyQuestionIntentFast(question, true);
-        if (fastIntent.intent === LegendAIQuestionIntent.METADATA &&
-            !fastIntent.ambiguous) {
-            await handleMetadataQuestion(question, metadata, context, startTime, true);
-            return;
-        }
-        if (fastIntent.ambiguous) {
-            await handleAmbiguousIntent(question, services, coordinates, metadata, context, startTime, orchestratorOpts);
-            return;
-        }
         const serviceNames = services.map((s) => s.title);
         const intent = await plugin.classifyQuestionIntent(question, true, config, serviceNames);
         if (intent === LegendAIQuestionIntent.METADATA) {
-            await handleMetadataQuestion(question, metadata, context, startTime, true);
+            await handleMetadataQuestion(question, metadata, context, startTime, true, services, modelContextEnrichment);
             return;
         }
         try {
-            await routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOpts);
+            await routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOpts, modelContextEnrichment);
         }
         catch (sqlError) {
             assertErrorThrown(sqlError);
-            addThinkingStep(setMessages, 'SQL generation failed, answering from product metadata...');
-            await handleMetadataQuestion(question, metadata, context, startTime, true);
+            addThinkingStep(setMessages, 'Query failed, answering from product metadata...');
+            await handleMetadataQuestion(question, metadata, context, startTime, true, services, modelContextEnrichment);
             appendFallbackSuggestions(setMessages, services);
         }
     }
@@ -1432,12 +1989,12 @@ export async function processQuestion(question, services, coordinates, metadata,
         appendFallbackSuggestions(setMessages, services);
     }
 }
-async function routeDataQueryWithErrorHandling(question, services, coordinates, metadata, context, orchestratorOptions) {
+async function routeDataQueryWithErrorHandling(question, services, coordinates, metadata, context, orchestratorOptions, modelContextEnrichment) {
     const { setMessages } = context;
     const startTime = Date.now();
     try {
         addThinkingStep(setMessages, 'Preparing data query...');
-        await routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOptions);
+        await routeToAccessPointOrData(question, services, coordinates, metadata, context, startTime, orchestratorOptions, modelContextEnrichment);
     }
     catch (error) {
         assertErrorThrown(error);
@@ -1445,13 +2002,16 @@ async function routeDataQueryWithErrorHandling(question, services, coordinates,
         finishWithThinkingError(setMessages, error.message, startTime, classifyError(error));
     }
 }
-export async function processQuestionWithIntent(question, intent, services, coordinates, metadata, context, orchestratorOptions) {
+export async function processQuestionWithIntent(question, intent, services, coordinates, metadata, context, orchestratorOptions, modelContext) {
     const { config, setMessages } = context;
     const dataProductCoordinates = orchestratorOptions?.dataProductCoordinates;
+    const modelContextEnrichment = modelContext
+        ? buildModelContextEnrichmentText(modelContext, services)
+        : undefined;
     if (intent === LegendAIQuestionIntent.METADATA) {
         const startTime = Date.now();
         try {
-            await handleMetadataQuestion(question, metadata, context, startTime, services.length > 0);
+            await handleMetadataQuestion(question, metadata, context, startTime, services.length > 0, services, modelContextEnrichment);
         }
         catch (error) {
             assertErrorThrown(error);
@@ -1464,11 +2024,11 @@ export async function processQuestionWithIntent(question, intent, services, coor
         config.orchestratorUrl &&
         dataProductCoordinates) {
         if (services.length > 0) {
-            await routeDataQueryWithErrorHandling(question, services, coordinates, metadata, context, orchestratorOptions);
+            await routeDataQueryWithErrorHandling(question, services, coordinates, metadata, context, orchestratorOptions, modelContextEnrichment);
             return;
         }
         const startTime = Date.now();
-        await handleMetadataQuestion(question, metadata, context, startTime, false);
+        await handleMetadataQuestion(question, metadata, context, startTime, false, services, modelContextEnrichment);
         updateLastAssistant(setMessages, () => ({
             fallbackAction: {
                 label: ORCHESTRATOR_FALLBACK_LABEL,
@@ -1477,6 +2037,6 @@ export async function processQuestionWithIntent(question, intent, services, coor
         }));
         return;
     }
-    await routeDataQueryWithErrorHandling(question, services, coordinates, metadata, context, orchestratorOptions);
+    await routeDataQueryWithErrorHandling(question, services, coordinates, metadata, context, orchestratorOptions, modelContextEnrichment);
 }
 //# sourceMappingURL=LegendAIChatProcessors.js.map