@kaelio/ktx 0.13.1 → 0.15.0

package/dist/context/scan/local-scan.js

@@ -5,10 +5,10 @@ import { KtxScanEmbeddingPortAdapter } from '../../context/llm/embedding-port.js
 import { ktxLocalStateDbPath } from '../project/local-state-db.js';
 import { redactKtxScanReport } from './credentials.js';
 import { resolveEnabledTables } from './enabled-tables.js';
-import { completedKtxScanEnrichmentStateSummary } from './enrichment-state.js';
+import { completedKtxScanEnrichmentStateSummary, } from './enrichment-state.js';
 import { failedKtxScanEnrichmentSummary, ktxScanErrorMessage } from './enrichment-summary.js';
 import { createDeterministicLocalScanEnrichmentProviders, runLocalScanEnrichment, } from './local-enrichment.js';
-import { writeLocalScanEnrichmentArtifacts, writeLocalScanManifestShards } from './local-enrichment-artifacts.js';
+import { createKtxScanDescriptionResumeStore, loadOnDiskDescriptionUpdates, writeLocalScanEnrichmentArtifacts, writeLocalScanEnrichmentCheckpoint, writeLocalScanManifestShards, } from './local-enrichment-artifacts.js';
 import { readLocalScanStructuralSnapshot } from './local-structural-artifacts.js';
 import { SqliteLocalScanEnrichmentStateStore } from './sqlite-local-enrichment-state-store.js';
 function enrichmentResolutionWarning(status) {
@@ -47,10 +47,11 @@ function normalizeDriver(driver) {
         normalized === 'clickhouse' ||
         normalized === 'sqlserver' ||
         normalized === 'bigquery' ||
-        normalized === 'snowflake') {
+        normalized === 'snowflake' ||
+        normalized === 'mongodb') {
         return normalized;
     }
-    throw new Error(`Standalone ktx scan supports postgres/sqlite/mysql/clickhouse/sqlserver/bigquery/snowflake in this phase, received "${driver ?? 'unknown'}"`);
+    throw new Error(`Standalone ktx scan supports postgres/sqlite/mysql/clickhouse/sqlserver/bigquery/snowflake/mongodb in this phase, received "${driver ?? 'unknown'}"`);
 }
 function tablePathCount(paths) {
     return paths.filter((path) => path.startsWith('tables/') && path.endsWith('.json')).length;
@@ -116,15 +117,17 @@ function createLocalScanEnrichmentStateStore(options) {
     }
     return new SqliteLocalScanEnrichmentStateStore({ dbPath: ktxLocalStateDbPath(options.project) });
 }
-function localScanProviderIdentity(config, llmConfig, relationships) {
+function localScanLlmIdentity(llmConfig) {
     return {
-        mode: config.mode,
-        embeddingDimensions: config.embeddings?.dimensions ?? null,
-        llmModel: llmConfig.models.default ?? null,
-        embeddingModel: config.embeddings?.model ?? null,
-        batchSize: config.embeddings?.batchSize ?? null,
+        model: llmConfig.models.default ?? null,
         baseUrlConfigured: Boolean(llmConfig.provider.gateway?.base_url),
-        relationships,
+    };
+}
+function localScanEmbeddingIdentity(config) {
+    return {
+        model: config.embeddings?.model ?? null,
+        dimensions: config.embeddings?.dimensions ?? null,
+        batchSize: config.embeddings?.batchSize ?? null,
     };
 }
 function reportFromIngest(input) {
@@ -299,6 +302,13 @@ export async function runLocalScan(options) {
         const enrichmentStateStore = connector ? createLocalScanEnrichmentStateStore(options) : null;
         let enrichmentState = completedKtxScanEnrichmentStateSummary();
         let enrichmentSnapshot = null;
+        // On a `--stages` subset run, the structural manifest write below (and the
+        // later enrichment write) merge with on-disk shards, but the merge treats ai/db
+        // descriptions as scan-managed and overwrites them with whatever the run emits.
+        // A subset that skips `descriptions` emits none, so without this the structural
+        // write would delete the prior descriptions before enrichment can preserve them.
+        // Capture them up front (only for subset runs) and feed them to both writes.
+        let priorDescriptionUpdates = null;
         if (!reusedExistingScanArtifacts && !report.dryRun && report.artifactPaths.rawSourcesDir) {
             await options.progress?.update(0.7, 'Writing schema artifacts');
             const rawSnapshot = await readLocalScanStructuralSnapshot({
@@ -312,12 +322,16 @@ export async function runLocalScan(options) {
             if (rawSnapshot.warnings?.length) {
                 report.warnings.push(...rawSnapshot.warnings);
             }
+            if (options.stages !== undefined && connector) {
+                priorDescriptionUpdates = await loadOnDiskDescriptionUpdates(options.project, options.connectionId, rawSnapshot);
+            }
             const manifestArtifacts = await writeLocalScanManifestShards({
                 project: options.project,
                 connectionId: options.connectionId,
                 syncId: record.syncId,
                 driver,
                 snapshot: rawSnapshot,
+                ...(priorDescriptionUpdates ? { descriptionUpdates: priorDescriptionUpdates } : {}),
                 dryRun: false,
             });
             report.artifactPaths.manifestShards = manifestArtifacts.manifestShards;
@@ -332,15 +346,42 @@ export async function runLocalScan(options) {
                     connectionId: options.connectionId,
                     mode,
                     detectRelationships: options.detectRelationships,
+                    ...(options.stages ? { stages: options.stages } : {}),
                     connector,
                     ...(enrichmentSnapshot ? { snapshot: enrichmentSnapshot } : {}),
-                    context: { runId: record.runId, progress: options.progress?.startPhase(0.18) },
+                    context: {
+                        runId: record.runId,
+                        ...(options.signal ? { signal: options.signal } : {}),
+                        ...(options.progress ? { progress: options.progress.startPhase(0.18) } : {}),
+                    },
                     providers: enrichmentProviders,
                     stateStore: enrichmentStateStore,
+                    descriptionResumeStore: options.dryRun
+                        ? null
+                        : createKtxScanDescriptionResumeStore({
+                            project: options.project,
+                            connectionId: options.connectionId,
+                            syncId: record.syncId,
+                            driver,
+                        }),
                     syncId: record.syncId,
-                    providerIdentity: localScanProviderIdentity(options.project.config.scan.enrichment, options.project.config.llm, options.project.config.scan.relationships),
+                    loadPriorDescriptions: (enrichedSnapshot) => priorDescriptionUpdates
+                        ? Promise.resolve(priorDescriptionUpdates)
+                        : loadOnDiskDescriptionUpdates(options.project, options.connectionId, enrichedSnapshot),
+                    llmIdentity: localScanLlmIdentity(options.project.config.llm),
+                    embeddingIdentity: localScanEmbeddingIdentity(options.project.config.scan.enrichment),
                     relationshipSettings: options.project.config.scan.relationships,
                     now: options.now,
+                    onCheckpoint: async (checkpoint) => {
+                        await writeLocalScanEnrichmentCheckpoint({
+                            project: options.project,
+                            connectionId: options.connectionId,
+                            syncId: record.syncId,
+                            driver,
+                            enrichment: checkpoint,
+                            dryRun: options.dryRun ?? false,
+                        });
+                    },
                 });
                 const artifacts = await writeLocalScanEnrichmentArtifacts({
                     project: options.project,

package/dist/context/scan/local-structural-artifacts.d.ts

@@ -1,5 +1,5 @@
 import type { KtxLocalProject } from '../../context/project/project.js';
-import type { KtxConnectionDriver, KtxSchemaSnapshot } from './types.js';
+import type { KtxConnectionDriver, KtxScanWarning, KtxSchemaSnapshot } from './types.js';
 export interface ReadLocalScanStructuralSnapshotInput {
     project: KtxLocalProject;
     connectionId: string;
@@ -7,4 +7,6 @@ export interface ReadLocalScanStructuralSnapshotInput {
     rawSourcesDir: string;
     extractedAtFallback: string;
 }
+/** @internal */
+export declare function isKtxScanWarningCode(code: string): code is KtxScanWarning['code'];
 export declare function readLocalScanStructuralSnapshot(input: ReadLocalScanStructuralSnapshotInput): Promise<KtxSchemaSnapshot>;

package/dist/context/scan/local-structural-artifacts.js

@@ -18,7 +18,12 @@ const scanWarningCodes = new Set([
     'enrichment_failed',
     'description_fallback_used',
     'constraint_discovery_unauthorized',
+    'object_introspection_failed',
 ]);
+/** @internal */
+export function isKtxScanWarningCode(code) {
+    return scanWarningCodes.has(code);
+}
 function parseWarning(rawWarning, path) {
     if (!isRecord(rawWarning) ||
         typeof rawWarning.code !== 'string' ||

package/dist/context/scan/object-introspection.d.ts

@@ -0,0 +1,21 @@
+import type { KtxScanWarning } from './types.js';
+export interface IntrospectObjectContext {
+    /** Bare object name (table or view). */
+    object: string;
+    catalog?: string | null;
+    db?: string | null;
+}
+export type IntrospectObjectOutcome<T> = {
+    ok: true;
+    table: T;
+} | {
+    ok: false;
+    warning: KtxScanWarning;
+};
+/**
+ * Runs a single-object metadata/profiling read and isolates its failure: a
+ * broken or inaccessible object becomes a recoverable warning instead of
+ * aborting the whole scan. Native programming faults (a ktx bug, not a broken
+ * object) still propagate so they are not masked as object skips.
+ */
+export declare function tryIntrospectObject<T>(ctx: IntrospectObjectContext, fn: () => T | Promise<T>): Promise<IntrospectObjectOutcome<T>>;

package/dist/context/scan/object-introspection.js

@@ -0,0 +1,35 @@
+import { isNativeProgrammingFault } from '../../errors.js';
+function objectLabel(ctx) {
+    return [ctx.catalog, ctx.db, ctx.object].filter((part) => Boolean(part)).join('.');
+}
+function objectIntrospectionWarning(ctx, error) {
+    const reason = error instanceof Error ? error.message : String(error);
+    return {
+        code: 'object_introspection_failed',
+        message: reason,
+        table: ctx.object,
+        recoverable: true,
+        metadata: {
+            object: objectLabel(ctx),
+            ...(ctx.db ? { db: ctx.db } : {}),
+            ...(ctx.catalog ? { catalog: ctx.catalog } : {}),
+        },
+    };
+}
+/**
+ * Runs a single-object metadata/profiling read and isolates its failure: a
+ * broken or inaccessible object becomes a recoverable warning instead of
+ * aborting the whole scan. Native programming faults (a ktx bug, not a broken
+ * object) still propagate so they are not masked as object skips.
+ */
+export async function tryIntrospectObject(ctx, fn) {
+    try {
+        return { ok: true, table: await fn() };
+    }
+    catch (error) {
+        if (isNativeProgrammingFault(error)) {
+            throw error;
+        }
+        return { ok: false, warning: objectIntrospectionWarning(ctx, error) };
+    }
+}

package/dist/context/scan/relationship-benchmarks.js

@@ -6,7 +6,7 @@ import { gunzipSync } from 'node:zlib';
 import Database from 'better-sqlite3';
 import YAML from 'yaml';
 import { z } from 'zod';
-import { getDialectForDriver } from '../connections/dialects.js';
+import { getSqlDialectForDriver } from '../connections/dialects.js';
 import { snapshotToKtxEnrichedSchema } from './local-enrichment.js';
 import { generateKtxRelationshipDiscoveryCandidates, mergeKtxRelationshipDiscoveryCandidates, } from './relationship-candidates.js';
 import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
@@ -319,7 +319,7 @@ export function ktxRelationshipBenchmarkDetectorWithLlm(llmRuntime) {
             const formalLinks = formalMetadata.accepted.map((relationship) => relationshipToBenchmarkLink(relationship));
             const acceptedKeys = new Set(formalLinks.map(fkKey));
             const sqliteDataAvailable = Boolean(input.dataPath && input.snapshot.driver === 'sqlite');
-            const dialect = getDialectForDriver(input.snapshot.driver);
+            const dialect = getSqlDialectForDriver(input.snapshot.driver);
             const profilingExecutor = sqliteDataAvailable && input.mode !== 'profiling_disabled'
                 ? new KtxRelationshipBenchmarkSqliteExecutor(input.dataPath)
                 : null;
@@ -332,6 +332,7 @@ export function ktxRelationshipBenchmarkDetectorWithLlm(llmRuntime) {
                 })
                 : await profileKtxRelationshipSchema({
                     connectionId: input.snapshot.connectionId,
+                    driver: input.snapshot.driver,
                     dialect,
                     schema: input.schema,
                     executor: profilingExecutor,
@@ -447,7 +448,7 @@ export function currentKtxRelationshipBenchmarkDetector() {
             const formalLinks = formalMetadata.accepted.map((relationship) => relationshipToBenchmarkLink(relationship));
             const acceptedKeys = new Set(formalLinks.map(fkKey));
             const sqliteDataAvailable = Boolean(input.dataPath && input.snapshot.driver === 'sqlite');
-            const dialect = getDialectForDriver(input.snapshot.driver);
+            const dialect = getSqlDialectForDriver(input.snapshot.driver);
             const profilingExecutor = sqliteDataAvailable && input.mode !== 'profiling_disabled'
                 ? new KtxRelationshipBenchmarkSqliteExecutor(input.dataPath)
                 : null;
@@ -460,6 +461,7 @@ export function currentKtxRelationshipBenchmarkDetector() {
                 })
                 : await profileKtxRelationshipSchema({
                     connectionId: input.snapshot.connectionId,
+                    driver: input.snapshot.driver,
                     dialect,
                     schema: input.schema,
                     executor: profilingExecutor,

package/dist/context/scan/relationship-composite-candidates.d.ts

@@ -1,7 +1,8 @@
-import type { KtxDialect } from '../connections/dialects.js';
+import type { KtxSqlDialect } from '../connections/dialects.js';
 import type { KtxEnrichedSchema, KtxRelationshipType } from './enrichment-types.js';
+import type { KtxRelationshipDetectionBudget } from './relationship-detection-budget.js';
 import { type KtxRelationshipProfileArtifact, type KtxRelationshipReadOnlyExecutor } from './relationship-profiling.js';
-import type { KtxScanContext, KtxTableRef } from './types.js';
+import type { KtxProgressPort, KtxScanContext, KtxTableRef } from './types.js';
 type KtxCompositeRelationshipStatus = 'accepted' | 'review' | 'rejected';
 interface KtxCompositeRelationshipTupleEndpoint {
     tableId: string;
@@ -47,7 +48,7 @@ export interface KtxCompositeRelationshipCandidate {
 }
 export interface DiscoverKtxCompositeRelationshipsInput {
     connectionId: string;
-    dialect: KtxDialect;
+    dialect: KtxSqlDialect;
     schema: KtxEnrichedSchema;
     profiles: KtxRelationshipProfileArtifact;
     executor: KtxRelationshipReadOnlyExecutor | null;
@@ -57,6 +58,8 @@ export interface DiscoverKtxCompositeRelationshipsInput {
     minPrimaryKeyUniqueness?: number;
     minSourceCoverage?: number;
     maxViolationRatio?: number;
+    budget?: KtxRelationshipDetectionBudget;
+    progress?: KtxProgressPort;
 }
 export interface DiscoverKtxCompositeRelationshipsResult {
     primaryKeys: KtxCompositePrimaryKeyCandidate[];

package/dist/context/scan/relationship-composite-candidates.js

@@ -341,7 +341,13 @@ export async function discoverKtxCompositeRelationships(input) {
     const tableByName = new Map(tables.map((table) => [table.ref.name, table]));
     const primaryKeys = [];
     let queryCount = 0;
-    for (const table of tables) {
+    for (const [index, table] of tables.entries()) {
+        if (input.budget?.check()) {
+            break;
+        }
+        await input.progress?.update((index + 1) / tables.length, `Probing composite keys ${index + 1}/${tables.length}`, {
+            transient: true,
+        });
         const result = await detectCompositePrimaryKeys({
             connectionId: input.connectionId,
             dialect: input.dialect,
@@ -358,6 +364,9 @@ export async function discoverKtxCompositeRelationships(input) {
     }
     const relationships = [];
     for (const targetKey of primaryKeys) {
+        if (input.budget?.check()) {
+            break;
+        }
         const targetTable = tableByName.get(targetKey.table.name);
         if (!targetTable) {
             continue;
@@ -371,6 +380,9 @@ export async function discoverKtxCompositeRelationships(input) {
             continue;
         }
         for (const sourceTable of tables) {
+            if (input.budget?.check()) {
+                break;
+            }
             if (sourceTable.id === targetTable.id) {
                 continue;
             }

package/dist/context/scan/relationship-detection-budget.d.ts

@@ -0,0 +1,35 @@
+export type KtxRelationshipDetectionStopReason = 'budget' | 'aborted';
+export interface KtxRelationshipDetectionBudget {
+    /**
+     * Returns a stop reason when the relationship stage must stop scheduling new
+     * work, else null. Calling it at a unit boundary records the first observed
+     * stop so the stage can be finalized as partial.
+     */
+    check(): KtxRelationshipDetectionStopReason | null;
+    /** The first stop reason observed via check(), or null if the stage ran to completion. */
+    stopReason(): KtxRelationshipDetectionStopReason | null;
+}
+export interface CreateKtxRelationshipDetectionBudgetInput {
+    budgetMs: number;
+    signal?: AbortSignal;
+    now?: () => number;
+}
+export declare function createKtxRelationshipDetectionBudget(input: CreateKtxRelationshipDetectionBudgetInput): KtxRelationshipDetectionBudget;
+export interface MapWithBudgetInput<TInput, TOutput> {
+    inputs: readonly TInput[];
+    concurrency: number;
+    budget?: KtxRelationshipDetectionBudget;
+    onStart?: (index: number, total: number, item: TInput) => Promise<void> | void;
+    mapOne: (item: TInput, index: number) => Promise<TOutput>;
+}
+export interface MapWithBudgetResult<TOutput> {
+    /** Output aligned with inputs; entries skipped on budget exhaustion are undefined. */
+    results: Array<TOutput | undefined>;
+    processedCount: number;
+}
+/**
+ * Concurrent map that stops claiming new items once the budget trips. In-flight
+ * items finish; pending items are left undefined. With no budget it is a plain
+ * bounded-concurrency map.
+ */
+export declare function mapWithBudget<TInput, TOutput>(input: MapWithBudgetInput<TInput, TOutput>): Promise<MapWithBudgetResult<TOutput>>;

package/dist/context/scan/relationship-detection-budget.js

@@ -0,0 +1,53 @@
+export function createKtxRelationshipDetectionBudget(input) {
+    const now = input.now ?? (() => Date.now());
+    const deadline = now() + Math.max(0, input.budgetMs);
+    let tripped = null;
+    return {
+        check() {
+            if (input.signal?.aborted) {
+                tripped = 'aborted';
+                return 'aborted';
+            }
+            if (now() >= deadline) {
+                tripped ??= 'budget';
+                return 'budget';
+            }
+            return null;
+        },
+        stopReason() {
+            return tripped;
+        },
+    };
+}
+/**
+ * Concurrent map that stops claiming new items once the budget trips. In-flight
+ * items finish; pending items are left undefined. With no budget it is a plain
+ * bounded-concurrency map.
+ */
+export async function mapWithBudget(input) {
+    const total = input.inputs.length;
+    const results = new Array(total);
+    const safeConcurrency = Math.max(1, Math.floor(input.concurrency));
+    let nextIndex = 0;
+    let processedCount = 0;
+    async function worker() {
+        while (true) {
+            const index = nextIndex;
+            if (index >= total) {
+                return;
+            }
+            // Check the budget only when work remains, so a deadline that elapses
+            // after the last item is claimed never marks a fully-processed stage partial.
+            if (input.budget?.check()) {
+                return;
+            }
+            nextIndex += 1;
+            const item = input.inputs[index];
+            await input.onStart?.(index, total, item);
+            results[index] = await input.mapOne(item, index);
+            processedCount += 1;
+        }
+    }
+    await Promise.all(Array.from({ length: Math.min(safeConcurrency, total) }, () => worker()));
+    return { results, processedCount };
+}

package/dist/context/scan/relationship-diagnostics.d.ts

@@ -63,6 +63,8 @@ export interface KtxRelationshipDiagnosticsArtifact {
     generatedAt: string;
     summary: KtxRelationshipDiagnosticsSummary;
     noAcceptedReason: string | null;
+    partial: boolean;
+    partialReason: string | null;
     candidateCountsBySource: Record<string, number>;
     validation: KtxRelationshipDiagnosticsValidation;
     thresholds: KtxRelationshipDiagnosticsThresholds;
@@ -83,6 +85,9 @@ export interface BuildKtxRelationshipDiagnosticsInput {
     warnings?: readonly KtxScanWarning[];
     thresholds?: Partial<KtxRelationshipDiagnosticsThresholds>;
     policy?: Partial<KtxRelationshipDiagnosticsPolicy>;
+    partial?: {
+        reason: string;
+    } | null;
     generatedAt?: string;
 }
 export interface EmptyKtxRelationshipProfileArtifactInput {

package/dist/context/scan/relationship-diagnostics.js

@@ -211,6 +211,8 @@ export function buildKtxRelationshipDiagnostics(input) {
         generatedAt: input.generatedAt ?? new Date().toISOString(),
         summary,
         noAcceptedReason: noAcceptedReason({ artifacts: input.artifacts, profile: input.profile }),
+        partial: Boolean(input.partial),
+        partialReason: input.partial?.reason ?? null,
         candidateCountsBySource: candidateCountsBySource(input.artifacts),
         validation: {
             available: input.profile.sqlAvailable,

package/dist/context/scan/relationship-discovery.d.ts

@@ -1,19 +1,22 @@
 import type { KtxLlmRuntimePort } from '../../context/llm/runtime-port.js';
-import type { KtxDialect } from '../connections/dialects.js';
+import type { KtxSqlDialect } from '../connections/dialects.js';
 import type { KtxScanRelationshipConfig } from '../project/config.js';
 import type { KtxEnrichedSchema, KtxRelationshipUpdate } from './enrichment-types.js';
 import { type KtxCompositeRelationshipCandidate } from './relationship-composite-candidates.js';
+import { type KtxRelationshipDetectionStopReason } from './relationship-detection-budget.js';
 import { type KtxResolvedRelationshipDiscoveryCandidate } from './relationship-graph-resolver.js';
 import { type KtxRelationshipProfileArtifact } from './relationship-profiling.js';
-import type { KtxScanConnector, KtxScanContext, KtxScanEnrichmentSummary, KtxScanRelationshipSummary, KtxScanWarning } from './types.js';
+import type { KtxProgressPort, KtxScanConnector, KtxScanContext, KtxScanEnrichmentSummary, KtxScanRelationshipSummary, KtxScanWarning } from './types.js';
 export interface DiscoverKtxRelationshipsInput {
     connectionId: string;
-    dialect: KtxDialect;
+    dialect: KtxSqlDialect | null;
     connector: KtxScanConnector;
     schema: KtxEnrichedSchema;
     context: KtxScanContext;
     settings: KtxScanRelationshipConfig;
     llmRuntime?: KtxLlmRuntimePort | null;
+    progress?: KtxProgressPort;
+    now?: () => number;
 }
 export interface DiscoverKtxRelationshipsResult {
     relationshipUpdate: KtxRelationshipUpdate;
@@ -24,5 +27,8 @@ export interface DiscoverKtxRelationshipsResult {
     statisticalValidation: KtxScanEnrichmentSummary['statisticalValidation'];
     llmRelationshipValidation: KtxScanEnrichmentSummary['llmRelationshipValidation'];
     warnings: KtxScanWarning[];
+    partial: {
+        reason: KtxRelationshipDetectionStopReason;
+    } | null;
 }
 export declare function discoverKtxRelationships(input: DiscoverKtxRelationshipsInput): Promise<DiscoverKtxRelationshipsResult>;

package/dist/context/scan/relationship-discovery.js

@@ -1,5 +1,6 @@
 import { generateKtxRelationshipDiscoveryCandidates, mergeKtxRelationshipDiscoveryCandidates, } from './relationship-candidates.js';
 import { discoverKtxCompositeRelationships, } from './relationship-composite-candidates.js';
+import { createKtxRelationshipDetectionBudget, } from './relationship-detection-budget.js';
 import { collectKtxFormalMetadataRelationships } from './relationship-formal-metadata.js';
 import { resolveKtxRelationshipGraph, } from './relationship-graph-resolver.js';
 import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
@@ -62,17 +63,20 @@ function compositeSummary(relationships) {
     };
 }
 async function detectCompositeRelationships(input) {
-    if (!input.executor || !input.profile.sqlAvailable) {
+    if (!input.executor || !input.profile.sqlAvailable || !input.dialect) {
         return [];
     }
+    const dialect = input.dialect;
     try {
         const compositeDetection = await discoverKtxCompositeRelationships({
             connectionId: input.connectionId,
-            dialect: input.dialect,
+            dialect,
             schema: input.schema,
             profiles: input.profile,
             executor: input.executor,
             ctx: input.context,
+            budget: input.budget,
+            ...(input.progress ? { progress: input.progress } : {}),
         });
         for (const warning of compositeDetection.warnings) {
             input.warnings.push({
@@ -139,10 +143,16 @@ function sqlExecutor(input) {
 }
 export async function discoverKtxRelationships(input) {
     const { executor, warnings } = sqlExecutor(input);
+    const budget = createKtxRelationshipDetectionBudget({
+        budgetMs: input.settings.detectionBudgetMs,
+        ...(input.context.signal ? { signal: input.context.signal } : {}),
+        ...(input.now ? { now: input.now } : {}),
+    });
     const formalMetadata = collectKtxFormalMetadataRelationships(input.schema);
     const profileCache = createKtxRelationshipProfileCache();
     const profile = await profileKtxRelationshipSchema({
         connectionId: input.connectionId,
+        driver: input.connector.driver,
         dialect: input.dialect,
         schema: input.schema,
         executor,
@@ -150,12 +160,17 @@ export async function discoverKtxRelationships(input) {
         profileSampleRows: input.settings.profileSampleRows,
         profileConcurrency: input.settings.profileConcurrency,
         cache: profileCache,
+        budget,
+        ...(input.progress ? { progress: input.progress } : {}),
     });
     const deterministicCandidates = generateKtxRelationshipDiscoveryCandidates(input.schema, {
         maxCandidatesPerColumn: input.settings.maxCandidatesPerColumn,
         profiles: profile,
     });
-    const llmProposalResult = input.settings.llmProposals
+    // The LLM proposal is one more unit of relationship work, so it honors the same
+    // budget/abort gate as profiling, validation, and composite probing — a stage
+    // that already exhausted its budget (or was aborted) must not start a fresh call.
+    const llmProposalResult = input.settings.llmProposals && !budget.check()
         ? await proposeKtxRelationshipCandidatesWithLlm({
             connectionId: input.connectionId,
             schema: input.schema,
@@ -186,6 +201,8 @@ export async function discoverKtxRelationships(input) {
             concurrency: input.settings.validationConcurrency,
             validationBudget: input.settings.validationBudget,
         },
+        budget,
+        ...(input.progress ? { progress: input.progress } : {}),
     });
     const graph = resolveKtxRelationshipGraph({
         schema: input.schema,
@@ -205,6 +222,8 @@ export async function discoverKtxRelationships(input) {
         executor,
         context: input.context,
         warnings,
+        budget,
+        ...(input.progress ? { progress: input.progress } : {}),
     });
     const inferredAccepted = nonFormalAcceptedRelationships({
         formalIds: formalMetadata.acceptedIds,
@@ -225,6 +244,7 @@ export async function discoverKtxRelationships(input) {
         resolvedRelationships: graph.relationships,
     });
     const compositeCounts = compositeSummary(compositeRelationships);
+    const stopReason = budget.stopReason();
     return {
         relationshipUpdate: {
             connectionId: input.connectionId,
@@ -241,8 +261,11 @@ export async function discoverKtxRelationships(input) {
         profile,
         resolvedRelationships: graph.relationships,
         compositeRelationships,
-        statisticalValidation: profile.sqlAvailable ? 'completed' : 'skipped',
+        // A budget/abort stop means profiling did not finish, so report it as not
+        // completed even though the SQL capability was available.
+        statisticalValidation: profile.sqlAvailable && !stopReason ? 'completed' : 'skipped',
         llmRelationshipValidation: llmProposalResult.summary,
         warnings,
+        partial: stopReason ? { reason: stopReason } : null,
     };
 }