diff --git a/src/app/api/ai/describe-schema/route.ts b/src/app/api/ai/describe-schema/route.ts index 8141b5e3b7..b3cff38738 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"; @@ -10,8 +11,14 @@ 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 (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 }); } const provider = await createLLMProvider(); diff --git a/src/components/DataProfiler.tsx b/src/components/DataProfiler.tsx index 52a08c85b4..47ff401e37 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, @@ -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, @@ -59,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; @@ -82,7 +82,6 @@ export function DataProfiler({ tablePath, tableSchema, connection, - schemaContext, databaseType, onProfile, onDescribeSchema, @@ -93,6 +92,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,9 +141,10 @@ 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. + // 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), @@ -164,7 +165,11 @@ export function DataProfiler({ ) .join("\n"); - const fullSchemaContext = `Table: ${tableName} (${data.totalRows} rows)\n\nColumn Profiles:\n${profileSummary}\n\nSchema:\n${schemaContext || ""}`; + 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); + } if (onDescribeSchema) { // Platform adapter: use callback instead of fetch @@ -182,7 +187,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 +204,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 +263,7 @@ export function DataProfiler({ return () => { setProfile(null); setAiSummary(""); + setAiError(null); setError(null); }; // eslint-disable-next-line react-hooks/exhaustive-deps @@ -523,7 +533,7 @@ export function DataProfiler({ {/* AI Summary */} - {(aiSummary || isAiLoading) && ( + {(aiSummary || isAiLoading || aiError) && (
@@ -533,6 +543,7 @@ export function DataProfiler({ {aiSummary && (
{aiSummary}
)} + {aiError &&
{aiError}
}
)} diff --git a/src/components/DatabaseDocs.tsx b/src/components/DatabaseDocs.tsx index a1d0c2231e..75c1dfd635 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/components/studio/StudioModals.tsx b/src/components/studio/StudioModals.tsx index c24d9eb9eb..c4a7b69567 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/api/errors.ts b/src/lib/api/errors.ts index 641f2dee22..bc1cd41242 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 request 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/db/detailed-object.ts b/src/lib/db/detailed-object.ts index e15e1cab93..853d2058ac 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/src/lib/llm/types.ts b/src/lib/llm/types.ts index 77990c9c29..40e37c3e74 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 3f64214b97..ed4c0361dd 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 ─────────────────────────────────────────────────────────── @@ -74,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( @@ -298,4 +307,61 @@ 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."); + 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 () => { + 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/api/ai/explain.test.ts b/tests/api/ai/explain.test.ts index 1bd1d88cd4..54306a3603 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 a743a47816..9e57668982 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/components/DataProfiler.test.tsx b/tests/components/DataProfiler.test.tsx index 4f414cc5ec..127954753c 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,6 +469,156 @@ describe("DataProfiler", () => { 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; + }); + + 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 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; + 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({ tableSchema: massiveTable }); + 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 () => { @@ -604,7 +753,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); @@ -634,7 +783,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(); @@ -1110,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); }); @@ -1122,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); @@ -1153,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, }); @@ -1191,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) // diff --git a/tests/components/DatabaseDocs.test.tsx b/tests/components/DatabaseDocs.test.tsx index ee50d4fe3a..a4162c9a8a 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/security/route-auth.test.ts b/tests/security/route-auth.test.ts index 377fa897f5..4aec84d8f2 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"); diff --git a/tests/unit/api-errors.test.ts b/tests/unit/api-errors.test.ts index 76d1088555..73e06102b8 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 request 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 cfc5f527c0..51447e10ba 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");