From 5d71c9676365f3f09735a4204386360166776813 Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Wed, 7 Oct 2026 17:17:03 +0530 Subject: [PATCH 01/10] fix(ai): bound schema context length and handle context overflow errors --- src/app/api/ai/describe-schema/route.ts | 4 ++ src/components/DatabaseDocs.tsx | 53 ++++++++++++++++++------- src/lib/api/errors.ts | 12 ++++++ src/lib/llm/types.ts | 15 +++++++ tests/api/ai/describe-schema.test.ts | 26 ++++++++++++ tests/components/DatabaseDocs.test.tsx | 42 ++++++++++++++++++++ tests/unit/api-errors.test.ts | 9 +++++ tests/unit/llm/retry.test.ts | 26 +++++++++++- 8 files changed, 172 insertions(+), 15 deletions(-) diff --git a/src/app/api/ai/describe-schema/route.ts b/src/app/api/ai/describe-schema/route.ts index 8141b5e3b..892eb8497 100644 --- a/src/app/api/ai/describe-schema/route.ts +++ b/src/app/api/ai/describe-schema/route.ts @@ -1,5 +1,6 @@ import { NextRequest, NextResponse } from "next/server"; import { createLLMProvider } from "@/lib/llm"; +import { MAX_SCHEMA_CONTEXT_CHARS } from "@/lib/llm/types"; import { createErrorResponse } from "@/lib/api/errors"; import { guardRoute } from "@/lib/api/require-session"; @@ -13,6 +14,9 @@ export async function POST(req: NextRequest) { if (!schemaContext) { return NextResponse.json({ error: "Schema context required" }, { status: 400 }); } + if (typeof schemaContext === "string" && schemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { + return NextResponse.json({ error: "The schema context is too large for AI documentation." }, { status: 413 }); + } const provider = await createLLMProvider(); diff --git a/src/components/DatabaseDocs.tsx b/src/components/DatabaseDocs.tsx index a1d0c2231..75c1dfd63 100644 --- a/src/components/DatabaseDocs.tsx +++ b/src/components/DatabaseDocs.tsx @@ -9,6 +9,7 @@ import { objectPathLabel, pathKey } from "@/lib/db/object-path"; import type { ProviderCapabilities } from "@/lib/db/types"; import { renderInline } from "@/components/rich-text"; import { downloadText } from "@/lib/export/download"; +import { MAX_SCHEMA_CONTEXT_CHARS } from "@/lib/llm/types"; interface DatabaseDocsProps { schema: readonly DetailedObject[]; @@ -57,20 +58,44 @@ export function DatabaseDocs({ schema, schemaContext, databaseType, capabilities let filteredSchemaStr = ""; if (schemaContext) { try { - const tables = JSON.parse(schemaContext); - filteredSchemaStr = tables - .slice(0, 50) - .map((t: ParsedSchemaTable) => { - const cols = - t.columns - ?.map( - (c) => - `${c.name} (${c.type}${c.isPrimary ? ", PK" : ""}${c.isNullable === false ? ", NOT NULL" : ""})`, - ) - .join(", ") || ""; - return `Table: ${t.name} (${t.rowCount || 0} rows)\nColumns: ${cols}`; - }) - .join("\n\n"); + const tables = JSON.parse(schemaContext) as ParsedSchemaTable[]; + const TRUNCATION_NOTE = "\n\n[Note: Schema was truncated to fit model context limit]"; + const budget = MAX_SCHEMA_CONTEXT_CHARS - TRUNCATION_NOTE.length; + let currentChars = 0; + const tableEntries: string[] = []; + let truncated = false; + + for (const t of tables.slice(0, 50)) { + const allCols = + t.columns + ?.map( + (c) => + `${c.name} (${c.type}${c.isPrimary ? ", PK" : ""}${c.isNullable === false ? ", NOT NULL" : ""})`, + ) + .join(", ") || ""; + const fullTableText = `Table: ${t.name} (${t.rowCount || 0} rows)\nColumns: ${allCols}`; + const separatorLen = tableEntries.length > 0 ? 2 : 0; + + if (currentChars + separatorLen + fullTableText.length <= budget) { + tableEntries.push(fullTableText); + currentChars += separatorLen + fullTableText.length; + } else { + truncated = true; + const colCount = t.columns?.length || 0; + const summaryText = `Table: ${t.name} (${t.rowCount || 0} rows) - [${colCount} columns omitted for length]`; + if (currentChars + separatorLen + summaryText.length <= budget) { + tableEntries.push(summaryText); + currentChars += separatorLen + summaryText.length; + } else { + break; + } + } + } + + filteredSchemaStr = tableEntries.join("\n\n"); + if (truncated) { + filteredSchemaStr += TRUNCATION_NOTE; + } } catch { filteredSchemaStr = schemaContext.substring(0, 5000); } diff --git a/src/lib/api/errors.ts b/src/lib/api/errors.ts index 641f2dee2..630ccf8eb 100644 --- a/src/lib/api/errors.ts +++ b/src/lib/api/errors.ts @@ -23,6 +23,7 @@ import { LLMRateLimitError, LLMSafetyError, LLMStreamError, + isContextLengthError, } from "@/lib/llm/types"; import { RateLimitError } from "@/lib/api/rate-limit"; import { SeedConnectionError } from "@/lib/seed/resolve-connection"; @@ -223,6 +224,17 @@ export function createErrorResponse(error: unknown, context?: { route?: string } // --- LLM: Stream --- if (error instanceof LLMStreamError) { logger.error("LLM stream error", error, { route, provider: error.provider }); + if (isContextLengthError(error)) { + return NextResponse.json( + { + error: "The schema is too large for the configured model's context.", + code: ApiErrorCode.LLM_STREAM, + statusCode: 502, + retryable: false, + }, + { status: 502 }, + ); + } return NextResponse.json( { error: error.message, code: ApiErrorCode.LLM_STREAM, statusCode: 502, retryable: true }, { status: 502 }, diff --git a/src/lib/llm/types.ts b/src/lib/llm/types.ts index 77990c9c2..40e37c3e7 100644 --- a/src/lib/llm/types.ts +++ b/src/lib/llm/types.ts @@ -184,10 +184,25 @@ export function isRetryableError(error: unknown): boolean { return false; } + // Context length exceeded errors are not retryable + if (isContextLengthError(error)) { + return false; + } + // Rate limit and stream errors may be retryable return true; } +/** Maximum character budget for schema context in AI documentation */ +export const MAX_SCHEMA_CONTEXT_CHARS = 30_000; + +/** Whether an error indicates the prompt exceeded the model's available context window */ +export function isContextLengthError(error: unknown): boolean { + if (!error) return false; + const msg = error instanceof Error ? error.message : String(error); + return msg.includes("exceeds the available context size"); +} + /** * Whether the endpoint refused the model's TOOL CALL because it could not parse it. * diff --git a/tests/api/ai/describe-schema.test.ts b/tests/api/ai/describe-schema.test.ts index 3f64214b9..52fb9ee9b 100644 --- a/tests/api/ai/describe-schema.test.ts +++ b/tests/api/ai/describe-schema.test.ts @@ -1,6 +1,7 @@ import { describe, test, expect, mock, beforeEach } from "bun:test"; import { createMockRequest, readStreamResponse, parseResponseJSON } from "../../helpers/mock-next"; import { clearRateLimitState } from "@/lib/api/rate-limit"; +import { MAX_SCHEMA_CONTEXT_CHARS } from "@/lib/llm/types"; // ─── Mock helpers ─────────────────────────────────────────────────────────── @@ -298,4 +299,29 @@ describe("POST /api/ai/describe-schema", () => { }; expect(callArgs.messages[0].content).toContain("Database type: SQL"); }); + + test("refuses over-budget schemaContext with 413 without calling provider", async () => { + const hugeContext = "a".repeat(MAX_SCHEMA_CONTEXT_CHARS + 1); + const req = createMockRequest("/api/ai/describe-schema", { + method: "POST", + body: { schemaContext: hugeContext }, + }); + + const res = await POST(req as never); + expect(res.status).toBe(413); + + const data = await parseResponseJSON<{ error: string }>(res); + expect(data.error).toBe("The schema context is too large for AI documentation."); + }); + + test("accepts schemaContext within the budget", async () => { + const validContext = "a".repeat(MAX_SCHEMA_CONTEXT_CHARS); + const req = createMockRequest("/api/ai/describe-schema", { + method: "POST", + body: { schemaContext: validContext }, + }); + + const res = await POST(req as never); + expect(res.status).toBe(200); + }); }); diff --git a/tests/components/DatabaseDocs.test.tsx b/tests/components/DatabaseDocs.test.tsx index ee50d4fe3..a4162c9a8 100644 --- a/tests/components/DatabaseDocs.test.tsx +++ b/tests/components/DatabaseDocs.test.tsx @@ -10,6 +10,7 @@ import { DatabaseDocs } from "@/components/DatabaseDocs"; import { schemaContextOf, type DetailedObject } from "@/lib/db/detailed-object"; import { SAMPLED_MARKER, sampledSchema } from "../fixtures/sampled-schema"; import type { ProviderCapabilities } from "@/lib/db/types"; +import { MAX_SCHEMA_CONTEXT_CHARS } from "@/lib/llm/types"; const schema: DetailedObject[] = [ { @@ -436,6 +437,47 @@ describe("DatabaseDocs", () => { expect(queryByText("AI Describe")).not.toBeNull(); expect(queryByText("Export MD")).not.toBeNull(); }); + + test("wide schema produces bounded schemaContext that names truncation", async () => { + const user = userEvent.setup(); + let capturedBody = ""; + globalThis.fetch = mock((_url, init) => { + capturedBody = (init?.body as string) || ""; + return Promise.resolve({ + ok: true, + body: new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("Doc")); + controller.close(); + }, + }), + }); + }) as unknown as typeof fetch; + + // Build a large schema with 50 tables of 50 columns each + const largeSchema = Array.from({ length: 50 }, (_, i) => ({ + name: `wide_table_${i}`, + rowCount: 100, + columns: Array.from({ length: 50 }, (_, j) => ({ + name: `col_${j}_with_a_very_long_descriptive_name`, + type: "VARCHAR(255)", + })), + })); + + const { queryByText } = render( + , + ); + + await user.click(queryByText("AI Describe")!); + + await waitFor(() => { + expect(capturedBody).not.toBe(""); + }); + + const parsed = JSON.parse(capturedBody); + expect(parsed.schemaContext).toContain("[Note: Schema was truncated to fit model context limit]"); + expect(parsed.schemaContext.length).toBeLessThanOrEqual(MAX_SCHEMA_CONTEXT_CHARS); + }); }); // ============================================================================= diff --git a/tests/unit/api-errors.test.ts b/tests/unit/api-errors.test.ts index 76d108855..5c7780d62 100644 --- a/tests/unit/api-errors.test.ts +++ b/tests/unit/api-errors.test.ts @@ -203,6 +203,15 @@ describe("createErrorResponse", () => { expect(body.statusCode).toBe(502); expect(body.retryable).toBe(true); }); + test("LLMStreamError with context length exceeded returns 502 with friendly error and retryable false", async () => { + const err = new LLMStreamError("request exceeds the available context size (32768 tokens)", "openai"); + const res = createErrorResponse(err); + expect(res.status).toBe(502); + const body = await res.json(); + expect(body.code).toBe("LLM_STREAM"); + expect(body.error).toBe("The schema is too large for the configured model's context."); + expect(body.retryable).toBe(false); + }); test("LLMError (base) uses statusCode or 500", async () => { const err = new LLMError("generic llm error", "gemini", 503); diff --git a/tests/unit/llm/retry.test.ts b/tests/unit/llm/retry.test.ts index cfc5f527c..51447e10b 100644 --- a/tests/unit/llm/retry.test.ts +++ b/tests/unit/llm/retry.test.ts @@ -1,6 +1,6 @@ import { describe, test, expect, beforeEach, afterEach, mock } from "bun:test"; import { withRetry, makeRetryable } from "@/lib/llm/utils/retry"; -import { LLMConfigError, LLMRateLimitError } from "@/lib/llm/types"; +import { LLMConfigError, LLMRateLimitError, LLMStreamError } from "@/lib/llm/types"; // Suppress console.error output during retry tests let originalConsoleError: typeof console.error; @@ -89,6 +89,30 @@ describe("withRetry", () => { } }); + test("context-length error is not retried and fails immediately", async () => { + let attempts = 0; + const fn = mock(async () => { + attempts++; + throw new LLMStreamError("request exceeds the available context size (32768 tokens)", "openai"); + }); + + await expect(withRetry(fn, { maxAttempts: 3, initialDelay: 1, maxDelay: 5 })).rejects.toThrow( + /exceeds the available context size/, + ); + expect(attempts).toBe(1); + }); + + test("standard stream error is retried up to maxAttempts", async () => { + let attempts = 0; + const fn = mock(async () => { + attempts++; + throw new LLMStreamError("stream interrupted", "openai"); + }); + + await expect(withRetry(fn, { maxAttempts: 3, initialDelay: 1, maxDelay: 5 })).rejects.toThrow(/stream interrupted/); + expect(attempts).toBe(3); + }); + test("maxAttempts=1 means no retry", async () => { const fn = mock(async () => { throw new LLMRateLimitError("rate limited", "openai"); From f780131d3c2c68b8b91485a28946d6bf5ea6ef6d Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Wed, 7 Oct 2026 18:02:49 +0530 Subject: [PATCH 02/10] test: mock isContextLengthError and MAX_SCHEMA_CONTEXT_CHARS in route tests --- tests/api/ai/describe-schema.test.ts | 8 ++++++++ tests/api/ai/explain.test.ts | 8 ++++++++ tests/api/ai/query-safety.test.ts | 8 ++++++++ tests/security/route-auth.test.ts | 8 ++++++++ 4 files changed, 32 insertions(+) diff --git a/tests/api/ai/describe-schema.test.ts b/tests/api/ai/describe-schema.test.ts index 52fb9ee9b..90350e051 100644 --- a/tests/api/ai/describe-schema.test.ts +++ b/tests/api/ai/describe-schema.test.ts @@ -75,6 +75,14 @@ mock.module("@/lib/llm/types", () => ({ LLMRateLimitError: MockLLMRateLimitError, LLMSafetyError: MockLLMSafetyError, LLMStreamError: MockLLMStreamError, + /* + Same whole-module mock requirement: `errors.ts` imports `isContextLengthError` + and route handlers import `MAX_SCHEMA_CONTEXT_CHARS`. + */ + isContextLengthError: (error: unknown): boolean => + (error instanceof MockLLMStreamError || error instanceof Error) && + error.message.includes("exceeds the available context size"), + MAX_SCHEMA_CONTEXT_CHARS: 30_000, })); const mockGetSession = mock( diff --git a/tests/api/ai/explain.test.ts b/tests/api/ai/explain.test.ts index 1bd1d88cd..54306a360 100644 --- a/tests/api/ai/explain.test.ts +++ b/tests/api/ai/explain.test.ts @@ -74,6 +74,14 @@ mock.module("@/lib/llm/types", () => ({ LLMRateLimitError: MockLLMRateLimitError, LLMSafetyError: MockLLMSafetyError, LLMStreamError: MockLLMStreamError, + /* + Same whole-module mock requirement: `errors.ts` imports `isContextLengthError` + and route handlers import `MAX_SCHEMA_CONTEXT_CHARS`. + */ + isContextLengthError: (error: unknown): boolean => + (error instanceof MockLLMStreamError || error instanceof Error) && + error.message.includes("exceeds the available context size"), + MAX_SCHEMA_CONTEXT_CHARS: 30_000, })); const mockGetSession = mock( diff --git a/tests/api/ai/query-safety.test.ts b/tests/api/ai/query-safety.test.ts index a743a4781..9e5766898 100644 --- a/tests/api/ai/query-safety.test.ts +++ b/tests/api/ai/query-safety.test.ts @@ -77,6 +77,14 @@ mock.module("@/lib/llm/types", () => ({ LLMRateLimitError: MockLLMRateLimitError, LLMSafetyError: MockLLMSafetyError, LLMStreamError: MockLLMStreamError, + /* + Same whole-module mock requirement: `errors.ts` imports `isContextLengthError` + and route handlers import `MAX_SCHEMA_CONTEXT_CHARS`. + */ + isContextLengthError: (error: unknown): boolean => + (error instanceof MockLLMStreamError || error instanceof Error) && + error.message.includes("exceeds the available context size"), + MAX_SCHEMA_CONTEXT_CHARS: 30_000, })); const mockGetSession = mock( diff --git a/tests/security/route-auth.test.ts b/tests/security/route-auth.test.ts index 377fa897f..4aec84d8f 100644 --- a/tests/security/route-auth.test.ts +++ b/tests/security/route-auth.test.ts @@ -110,6 +110,14 @@ mock.module("@/lib/llm/types", () => ({ */ isToolCallParseError: (error: unknown): boolean => error instanceof MockLLMStreamError && /parsing tool call/i.test(error.message), + /* + Same whole-module mock requirement: `errors.ts` imports `isContextLengthError` + and route handlers import `MAX_SCHEMA_CONTEXT_CHARS`. + */ + isContextLengthError: (error: unknown): boolean => + (error instanceof MockLLMStreamError || error instanceof Error) && + error.message.includes("exceeds the available context size"), + MAX_SCHEMA_CONTEXT_CHARS: 30_000, })); const { guardRoute } = await import("@/lib/api/require-session"); From 1c0873fabf981cc221181d8363d5d41a5dc04d84 Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Thu, 8 Oct 2026 21:07:56 +0530 Subject: [PATCH 03/10] fix(ai): use neutral context length message and validate input types on describe-schema route --- src/app/api/ai/describe-schema/route.ts | 9 ++++--- src/lib/api/errors.ts | 2 +- tests/api/ai/describe-schema.test.ts | 32 +++++++++++++++++++++++++ tests/unit/api-errors.test.ts | 2 +- 4 files changed, 40 insertions(+), 5 deletions(-) diff --git a/src/app/api/ai/describe-schema/route.ts b/src/app/api/ai/describe-schema/route.ts index 892eb8497..b3cff3873 100644 --- a/src/app/api/ai/describe-schema/route.ts +++ b/src/app/api/ai/describe-schema/route.ts @@ -11,10 +11,13 @@ export async function POST(req: NextRequest) { try { const { schemaContext, databaseType, mode } = await req.json(); - if (!schemaContext) { - return NextResponse.json({ error: "Schema context required" }, { status: 400 }); + if (typeof schemaContext !== "string" || !schemaContext) { + return NextResponse.json({ error: "Schema context is required and must be a string." }, { status: 400 }); } - if (typeof schemaContext === "string" && schemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { + if (databaseType !== undefined && (typeof databaseType !== "string" || databaseType.length > 64)) { + return NextResponse.json({ error: "Database type must be a string of at most 64 characters." }, { status: 400 }); + } + if (schemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { return NextResponse.json({ error: "The schema context is too large for AI documentation." }, { status: 413 }); } diff --git a/src/lib/api/errors.ts b/src/lib/api/errors.ts index 630ccf8eb..bc1cd4124 100644 --- a/src/lib/api/errors.ts +++ b/src/lib/api/errors.ts @@ -227,7 +227,7 @@ export function createErrorResponse(error: unknown, context?: { route?: string } if (isContextLengthError(error)) { return NextResponse.json( { - error: "The schema is too large for the configured model's context.", + error: "The request is too large for the configured model's context.", code: ApiErrorCode.LLM_STREAM, statusCode: 502, retryable: false, diff --git a/tests/api/ai/describe-schema.test.ts b/tests/api/ai/describe-schema.test.ts index 90350e051..ed4c0361d 100644 --- a/tests/api/ai/describe-schema.test.ts +++ b/tests/api/ai/describe-schema.test.ts @@ -320,6 +320,38 @@ describe("POST /api/ai/describe-schema", () => { const data = await parseResponseJSON<{ error: string }>(res); expect(data.error).toBe("The schema context is too large for AI documentation."); + expect(mockCreateLLMProvider).not.toHaveBeenCalled(); + }); + + test("returns 400 when schemaContext is not a string without calling provider", async () => { + const req = createMockRequest("/api/ai/describe-schema", { + method: "POST", + body: { schemaContext: ["x".repeat(200_000)] }, + }); + + const res = await POST(req as never); + expect(res.status).toBe(400); + + const data = await parseResponseJSON<{ error: string }>(res); + expect(data.error).toBeDefined(); + expect(mockCreateLLMProvider).not.toHaveBeenCalled(); + }); + + test("returns 400 when databaseType is oversized without calling provider", async () => { + const req = createMockRequest("/api/ai/describe-schema", { + method: "POST", + body: { + schemaContext: "users(id, name)", + databaseType: "x".repeat(200_000), + }, + }); + + const res = await POST(req as never); + expect(res.status).toBe(400); + + const data = await parseResponseJSON<{ error: string }>(res); + expect(data.error).toBeDefined(); + expect(mockCreateLLMProvider).not.toHaveBeenCalled(); }); test("accepts schemaContext within the budget", async () => { diff --git a/tests/unit/api-errors.test.ts b/tests/unit/api-errors.test.ts index 5c7780d62..73e06102b 100644 --- a/tests/unit/api-errors.test.ts +++ b/tests/unit/api-errors.test.ts @@ -209,7 +209,7 @@ describe("createErrorResponse", () => { expect(res.status).toBe(502); const body = await res.json(); expect(body.code).toBe("LLM_STREAM"); - expect(body.error).toBe("The schema is too large for the configured model's context."); + expect(body.error).toBe("The request is too large for the configured model's context."); expect(body.retryable).toBe(false); }); From 30c8dc4a26fd9e20f4194ca81a3b09baacd27a5c Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Fri, 9 Oct 2026 18:03:06 +0530 Subject: [PATCH 04/10] fix(profiler): send only profiled table schema and display AI summary errors --- src/components/DataProfiler.tsx | 59 +++++++++++- tests/components/DataProfiler.test.tsx | 122 +++++++++++++++++++++++++ 2 files changed, 176 insertions(+), 5 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index 52a08c85b..6e3292d1e 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -18,6 +18,7 @@ import { import { buildConnectionPayload } from "@/hooks/use-connection-payload"; import { dataProfileText, type ColumnProfile, type ProfileData } from "@/lib/export/data-profile"; import { downloadText } from "@/lib/export/download"; +import { MAX_SCHEMA_CONTEXT_CHARS } from "@/lib/llm/types"; import { Button } from "./ui/button"; import { DropdownMenu, @@ -76,6 +77,42 @@ interface DataProfilerProps { userRole?: string; } +function resolveTableSchemaContext( + schemaContext: string | undefined, + tableSchema: DetailedObject | null, + tableName: string, +): string { + if (schemaContext) { + try { + const parsed = JSON.parse(schemaContext); + if (Array.isArray(parsed)) { + const found = parsed.find( + (t: { name?: string; path?: readonly string[] }) => + t.name === tableName || + (Array.isArray(t.path) && (t.path[t.path.length - 1] === tableName || t.path.join(".") === tableName)), + ); + if (found) { + return JSON.stringify(found); + } + if (tableSchema) { + return JSON.stringify(tableSchema); + } + } else { + return schemaContext; + } + } catch { + if (schemaContext.length > MAX_SCHEMA_CONTEXT_CHARS && tableSchema) { + return JSON.stringify(tableSchema); + } + return schemaContext; + } + } + if (tableSchema) { + return JSON.stringify(tableSchema); + } + return ""; +} + export function DataProfiler({ isOpen, onClose, @@ -93,6 +130,7 @@ export function DataProfiler({ const [isLoading, setIsLoading] = useState(false); const [profile, setProfile] = useState(null); const [aiSummary, setAiSummary] = useState(""); + const [aiError, setAiError] = useState(null); const [isAiLoading, setIsAiLoading] = useState(false); const [error, setError] = useState(null); // The address as one string, for the header and for the AI prompt. `pathKey` is the @@ -141,6 +179,7 @@ export function DataProfiler({ const fetchAiSummary = async (data: ProfileData) => { setIsAiLoading(true); + setAiError(null); try { // The profile table lists every column, but the summary goes to a model, which never sees a column the // engine only inferred from sampled data: `machineColumns` decides which, as for `schemaContext` below. @@ -164,7 +203,11 @@ export function DataProfiler({ ) .join("\n"); - const fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${schemaContext || ""}`; + const tableSchemaSnippet = resolveTableSchemaContext(schemaContext, tableSchema, tableName); + let fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${tableSchemaSnippet}`; + if (fullSchemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { + fullSchemaContext = fullSchemaContext.slice(0, MAX_SCHEMA_CONTEXT_CHARS); + } if (onDescribeSchema) { // Platform adapter: use callback instead of fetch @@ -182,7 +225,11 @@ export function DataProfiler({ }), }); - if (!response.ok) return; + if (!response.ok) { + const err = await response.json().catch(() => ({ error: "AI summary failed" })); + setAiError(err.error || "AI summary failed"); + return; + } const reader = response.body?.getReader(); if (!reader) return; @@ -195,8 +242,8 @@ export function DataProfiler({ setAiSummary(full); } } - } catch { - // AI summary is optional, don't show error + } catch (err) { + setAiError(err instanceof Error ? err.message : "AI summary failed"); } finally { setIsAiLoading(false); } @@ -254,6 +301,7 @@ export function DataProfiler({ return () => { setProfile(null); setAiSummary(""); + setAiError(null); setError(null); }; // eslint-disable-next-line react-hooks/exhaustive-deps @@ -523,7 +571,7 @@ export function DataProfiler({ {/* AI Summary */} - {(aiSummary || isAiLoading) && ( + {(aiSummary || isAiLoading || aiError) && (
@@ -533,6 +581,7 @@ export function DataProfiler({ {aiSummary && (
{aiSummary}
)} + {aiError &&
{aiError}
}
)} diff --git a/tests/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 4f414cc5e..257da8b65 100644 --- a/tests/components/DataProfiler.test.tsx +++ b/tests/components/DataProfiler.test.tsx @@ -470,6 +470,128 @@ describe("DataProfiler", () => { globalThis.fetch = originalFetch; }); + test("sends profiler summary request under the limit when schemaContext exceeds 30,000 characters", async () => { + const aiText = "Profiled summary for users table."; + restoreGlobalFetch(); + + const originalFetch = globalThis.fetch; + let capturedBody: { schemaContext: string } | null = null; + + const largeSchema = JSON.stringify( + Array.from({ length: 60 }, (_, i) => ({ + name: i === 0 ? "users" : `table_${i}`, + columns: Array.from({ length: 25 }, (_, j) => ({ + name: `col_${j}`, + type: "varchar(255)", + nullable: true, + isPrimary: j === 0, + })), + rowCount: 100, + })), + ); + expect(largeSchema.length).toBeGreaterThan(30_000); + + globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + const pathname = new URL(url, "http://localhost:3000").pathname; + + if (pathname.includes("/api/db/profile")) { + return new Response(JSON.stringify(mockProfileResponse), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + if (pathname.includes("/api/ai/describe-schema")) { + const bodyStr = typeof init?.body === "string" ? init.body : await (input as Request).text(); + capturedBody = JSON.parse(bodyStr); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(aiText)); + controller.close(); + }, + }); + return new Response(stream, { + status: 200, + headers: { "content-type": "text/plain" }, + }); + } + + return new Response("Not found", { status: 404 }); + }) as unknown as typeof fetch; + + const props = createDefaultProps({ schemaContext: largeSchema }); + const { container } = render(); + const view = within(container); + + await waitFor(() => { + expect(view.queryByText("Column Profiles")).not.toBeNull(); + }); + + await waitFor( + () => { + expect(view.queryByText("AI Analysis")).not.toBeNull(); + }, + { timeout: 3000 }, + ); + + await waitFor( + () => { + expect(view.queryByText(aiText)).not.toBeNull(); + }, + { timeout: 3000 }, + ); + + expect(capturedBody).not.toBeNull(); + expect(capturedBody!.schemaContext.length).toBeLessThanOrEqual(30_000); + + globalThis.fetch = originalFetch; + }); + + test("displays AI error when describe-schema route refuses", async () => { + const originalFetch = globalThis.fetch; + const refusalError = "The request is too large for the configured model's context."; + restoreGlobalFetch(); + + globalThis.fetch = mock(async (input: RequestInfo | URL) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + const pathname = new URL(url, "http://localhost:3000").pathname; + + if (pathname.includes("/api/db/profile")) { + return new Response(JSON.stringify(mockProfileResponse), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + if (pathname.includes("/api/ai/describe-schema")) { + return new Response(JSON.stringify({ error: refusalError }), { + status: 502, + headers: { "content-type": "application/json" }, + }); + } + + return new Response("Not found", { status: 404 }); + }) as unknown as typeof fetch; + + const props = createDefaultProps(); + const { container } = render(); + const view = within(container); + + await waitFor(() => { + expect(view.queryByText("Column Profiles")).not.toBeNull(); + }); + + await waitFor( + () => { + expect(view.queryByText(refusalError)).not.toBeNull(); + }, + { timeout: 3000 }, + ); + + globalThis.fetch = originalFetch; + }); + // ── Profile fetch error handling ───────────────────────────────────────── test("displays error message when profile fetch fails", async () => { From bebd539f9b48e2a77681f8faf23d35bb358ac0c5 Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Fri, 9 Oct 2026 18:34:33 +0530 Subject: [PATCH 05/10] fix(profiler): preserve single-table schemaContext and match tablePath segments --- src/components/DataProfiler.tsx | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index 6e3292d1e..d904ee190 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -81,15 +81,21 @@ function resolveTableSchemaContext( schemaContext: string | undefined, tableSchema: DetailedObject | null, tableName: string, + tablePath?: readonly string[], ): string { if (schemaContext) { try { const parsed = JSON.parse(schemaContext); if (Array.isArray(parsed)) { + if (parsed.length <= 1) { + return schemaContext; + } + const lastSegment = tablePath && tablePath.length > 0 ? tablePath[tablePath.length - 1] : tableName; const found = parsed.find( (t: { name?: string; path?: readonly string[] }) => t.name === tableName || - (Array.isArray(t.path) && (t.path[t.path.length - 1] === tableName || t.path.join(".") === tableName)), + t.name === lastSegment || + (Array.isArray(t.path) && (t.path[t.path.length - 1] === lastSegment || t.path.join(".") === tableName)), ); if (found) { return JSON.stringify(found); @@ -97,9 +103,9 @@ function resolveTableSchemaContext( if (tableSchema) { return JSON.stringify(tableSchema); } - } else { return schemaContext; } + return schemaContext; } catch { if (schemaContext.length > MAX_SCHEMA_CONTEXT_CHARS && tableSchema) { return JSON.stringify(tableSchema); @@ -203,7 +209,7 @@ export function DataProfiler({ ) .join("\n"); - const tableSchemaSnippet = resolveTableSchemaContext(schemaContext, tableSchema, tableName); + const tableSchemaSnippet = resolveTableSchemaContext(schemaContext, tableSchema, tableName, tablePath); let fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${tableSchemaSnippet}`; if (fullSchemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { fullSchemaContext = fullSchemaContext.slice(0, MAX_SCHEMA_CONTEXT_CHARS); From ea4dbb64a800ff45fba899c976a5fb0bc46fff81 Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Fri, 9 Oct 2026 18:58:49 +0530 Subject: [PATCH 06/10] fix(profiler): cover all schemaContext and AI error paths --- src/components/DataProfiler.tsx | 16 +---- tests/components/DataProfiler.test.tsx | 98 ++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 13 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index d904ee190..8332132d1 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -86,10 +86,7 @@ function resolveTableSchemaContext( if (schemaContext) { try { const parsed = JSON.parse(schemaContext); - if (Array.isArray(parsed)) { - if (parsed.length <= 1) { - return schemaContext; - } + if (Array.isArray(parsed) && parsed.length > 1) { const lastSegment = tablePath && tablePath.length > 0 ? tablePath[tablePath.length - 1] : tableName; const found = parsed.find( (t: { name?: string; path?: readonly string[] }) => @@ -100,18 +97,11 @@ function resolveTableSchemaContext( if (found) { return JSON.stringify(found); } - if (tableSchema) { - return JSON.stringify(tableSchema); - } - return schemaContext; } - return schemaContext; } catch { - if (schemaContext.length > MAX_SCHEMA_CONTEXT_CHARS && tableSchema) { - return JSON.stringify(tableSchema); - } - return schemaContext; + // Non-JSON schemaContext falls back to raw string } + return schemaContext; } if (tableSchema) { return JSON.stringify(tableSchema); diff --git a/tests/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 257da8b65..622d084a4 100644 --- a/tests/components/DataProfiler.test.tsx +++ b/tests/components/DataProfiler.test.tsx @@ -592,6 +592,104 @@ describe("DataProfiler", () => { globalThis.fetch = originalFetch; }); + test("truncates fullSchemaContext to MAX_SCHEMA_CONTEXT_CHARS when schemaContext exceeds limit", async () => { + const aiText = "Profiled summary for massive table."; + restoreGlobalFetch(); + + const originalFetch = globalThis.fetch; + let capturedBody: { schemaContext: string } | null = null; + const massiveSchema = "x".repeat(35_000); + + globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + const pathname = new URL(url, "http://localhost:3000").pathname; + + if (pathname.includes("/api/db/profile")) { + return new Response(JSON.stringify(mockProfileResponse), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + if (pathname.includes("/api/ai/describe-schema")) { + const bodyStr = typeof init?.body === "string" ? init.body : await (input as Request).text(); + capturedBody = JSON.parse(bodyStr); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(aiText)); + controller.close(); + }, + }); + return new Response(stream, { + status: 200, + headers: { "content-type": "text/plain" }, + }); + } + + return new Response("Not found", { status: 404 }); + }) as unknown as typeof fetch; + + const props = createDefaultProps({ schemaContext: massiveSchema }); + const { container } = render(); + const view = within(container); + + await waitFor(() => { + expect(view.queryByText("Column Profiles")).not.toBeNull(); + }); + + await waitFor( + () => { + expect(view.queryByText(aiText)).not.toBeNull(); + }, + { timeout: 3000 }, + ); + + expect(capturedBody).not.toBeNull(); + expect(capturedBody!.schemaContext.length).toBe(30_000); + + globalThis.fetch = originalFetch; + }); + + test("displays AI error when describe-schema fetch throws a network error", async () => { + const originalFetch = globalThis.fetch; + restoreGlobalFetch(); + + globalThis.fetch = mock(async (input: RequestInfo | URL) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + const pathname = new URL(url, "http://localhost:3000").pathname; + + if (pathname.includes("/api/db/profile")) { + return new Response(JSON.stringify(mockProfileResponse), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + + if (pathname.includes("/api/ai/describe-schema")) { + throw new Error("Network offline"); + } + + return new Response("Not found", { status: 404 }); + }) as unknown as typeof fetch; + + const props = createDefaultProps(); + const { container } = render(); + const view = within(container); + + await waitFor(() => { + expect(view.queryByText("Column Profiles")).not.toBeNull(); + }); + + await waitFor( + () => { + expect(view.queryByText("Network offline")).not.toBeNull(); + }, + { timeout: 3000 }, + ); + + globalThis.fetch = originalFetch; + }); + // ── Profile fetch error handling ───────────────────────────────────────── test("displays error message when profile fetch fails", async () => { From 329f567d75e6aa786483ec03a4c0e8bf99ea1ee6 Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Sat, 10 Oct 2026 04:00:34 +0530 Subject: [PATCH 07/10] fix(profiler): send tableSchema directly instead of parsing schemaContext --- src/components/DataProfiler.tsx | 36 +----------- tests/components/DataProfiler.test.tsx | 79 +++++++++++++++++++------- 2 files changed, 61 insertions(+), 54 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index 8332132d1..148e44882 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -77,45 +77,13 @@ interface DataProfilerProps { userRole?: string; } -function resolveTableSchemaContext( - schemaContext: string | undefined, - tableSchema: DetailedObject | null, - tableName: string, - tablePath?: readonly string[], -): string { - if (schemaContext) { - try { - const parsed = JSON.parse(schemaContext); - if (Array.isArray(parsed) && parsed.length > 1) { - const lastSegment = tablePath && tablePath.length > 0 ? tablePath[tablePath.length - 1] : tableName; - const found = parsed.find( - (t: { name?: string; path?: readonly string[] }) => - t.name === tableName || - t.name === lastSegment || - (Array.isArray(t.path) && (t.path[t.path.length - 1] === lastSegment || t.path.join(".") === tableName)), - ); - if (found) { - return JSON.stringify(found); - } - } - } catch { - // Non-JSON schemaContext falls back to raw string - } - return schemaContext; - } - if (tableSchema) { - return JSON.stringify(tableSchema); - } - return ""; -} - export function DataProfiler({ isOpen, onClose, tablePath, tableSchema, connection, - schemaContext, + schemaContext: _schemaContext, databaseType, onProfile, onDescribeSchema, @@ -199,7 +167,7 @@ export function DataProfiler({ ) .join("\n"); - const tableSchemaSnippet = resolveTableSchemaContext(schemaContext, tableSchema, tableName, tablePath); + const tableSchemaSnippet = tableSchema ? JSON.stringify(tableSchema) : ""; let fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${tableSchemaSnippet}`; if (fullSchemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { fullSchemaContext = fullSchemaContext.slice(0, MAX_SCHEMA_CONTEXT_CHARS); diff --git a/tests/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 622d084a4..374defb65 100644 --- a/tests/components/DataProfiler.test.tsx +++ b/tests/components/DataProfiler.test.tsx @@ -470,26 +470,47 @@ describe("DataProfiler", () => { globalThis.fetch = originalFetch; }); - test("sends profiler summary request under the limit when schemaContext exceeds 30,000 characters", async () => { - const aiText = "Profiled summary for users table."; + test("sends only the profiled table schema when multiple schemas contain a table with the same name", async () => { + const aiText = "Profiled summary for app.users table."; restoreGlobalFetch(); const originalFetch = globalThis.fetch; let capturedBody: { schemaContext: string } | null = null; - const largeSchema = JSON.stringify( - Array.from({ length: 60 }, (_, i) => ({ - name: i === 0 ? "users" : `table_${i}`, - columns: Array.from({ length: 25 }, (_, j) => ({ - name: `col_${j}`, - type: "varchar(255)", - nullable: true, - isPrimary: j === 0, - })), - rowCount: 100, + const publicUsersTable = { + name: "users", + kind: "table" as const, + path: ["public", "users"], + rowCount: 50, + indexes: [], + columns: [ + { name: "public_id", type: "integer", nullable: false, isPrimary: true }, + { name: "public_profile", type: "text", nullable: true, isPrimary: false }, + ], + }; + + const appUsersTable = { + name: "users", + kind: "table" as const, + path: ["app", "users"], + rowCount: 100, + indexes: [], + columns: [ + { name: "app_id", type: "integer", nullable: false, isPrimary: true }, + { name: "app_role", type: "varchar(50)", nullable: false, isPrimary: false }, + ], + }; + + // rawSchemaContext contains both schemas and exceeds 30,000 characters + const largePadding = Array.from({ length: 60 }, (_, i) => ({ + name: `other_table_${i}`, + columns: Array.from({ length: 20 }, (_, j) => ({ + name: `col_${j}`, + type: "varchar(255)", })), - ); - expect(largeSchema.length).toBeGreaterThan(30_000); + })); + const rawSchemaContext = JSON.stringify([publicUsersTable, appUsersTable, ...largePadding]); + expect(rawSchemaContext.length).toBeGreaterThan(30_000); globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; @@ -520,7 +541,11 @@ describe("DataProfiler", () => { return new Response("Not found", { status: 404 }); }) as unknown as typeof fetch; - const props = createDefaultProps({ schemaContext: largeSchema }); + const props = createDefaultProps({ + tablePath: ["app", "users"], + tableSchema: appUsersTable, + schemaContext: rawSchemaContext, + }); const { container } = render(); const view = within(container); @@ -544,6 +569,12 @@ describe("DataProfiler", () => { expect(capturedBody).not.toBeNull(); expect(capturedBody!.schemaContext.length).toBeLessThanOrEqual(30_000); + // Sends only the profiled table's columns + expect(capturedBody!.schemaContext).toContain("app_id"); + expect(capturedBody!.schemaContext).toContain("app_role"); + // Does NOT contain the other schema's columns from rawSchemaContext + expect(capturedBody!.schemaContext).not.toContain("public_id"); + expect(capturedBody!.schemaContext).not.toContain("public_profile"); globalThis.fetch = originalFetch; }); @@ -592,13 +623,21 @@ describe("DataProfiler", () => { globalThis.fetch = originalFetch; }); - test("truncates fullSchemaContext to MAX_SCHEMA_CONTEXT_CHARS when schemaContext exceeds limit", async () => { + test("truncates fullSchemaContext to MAX_SCHEMA_CONTEXT_CHARS when tableSchema exceeds limit", async () => { const aiText = "Profiled summary for massive table."; restoreGlobalFetch(); const originalFetch = globalThis.fetch; let capturedBody: { schemaContext: string } | null = null; - const massiveSchema = "x".repeat(35_000); + const massiveTable = { + ...mockUsersTable, + columns: Array.from({ length: 800 }, (_, i) => ({ + name: `col_${i}_${"x".repeat(35)}`, + type: "varchar(255)", + nullable: true, + isPrimary: false, + })), + }; globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; @@ -629,7 +668,7 @@ describe("DataProfiler", () => { return new Response("Not found", { status: 404 }); }) as unknown as typeof fetch; - const props = createDefaultProps({ schemaContext: massiveSchema }); + const props = createDefaultProps({ tableSchema: massiveTable }); const { container } = render(); const view = within(container); @@ -824,7 +863,7 @@ describe("DataProfiler", () => { const onProfile = mock(async () => mockProfileResponse); const onDescribeSchema = mock(async () => "Adapter AI summary"); - const props = createDefaultProps({ onProfile, onDescribeSchema, schemaContext: "schema ctx" }); + const props = createDefaultProps({ onProfile, onDescribeSchema }); const { container } = render(); const view = within(container); @@ -854,7 +893,7 @@ describe("DataProfiler", () => { }; expect(describeArg.tableName).toBe("app.users"); expect(describeArg.schemaContext).toContain("Column Profiles:"); - expect(describeArg.schemaContext).toContain("schema ctx"); + expect(describeArg.schemaContext).toContain("users"); // Neither endpoint is hit when the adapters are provided expect(fetchSpy).not.toHaveBeenCalled(); From ffce00ea9e70da564ccbe97ddc30703e170fc0fb Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Sat, 10 Oct 2026 04:22:14 +0530 Subject: [PATCH 08/10] fix(profiler): preserve single-table schemaContext and filter sampled columns --- src/components/DataProfiler.tsx | 26 ++++++++++++++++++++++-- tests/components/DataProfiler.test.tsx | 28 ++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index 148e44882..0972ede45 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -77,13 +77,35 @@ interface DataProfilerProps { userRole?: string; } +function resolveTableSchemaSnippet( + schemaContext: string | undefined, + tableSchema: DetailedObject | null, +): string { + if (schemaContext) { + try { + const parsed = JSON.parse(schemaContext); + if (Array.isArray(parsed) && parsed.length === 1) { + return schemaContext; + } + } catch { + // Non-JSON + } + } + if (!tableSchema) return ""; + const cleanTable = { + ...tableSchema, + columns: tableSchema.columns?.filter((c) => (c as { provenance?: unknown }).provenance !== "sampled") ?? [], + }; + return JSON.stringify(cleanTable); +} + export function DataProfiler({ isOpen, onClose, tablePath, tableSchema, connection, - schemaContext: _schemaContext, + schemaContext, databaseType, onProfile, onDescribeSchema, @@ -167,7 +189,7 @@ export function DataProfiler({ ) .join("\n"); - const tableSchemaSnippet = tableSchema ? JSON.stringify(tableSchema) : ""; + const tableSchemaSnippet = resolveTableSchemaSnippet(schemaContext, tableSchema); let fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${tableSchemaSnippet}`; if (fullSchemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { fullSchemaContext = fullSchemaContext.slice(0, MAX_SCHEMA_CONTEXT_CHARS); diff --git a/tests/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 374defb65..124c7c69d 100644 --- a/tests/components/DataProfiler.test.tsx +++ b/tests/components/DataProfiler.test.tsx @@ -579,6 +579,34 @@ describe("DataProfiler", () => { globalThis.fetch = originalFetch; }); + test("preserves single-table schemaContext when provided", async () => { + restoreGlobalFetch(); + let capturedBody: { schemaContext: string } | null = null; + const singleTableContext = JSON.stringify([{ name: "single_table", columns: [] }]); + + const originalFetch = globalThis.fetch; + globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (url.includes("/api/db/profile")) { + return new Response(JSON.stringify(mockProfileResponse), { status: 200, headers: { "content-type": "application/json" } }); + } + if (url.includes("/api/ai/describe-schema")) { + const bodyStr = typeof init?.body === "string" ? init.body : await (input as Request).text(); + capturedBody = JSON.parse(bodyStr); + return new Response("summary", { status: 200 }); + } + return new Response("Not found", { status: 404 }); + }) as unknown as typeof fetch; + + const props = createDefaultProps({ schemaContext: singleTableContext }); + render(); + await waitFor(() => { + expect(capturedBody).not.toBeNull(); + }); + expect(capturedBody!.schemaContext).toContain(singleTableContext); + globalThis.fetch = originalFetch; + }); + test("displays AI error when describe-schema route refuses", async () => { const originalFetch = globalThis.fetch; const refusalError = "The request is too large for the configured model's context."; From 0b30d88d3a98ec600e7c27b8db32d6c4e242394b Mon Sep 17 00:00:00 2001 From: Maaz2212 Date: Sat, 10 Oct 2026 04:32:46 +0530 Subject: [PATCH 09/10] style: format DataProfiler and DataProfiler.test --- src/components/DataProfiler.tsx | 5 +---- tests/components/DataProfiler.test.tsx | 5 ++++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index 0972ede45..a8ed22b33 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -77,10 +77,7 @@ interface DataProfilerProps { userRole?: string; } -function resolveTableSchemaSnippet( - schemaContext: string | undefined, - tableSchema: DetailedObject | null, -): string { +function resolveTableSchemaSnippet(schemaContext: string | undefined, tableSchema: DetailedObject | null): string { if (schemaContext) { try { const parsed = JSON.parse(schemaContext); diff --git a/tests/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 124c7c69d..20cd0f452 100644 --- a/tests/components/DataProfiler.test.tsx +++ b/tests/components/DataProfiler.test.tsx @@ -588,7 +588,10 @@ describe("DataProfiler", () => { globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; if (url.includes("/api/db/profile")) { - return new Response(JSON.stringify(mockProfileResponse), { status: 200, headers: { "content-type": "application/json" } }); + return new Response(JSON.stringify(mockProfileResponse), { + status: 200, + headers: { "content-type": "application/json" }, + }); } if (url.includes("/api/ai/describe-schema")) { const bodyStr = typeof init?.body === "string" ? init.body : await (input as Request).text(); From 6a9d7343c7e4a1fd0875de23a12a24c27307e090 Mon Sep 17 00:00:00 2001 From: cevheri Date: Sun, 11 Oct 2026 02:35:07 +0300 Subject: [PATCH 10/10] fix(profiler): hand the model the profiled object through schemaContextOf The AI summary built its schema snippet by hand: a one-object schemaContext went through raw, anything else went as the table's JSON with sampled columns filtered inline. That bypassed schemaContextOf, the projection every AI-bound schema takes, so an object's readRanges would have reached the model. Send schemaContextOf([tableSchema]) and drop the schemaContext prop, which nothing reads any more. Main's sampled-column tests now profile the sampled table itself, and a new test pins the snippet to schemaContextOf. --- src/components/DataProfiler.tsx | 27 +--- src/components/studio/StudioModals.tsx | 1 - src/lib/db/detailed-object.ts | 3 +- tests/components/DataProfiler.test.tsx | 177 +++++-------------------- 4 files changed, 38 insertions(+), 170 deletions(-) diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index a8ed22b33..47ff401e3 100644 --- a/src/components/DataProfiler.tsx +++ b/src/components/DataProfiler.tsx @@ -6,7 +6,7 @@ import { LoaderCircle, ChartColumn, X, CircleAlert, Sparkles, Lock, Download } f import { cn } from "@/lib/utils"; import { DatabaseConnection } from "@/lib/types"; import { objectPathLabel, pathKey } from "@/lib/db/object-path"; -import { machineColumns, type DetailedObject } from "@/lib/db/detailed-object"; +import { machineColumns, schemaContextOf, type DetailedObject } from "@/lib/db/detailed-object"; import { detectSensitiveColumnsFromConfig, loadMaskingConfig, @@ -60,7 +60,6 @@ interface DataProfilerProps { tablePath: readonly string[]; tableSchema: DetailedObject | null; connection: DatabaseConnection | null; - schemaContext?: string; databaseType?: string; /** Optional API adapter: when provided, bypasses the built-in /api/db/profile fetch. */ onProfile?: (params: { connectionId: string; tablePath: readonly string[] }) => Promise; @@ -77,32 +76,12 @@ interface DataProfilerProps { userRole?: string; } -function resolveTableSchemaSnippet(schemaContext: string | undefined, tableSchema: DetailedObject | null): string { - if (schemaContext) { - try { - const parsed = JSON.parse(schemaContext); - if (Array.isArray(parsed) && parsed.length === 1) { - return schemaContext; - } - } catch { - // Non-JSON - } - } - if (!tableSchema) return ""; - const cleanTable = { - ...tableSchema, - columns: tableSchema.columns?.filter((c) => (c as { provenance?: unknown }).provenance !== "sampled") ?? [], - }; - return JSON.stringify(cleanTable); -} - export function DataProfiler({ isOpen, onClose, tablePath, tableSchema, connection, - schemaContext, databaseType, onProfile, onDescribeSchema, @@ -165,7 +144,7 @@ export function DataProfiler({ setAiError(null); try { // The profile table lists every column, but the summary goes to a model, which never sees a column the - // engine only inferred from sampled data: `machineColumns` decides which, as for `schemaContext` below. + // engine only inferred from sampled data: `machineColumns` decides which, as it does in `schemaContextOf` below. const declared = new Set(machineColumns(tableSchema?.columns ?? [])); const withheld = new Set( (tableSchema?.columns ?? []).filter((column) => !declared.has(column)).map((column) => column.name), @@ -186,7 +165,7 @@ export function DataProfiler({ ) .join("\n"); - const tableSchemaSnippet = resolveTableSchemaSnippet(schemaContext, tableSchema); + const tableSchemaSnippet = tableSchema ? schemaContextOf([tableSchema]) : ""; let fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${tableSchemaSnippet}`; if (fullSchemaContext.length > MAX_SCHEMA_CONTEXT_CHARS) { fullSchemaContext = fullSchemaContext.slice(0, MAX_SCHEMA_CONTEXT_CHARS); diff --git a/src/components/studio/StudioModals.tsx b/src/components/studio/StudioModals.tsx index c24d9eb9e..c4a7b6956 100644 --- a/src/components/studio/StudioModals.tsx +++ b/src/components/studio/StudioModals.tsx @@ -164,7 +164,6 @@ export function StudioModals({ tablePath={profilerPath ?? []} tableSchema={objectAtPath(schema, profilerPath)} connection={activeConnection} - schemaContext={schemaContext} databaseType={databaseType} maskingConfig={profilerMasking.config} maskingEnabled={profilerMasking.enabled} diff --git a/src/lib/db/detailed-object.ts b/src/lib/db/detailed-object.ts index e15e1cab9..853d2058a 100644 --- a/src/lib/db/detailed-object.ts +++ b/src/lib/db/detailed-object.ts @@ -200,7 +200,8 @@ export function sampledMark(column: Pick): string { * The ranges are the connection's own grants, and a piece can be a single key, so they are the * one part of a schema entry that may name a key. They travel with the entry for the two * generators alone, and this string is what the query editor's completions, the docs, the safety - * review, the profiler and the explain view read, the last two posting it to the model whole. + * review and the explain view read, the explain view posting it to the model whole. The profiler + * posts only the object it profiled, run through this function on its own. * * A REPLACER rather than a projection, and the choice is what keeps every other engine's context * byte for byte the `JSON.stringify(schema)` it has always been: the replacer answers every other diff --git a/tests/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 20cd0f452..127954753 100644 --- a/tests/components/DataProfiler.test.tsx +++ b/tests/components/DataProfiler.test.tsx @@ -66,7 +66,6 @@ function createDefaultProps(overrides: Partial[0 tablePath: ["app", "users"], tableSchema: mockUsersTable, connection: mockPostgresConnection, - schemaContext: "", databaseType: "postgres", // Masking off unless a test turns it on, so the values below read as the engine sent them. maskingConfig: DEFAULT_MASKING_CONFIG, @@ -470,146 +469,6 @@ describe("DataProfiler", () => { globalThis.fetch = originalFetch; }); - test("sends only the profiled table schema when multiple schemas contain a table with the same name", async () => { - const aiText = "Profiled summary for app.users table."; - restoreGlobalFetch(); - - const originalFetch = globalThis.fetch; - let capturedBody: { schemaContext: string } | null = null; - - const publicUsersTable = { - name: "users", - kind: "table" as const, - path: ["public", "users"], - rowCount: 50, - indexes: [], - columns: [ - { name: "public_id", type: "integer", nullable: false, isPrimary: true }, - { name: "public_profile", type: "text", nullable: true, isPrimary: false }, - ], - }; - - const appUsersTable = { - name: "users", - kind: "table" as const, - path: ["app", "users"], - rowCount: 100, - indexes: [], - columns: [ - { name: "app_id", type: "integer", nullable: false, isPrimary: true }, - { name: "app_role", type: "varchar(50)", nullable: false, isPrimary: false }, - ], - }; - - // rawSchemaContext contains both schemas and exceeds 30,000 characters - const largePadding = Array.from({ length: 60 }, (_, i) => ({ - name: `other_table_${i}`, - columns: Array.from({ length: 20 }, (_, j) => ({ - name: `col_${j}`, - type: "varchar(255)", - })), - })); - const rawSchemaContext = JSON.stringify([publicUsersTable, appUsersTable, ...largePadding]); - expect(rawSchemaContext.length).toBeGreaterThan(30_000); - - globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { - const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; - const pathname = new URL(url, "http://localhost:3000").pathname; - - if (pathname.includes("/api/db/profile")) { - return new Response(JSON.stringify(mockProfileResponse), { - status: 200, - headers: { "content-type": "application/json" }, - }); - } - - if (pathname.includes("/api/ai/describe-schema")) { - const bodyStr = typeof init?.body === "string" ? init.body : await (input as Request).text(); - capturedBody = JSON.parse(bodyStr); - const stream = new ReadableStream({ - start(controller) { - controller.enqueue(new TextEncoder().encode(aiText)); - controller.close(); - }, - }); - return new Response(stream, { - status: 200, - headers: { "content-type": "text/plain" }, - }); - } - - return new Response("Not found", { status: 404 }); - }) as unknown as typeof fetch; - - const props = createDefaultProps({ - tablePath: ["app", "users"], - tableSchema: appUsersTable, - schemaContext: rawSchemaContext, - }); - const { container } = render(); - const view = within(container); - - await waitFor(() => { - expect(view.queryByText("Column Profiles")).not.toBeNull(); - }); - - await waitFor( - () => { - expect(view.queryByText("AI Analysis")).not.toBeNull(); - }, - { timeout: 3000 }, - ); - - await waitFor( - () => { - expect(view.queryByText(aiText)).not.toBeNull(); - }, - { timeout: 3000 }, - ); - - expect(capturedBody).not.toBeNull(); - expect(capturedBody!.schemaContext.length).toBeLessThanOrEqual(30_000); - // Sends only the profiled table's columns - expect(capturedBody!.schemaContext).toContain("app_id"); - expect(capturedBody!.schemaContext).toContain("app_role"); - // Does NOT contain the other schema's columns from rawSchemaContext - expect(capturedBody!.schemaContext).not.toContain("public_id"); - expect(capturedBody!.schemaContext).not.toContain("public_profile"); - - globalThis.fetch = originalFetch; - }); - - test("preserves single-table schemaContext when provided", async () => { - restoreGlobalFetch(); - let capturedBody: { schemaContext: string } | null = null; - const singleTableContext = JSON.stringify([{ name: "single_table", columns: [] }]); - - const originalFetch = globalThis.fetch; - globalThis.fetch = mock(async (input: RequestInfo | URL, init?: RequestInit) => { - const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; - if (url.includes("/api/db/profile")) { - return new Response(JSON.stringify(mockProfileResponse), { - status: 200, - headers: { "content-type": "application/json" }, - }); - } - if (url.includes("/api/ai/describe-schema")) { - const bodyStr = typeof init?.body === "string" ? init.body : await (input as Request).text(); - capturedBody = JSON.parse(bodyStr); - return new Response("summary", { status: 200 }); - } - return new Response("Not found", { status: 404 }); - }) as unknown as typeof fetch; - - const props = createDefaultProps({ schemaContext: singleTableContext }); - render(); - await waitFor(() => { - expect(capturedBody).not.toBeNull(); - }); - expect(capturedBody!.schemaContext).toContain(singleTableContext); - globalThis.fetch = originalFetch; - }); - test("displays AI error when describe-schema route refuses", async () => { const originalFetch = globalThis.fetch; const refusalError = "The request is too large for the configured model's context."; @@ -1400,7 +1259,7 @@ describe("DataProfiler and a column the engine only inferred from sampled data", return { ok: false, status: 500, json: { error: "AI not configured" } }; }, }); - render(); + render(); await waitFor(() => { expect(bodies).toHaveLength(1); }); @@ -1412,7 +1271,9 @@ describe("DataProfiler and a column the engine only inferred from sampled data", mockGlobalFetch({ "/api/db/profile": { ok: true, json: mockProfileResponse } }); const onDescribeSchema = mock(async () => "Adapter AI summary"); render( - , + , ); await waitFor(() => { expect(onDescribeSchema).toHaveBeenCalledTimes(1); @@ -1443,7 +1304,6 @@ describe("DataProfiler and a column the engine only inferred from sampled data", createDefaultProps({ tablePath: ["articles"], tableSchema: sampledSchema[0], - schemaContext: schemaContextOf(sampledSchema), ...overrides, }); @@ -1481,6 +1341,35 @@ describe("DataProfiler and a column the engine only inferred from sampled data", }); }); +describe("DataProfiler and the schema it hands a model", () => { + afterEach(() => { + cleanup(); + restoreGlobalFetch(); + }); + + // The profiled object goes through `schemaContextOf`, the projection every AI panel's schema takes, so what that + // projection withholds (a sampled column, a read range kept for the generators) stays out of this request too. + test("the AI summary request carries the profiled object as schemaContextOf projects it", async () => { + const RANGE_MARKER = "a-read-range-the-model-never-sees"; + const tableSchema = { ...mockUsersTable, readRanges: [{ prefix: RANGE_MARKER }] }; + const bodies: string[] = []; + mockGlobalFetch({ + "/api/db/profile": { ok: true, json: mockProfileResponse }, + "/api/ai/describe-schema": async (req) => { + bodies.push(await req.text()); + return { ok: false, status: 500, json: { error: "AI not configured" } }; + }, + }); + render(); + await waitFor(() => { + expect(bodies).toHaveLength(1); + }); + const { schemaContext } = JSON.parse(bodies[0]) as { schemaContext: string }; + expect(schemaContext.endsWith(`\n\nSchema:\n${schemaContextOf([tableSchema])}`)).toBe(true); + expect(schemaContext).not.toContain(RANGE_MARKER); + }); +}); + // ============================================================================= // The profiler masks with the grid's decision (#1421) //