Improve AI SQL correction and stacked charts

Author: sorenbs

Architecture/sql-ai-generation.md

@@ -62,6 +62,13 @@ This architecture governs:
 - Successful responses MAY surface the rationale inline to explain what was generated.
 - Successful responses MUST carry the visualization decision forward so the next manual execution of that same AI-generated SQL can auto-generate a chart for graph-worthy results.
 
+## Database Error Correction Contract
+
+- If the user manually executes SQL that was generated by AI and the database returns an execution error, Studio MUST feed the original natural-language request, the failed SQL, and the database error message back through the same `sql-generation` LLM transport.
+- The correction request MUST preserve the original explicit-execution contract: corrected SQL replaces the editor contents, updates the rationale and visualization decision, and focuses the editor, but Studio MUST NOT auto-run the corrected SQL.
+- Database-error correction MUST only run for editor contents that still match the pending AI-generated SQL. Manually edited SQL errors MUST stay as normal inline query errors and MUST NOT trigger hidden AI calls.
+- If database-error correction fails, Studio MUST keep the database error visible and surface the AI correction failure inline without discarding the user's current SQL.
+
 ## Embedder and Demo Contract
 
 - Embedders own the actual provider call through one shared `llm({ task, prompt })` hook.
@@ -81,3 +88,5 @@ Changes to AI SQL generation MUST include tests covering:
 - preserving the visualization decision until the user manually runs the generated SQL
 - surfacing the visualization decision from the AI response
 - inline error rendering for failed AI generation
+- feeding execution errors from AI-generated SQL back into the model without auto-running the corrected SQL
+- leaving manually written SQL execution errors as normal inline query errors without invoking AI correction

Architecture/sql-result-visualization.md

@@ -57,8 +57,10 @@ This architecture governs:
   - `title`: optional short display title
   - `data`: plain JSON objects with primitive field values only
   - `xKey` plus `series[]` for `bar` and `line` charts
+  - optional `stacked: true` for `bar` charts only
   - `labelKey` plus `valueKey` for `pie` and `doughnut` charts
 - `line` charts MUST use date-like `xKey` values: ISO dates, ISO datetimes, or epoch milliseconds. Generic categorical data should use `bar`.
+- Stacked bar charts MUST use one data row per category and separate numeric series fields for each stack segment. Long/tidy SQL result rows MAY be pivoted by the AI visualization response into this Studio-owned chart config.
 - The supported chart types MUST be constrained to the known allowlist above.
 - The implementation MUST retry up to two times when the model returns malformed JSON or an invalid chart config, and each correction prompt MUST include the latest validation error.
 
@@ -78,3 +80,4 @@ Changes to SQL result visualization MUST include tests covering:
 - automatic chart generation for AI-generated SQL results that request visualization
 - replacement of the action with a mounted chart
 - reset behavior when another query execution starts
+- validation and rendering support for stacked bar chart configs

FEATURES.md

@@ -221,12 +221,13 @@ The prompt context is built from live introspection metadata, including the conc
 AI responses must satisfy a strict JSON contract with generated SQL, a short rationale, and a yes/no visualization decision, and Studio retries once if the model returns malformed JSON.
 Submitted AI requests are also stored locally in the SQL-view TanStack collection, so an empty focused prompt field can browse older requests with `ArrowUp` / `ArrowDown` as placeholder-only previews before committing one back into the input for editing.
 Provider output-limit failures are surfaced explicitly and can feed into the next JSON-correction prompt instead of showing up as a vague parse failure. The visualization decision from AI generation is also preserved so the later manual run can still auto-chart graph-worthy results.
+If AI-generated SQL fails when the user runs it, Studio sends the original request, failed SQL, and database error back to the model, then replaces the editor with a corrected query without auto-running it.
 
 ## AI SQL Result Visualization
 
 When SQL query results are visible and `llm` is configured, Studio can also turn the returned rows into an in-grid Bklit visualization.
 The visualization uses a minimal summary-row trigger labeled `Visualize data with AI`, right-aligned beside the query result count, and mounts the generated chart above the SQL result headers inside the shared scrollable grid without a regenerate control.
-Studio sends the executed SQL, the concrete database engine, and the full result row set to the model, and when the result came from `Generate SQL with AI` it also includes the original natural-language request for extra visualization context. The model is asked for a strict Bklit chart config for bar, line, pie, or doughnut charts, and Studio validates that schema before rendering.
+Studio sends the executed SQL, the concrete database engine, and the full result row set to the model, and when the result came from `Generate SQL with AI` it also includes the original natural-language request for extra visualization context. The model is asked for a strict Bklit chart config for bar, line, pie, doughnut, or stacked bar charts, and Studio validates that schema before rendering.
 Mounted charts sit inside an in-grid background band that stays tied to the visible result viewport instead of the total table width, while the chart itself stays centered and width-clamped between 300px and 1200px so wide result grids do not force giant charts.
 When SQL is generated through AI, the same model call also decides whether the expected result is graph-worthy; if it says yes, Studio auto-generates the chart after the user manually runs that generated SQL instead of waiting for a separate chart button click.
 If another query starts running, the visualization resets immediately so stale charts do not persist across changing result sets. Visualization generation also retries up to two times on malformed JSON, invalid chart configs, or explicit provider output-limit failures.

README.md

@@ -109,7 +109,7 @@ export function EmbeddedStudio() {
 - `queryInsights` is optional. Pass it only when your BFF implements the `query-insights` procedure; otherwise omit it so Studio hides the `Queries` view.
 - Studio does not render a built-in fullscreen header button. If your host needs fullscreen behavior, render that control at the host container level, as the local demo does.
 
-Studio handles prompt construction, type-aware validation, correction retries, SQL execution retries, and conversion into the normal filter, SQL, and visualization surfaces. The host transport only needs to forward the prepared request to an LLM provider and return the typed result.
+Studio handles prompt construction, type-aware validation, correction retries, database-error correction for AI-generated SQL, and conversion into the normal filter, SQL, and visualization surfaces. The host transport only needs to forward the prepared request to an LLM provider and return the typed result.
 
 ## AI Contract
 
@@ -136,7 +136,7 @@ type StudioLlmResponse =
     };
 ```
 
-Studio treats `output-limit-exceeded` as a first-class retry signal for SQL generation and visualization correction loops. All prompting and retry behavior live in Studio itself, so host implementations should stay transport-only.
+Studio treats `output-limit-exceeded` as a first-class retry signal for SQL generation and visualization correction loops. If AI-generated SQL fails after the user manually runs it, Studio also sends the failed SQL and database error back through the same `sql-generation` transport so the model can propose corrected SQL without auto-running it. All prompting and retry behavior live in Studio itself, so host implementations should stay transport-only.
 
 ## Integration Checklist
 

ui/studio/views/sql/SqlResultVisualization.tsx

@@ -212,6 +212,8 @@ function SqlResultBarChart({
       className="h-full min-h-64"
       data={config.data}
       margin={{ bottom: 38, left: 40, right: 24, top: 12 }}
+      stacked={config.stacked === true}
+      stackGap={config.stacked === true ? 1 : 0}
       xDataKey={config.xKey}
     >
       <Grid strokeDasharray="0" vertical={false} />

ui/studio/views/sql/SqlView.test.tsx

@@ -676,8 +676,9 @@ describe("SqlView", () => {
     harness.cleanup();
   });
 
-  it("surfaces query errors only after the user manually runs AI-generated SQL", async () => {
+  it("feeds AI-generated SQL execution errors back into the model without auto-running the correction", async () => {
     const badSql = "select typeof(json_col) from public.organizations limit 5;";
+    const correctedSql = "select id from public.organizations limit 5;";
     const raw: Adapter["raw"] = async (details) => {
       const error = new Error(
         "function typeof(json) does not exist",
@@ -689,10 +690,18 @@ describe("SqlView", () => {
     const studio = createStudioMock(adapter);
     const llmMock = vi
       .fn<(request: StudioLlmRequest) => Promise<string>>()
-      .mockResolvedValue(
+      .mockResolvedValueOnce(
         JSON.stringify({
           rationale: "Tried a typeof helper.",
           sql: badSql,
+          shouldGenerateVisualization: false,
+        }),
+      )
+      .mockResolvedValueOnce(
+        JSON.stringify({
+          rationale: "Uses plain selectable columns after the database error.",
+          sql: correctedSql,
+          shouldGenerateVisualization: true,
         }),
       );
     studio.llm = llmMock;
@@ -739,21 +748,75 @@ describe("SqlView", () => {
       runButton.dispatchEvent(new MouseEvent("click", { bubbles: true }));
     });
 
+    await waitFor(() => editor.value === correctedSql);
+
+    expect(llmMock).toHaveBeenCalledTimes(2);
+    expect(llmMock.mock.calls[1]?.[0]).toMatchObject({
+      task: "sql-generation",
+    });
+    expect(llmMock.mock.calls[1]?.[0].prompt).toContain(
+      "Previous SQL statement: select typeof(json_col) from public.organizations limit 5",
+    );
+    expect(llmMock.mock.calls[1]?.[0].prompt).toContain(
+      "Database error from that SQL: function typeof(json) does not exist",
+    );
+    expect(rawSpy).toHaveBeenCalledTimes(1);
+    expect(rawSpy).toHaveBeenCalledWith(
+      { sql: "select typeof(json_col) from public.organizations limit 5" },
+      expect.any(Object),
+    );
+    expect(harness.container.textContent).toContain(
+      "Uses plain selectable columns after the database error.",
+    );
+    expect(harness.container.textContent).not.toContain("Query error:");
+
+    harness.cleanup();
+  });
+
+  it("keeps normal manual SQL execution errors inline without asking AI for a correction", async () => {
+    const raw: Adapter["raw"] = async (details) => {
+      const error = new Error("relation missing_table does not exist") as AdapterError;
+      error.query = { parameters: [], sql: details.sql };
+      return [error];
+    };
+    const { adapter, rawSpy } = createAdapterMock({ raw });
+    const studio = createStudioMock(adapter);
+    const llmMock = vi.fn<(request: StudioLlmRequest) => Promise<string>>();
+    studio.llm = llmMock;
+    useStudioMock.mockReturnValue(studio);
+
+    const harness = renderSqlView();
+    const editor = harness.container.querySelector<HTMLTextAreaElement>(
+      'textarea[aria-label="SQL editor"]',
+    );
+    const runButton = [...harness.container.querySelectorAll("button")].find(
+      (button) => button.textContent?.includes("Run SQL"),
+    );
+
+    if (!editor || !runButton) {
+      throw new Error("Expected SQL editor and Run SQL button");
+    }
+
+    act(() => {
+      editor.value = "select * from missing_table";
+      mockCodeMirrorOnChange?.("select * from missing_table");
+    });
+
+    act(() => {
+      runButton.dispatchEvent(new MouseEvent("click", { bubbles: true }));
+    });
+
     await waitFor(() => {
       return (
         harness.container.textContent?.includes(
-          "function typeof(json) does not exist",
+          "relation missing_table does not exist",
         ) ?? false
       );
     });
 
-    expect(llmMock).toHaveBeenCalledTimes(1);
+    expect(llmMock).not.toHaveBeenCalled();
     expect(rawSpy).toHaveBeenCalledTimes(1);
-    expect(rawSpy).toHaveBeenCalledWith(
-      { sql: "select typeof(json_col) from public.organizations limit 5" },
-      expect.any(Object),
-    );
-    expect(harness.container.textContent).toContain("Tried a typeof helper.");
+    expect(harness.container.textContent).toContain("Query error:");
 
     harness.cleanup();
   });