fix(ai): bound schema context length and handle context overflow errors

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();

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<ProfileData>;
@@ -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<ProfileData | null>(null);
   const [aiSummary, setAiSummary] = useState("");
+  const [aiError, setAiError] = useState<string | null>(null);
   const [isAiLoading, setIsAiLoading] = useState(false);
   const [error, setError] = useState<string | null>(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({
                 </div>
 
                 {/* AI Summary */}
-                {(aiSummary || isAiLoading) && (
+                {(aiSummary || isAiLoading || aiError) && (
                   <div className="bg-hue-cyan-tint/5 border border-hue-cyan-tint/10 rounded-lg p-4">
                     <div className="flex items-center gap-2 mb-2">
                       <Sparkles strokeWidth={1.5} className="w-3.5 h-3.5 text-hue-cyan" />
@@ -533,6 +543,7 @@ export function DataProfiler({
                     {aiSummary && (
                       <div className="text-xs text-fg-tertiary leading-relaxed whitespace-pre-wrap">{aiSummary}</div>
                     )}
+                    {aiError && <div className="text-xs text-danger leading-relaxed">{aiError}</div>}
                   </div>
                 )}
               </>

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);
         }

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}

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 },

src/lib/db/detailed-object.ts

@@ -200,7 +200,8 @@ export function sampledMark(column: Pick<ColumnSchema, "provenance">): 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

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.
  *

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);
+  });
 });

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(

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(