From ebd616fa901db16d4dd1272effa238cd022f25c9 Mon Sep 17 00:00:00 2001 From: Charis <26616127+charislam@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:08:19 -0400 Subject: [PATCH] fix(studio): auto-retry notebook updates on stale/invalid conflicts (#49323) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - **Removed dead client-side refresh UI** in `NotebookProposalRenderer.tsx` and its test — the diff preview is always computed from live data, so the check was redundant with the tool's server-side re-validation - **Added typed `NotebookToolError`** in `notebook-tools.ts` with structured metadata (`{ exposeToAssistant: boolean }`) validated by a zod schema with a literal discriminant tag (`tag: 'notebook_tool_error'`) — tracks the two retryable failures: staleness conflict and invalid operations (unknown cell id) - **Encoded errors in `generate-v4.ts` onError** — the one place in the pipeline that holds the live `Error` before it becomes a string in the persisted message - **Extracted and fixed message history filter** into new `generate-assistant-response.utils.ts` — any tool-error whose `errorText` decodes against the `NotebookToolError` schema is let through (with `errorText` rewritten to plain prose so the model sees the message, not JSON), while other errors stay filtered as before Net effect: the assistant detects the specific, actionable rejection reason and retries on its own with no dead button or human intervention needed. ## Test plan - Existing unit tests in `NotebookProposalRenderer.test.tsx` pass (dead button test removed) - New unit tests in `notebook-tools.test.ts` cover encode/decode round-trips and error discrimination - New unit tests in `generate-assistant-response.utils.test.ts` cover message history filtering with all error states - `pnpm typecheck` is clean - `pnpm --filter studio run lint:ratchet` passes (no new ESLint warnings) ## Summary by CodeRabbit * **New Features** * Notebook update errors now provide clearer, structured explanations to the AI assistant. * Assistant responses preserve relevant notebook error details while filtering invalid or temporary tool states. * **Bug Fixes** * Improved handling of stale notebook revisions and invalid notebook update operations. * Notebook proposal rendering proceeds without an unnecessary refresh step. * **Tests** * Expanded coverage for notebook errors, message filtering, serialization, and error handling. --- .../NotebookProposalRenderer.test.tsx | 27 ---- .../NotebookProposalRenderer.tsx | 26 ---- .../lib/ai/generate-assistant-response.ts | 34 +---- .../generate-assistant-response.utils.test.ts | 125 ++++++++++++++++++ .../ai/generate-assistant-response.utils.ts | 50 +++++++ .../lib/ai/tools/notebook-tools.test.ts | 91 +++++++++---- apps/studio/lib/ai/tools/notebook-tools.ts | 57 +++++++- apps/studio/pages/api/ai/sql/generate-v4.ts | 4 + 8 files changed, 303 insertions(+), 111 deletions(-) create mode 100644 apps/studio/lib/ai/generate-assistant-response.utils.test.ts create mode 100644 apps/studio/lib/ai/generate-assistant-response.utils.ts diff --git a/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.test.tsx b/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.test.tsx index 5d591f90644..df86f4795e2 100644 --- a/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.test.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.test.tsx @@ -98,33 +98,6 @@ describe('NotebookProposalRenderer', () => { expect(onApprove).toHaveBeenCalledTimes(1) }) - it('warns and withholds the diff when the notebook changed since expected_updated_at', async () => { - const onApprove = vi.fn() - mockContentItem(mockNotebookRow({ updated_at: '2024-06-01T00:00:00.000Z' })) - - render( - - ) - - expect( - await screen.findByText('This notebook changed since the assistant planned this update') - ).toBeInTheDocument() - expect(screen.getByRole('button', { name: 'Refresh' })).toBeInTheDocument() - expect(onApprove).not.toHaveBeenCalled() - }) - it('falls back to a raw-input admonition without dropping the confirm footer on a parse failure', async () => { const user = userEvent.setup() const onApprove = vi.fn() diff --git a/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx index 0f2c665de51..e59c5d6aa06 100644 --- a/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx @@ -238,8 +238,6 @@ function UpdateNotebookProposal({ isLoading, isError, error, - refetch, - isFetching, } = useNotebookQuery( { projectRef: ref, id: parsedInput.success ? parsedInput.data.id : undefined }, { enabled: parsedInput.success } @@ -284,30 +282,6 @@ function UpdateNotebookProposal({ ) } - const isStale = notebook.updated_at !== parsedInput.data.expected_updated_at - - if (isStale) { - return ( - refetch()} - > -
- -
- - ) - } - const diff = deriveNotebookDiff(toWireNotebook(notebook.content), parsedInput.data.operations) if (!diff.success) { diff --git a/apps/studio/lib/ai/generate-assistant-response.ts b/apps/studio/lib/ai/generate-assistant-response.ts index 9f457df4bc5..1488882f504 100644 --- a/apps/studio/lib/ai/generate-assistant-response.ts +++ b/apps/studio/lib/ai/generate-assistant-response.ts @@ -2,7 +2,6 @@ import * as ai from 'ai' import { convertToModelMessages, isStepCount, - isToolUIPart, type LanguageModel, type ModelMessage, type SystemModelMessage, @@ -16,6 +15,7 @@ import type { AssistantEvalInput } from '@/evals/scorer' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' import { buildAssistantContextMessages, NO_SCHEMA_ACCESS_MESSAGE } from '@/lib/ai/assistant-context' import { IS_TRACING_ENABLED } from '@/lib/ai/braintrust-logger' +import { prepareMessagesForModel } from '@/lib/ai/generate-assistant-response.utils' import { CHAT_PROMPT, GENERAL_PROMPT, @@ -23,7 +23,6 @@ import { NOTEBOOKS_PROMPT, SECURITY_PROMPT, } from '@/lib/ai/prompts' -import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer' const { streamText: tracedStreamText } = wrapAISDK(ai) @@ -74,36 +73,7 @@ export async function generateAssistantResponse({ const shouldTrace = allowTracing ?? IS_TRACING_ENABLED const run = async (span?: Span) => { - // Only returns last 7 messages - // Filters out tools with invalid states - // Filters out tool outputs based on opt-in level - const messages = (rawMessages || []).slice(-7).map((msg) => { - if (msg && msg.role === 'assistant' && 'results' in msg) { - const cleanedMsg = { ...msg } - delete cleanedMsg.results - return cleanedMsg - } - if (msg && msg.role === 'assistant' && msg.parts) { - const cleanedParts = msg.parts - .filter((part) => { - if (isToolUIPart(part)) { - const invalidStates = [ - 'input-streaming', - 'input-available', - 'approval-requested', - 'output-error', - ] - return !invalidStates.includes(part.state) - } - return true - }) - .map((part) => { - return sanitizeMessagePart(part, aiOptInLevel) - }) - return { ...msg, parts: cleanedParts } - } - return msg - }) + const messages = prepareMessagesForModel(rawMessages, aiOptInLevel) const schemasString = aiOptInLevel !== 'disabled' && getSchemas diff --git a/apps/studio/lib/ai/generate-assistant-response.utils.test.ts b/apps/studio/lib/ai/generate-assistant-response.utils.test.ts new file mode 100644 index 00000000000..7e5d23429df --- /dev/null +++ b/apps/studio/lib/ai/generate-assistant-response.utils.test.ts @@ -0,0 +1,125 @@ +import type { ToolUIPart, UIMessage } from 'ai' +import { describe, expect, it } from 'vitest' + +import { prepareMessagesForModel } from './generate-assistant-response.utils' +import { encodeNotebookToolError, NotebookToolError } from './tools/notebook-tools' + +function assistantMessage(parts: UIMessage['parts']): UIMessage { + return { id: 'msg-1', role: 'assistant', parts } +} + +function toolPart(overrides: Partial): ToolUIPart { + return { + type: 'tool-update_notebook', + toolCallId: 'call-1', + state: 'output-available', + input: {}, + ...overrides, + } as ToolUIPart +} + +describe('prepareMessagesForModel', () => { + it('filters out a plain output-error tool part', () => { + const messages = [ + assistantMessage([toolPart({ state: 'output-error', errorText: 'Network error' })]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([]) + }) + + it('keeps a stale-update_notebook conflict and rewrites errorText to the plain message', () => { + const error = new NotebookToolError('Notebook changed since expected_updated_at', { + exposeToAssistant: true, + }) + const messages = [ + assistantMessage([ + toolPart({ state: 'output-error', errorText: encodeNotebookToolError(error)! }), + ]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([ + toolPart({ state: 'output-error', errorText: 'Notebook changed since expected_updated_at' }), + ]) + }) + + it('still filters out an unrelated update_notebook output-error', () => { + const messages = [ + assistantMessage([ + toolPart({ state: 'output-error', errorText: 'Unexpected upstream failure' }), + ]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([]) + }) + + it('still filters out JSON-shaped output-error text lacking the notebook_tool_error tag', () => { + const messages = [ + assistantMessage([ + toolPart({ + state: 'output-error', + errorText: JSON.stringify({ + exposeToAssistant: true, + message: 'not a real notebook error', + }), + }), + ]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([]) + }) + + it('exposes any tool part carrying a validly-tagged NotebookToolError, not just update_notebook', () => { + const error = new NotebookToolError('Notebook changed since expected_updated_at', { + exposeToAssistant: true, + }) + const messages = [ + assistantMessage([ + toolPart({ + type: 'tool-execute_sql', + state: 'output-error', + errorText: encodeNotebookToolError(error)!, + }), + ]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([ + toolPart({ + type: 'tool-execute_sql', + state: 'output-error', + errorText: 'Notebook changed since expected_updated_at', + }), + ]) + }) + + it('still filters out input-streaming, input-available, and approval-requested parts', () => { + const messages = [ + assistantMessage([ + toolPart({ state: 'input-streaming' }), + toolPart({ state: 'input-available' }), + toolPart({ state: 'approval-requested', approval: { id: 'a1' } } as Partial), + ]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([]) + }) + + it('keeps non-tool parts untouched', () => { + const messages = [assistantMessage([{ type: 'text', text: 'hello' }])] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([{ type: 'text', text: 'hello' }]) + }) +}) diff --git a/apps/studio/lib/ai/generate-assistant-response.utils.ts b/apps/studio/lib/ai/generate-assistant-response.utils.ts new file mode 100644 index 00000000000..3d2851656ec --- /dev/null +++ b/apps/studio/lib/ai/generate-assistant-response.utils.ts @@ -0,0 +1,50 @@ +import { isToolUIPart, type DynamicToolUIPart, type ToolUIPart, type UIMessage } from 'ai' + +import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' +import { decodeNotebookToolError } from '@/lib/ai/tools/notebook-tools' +import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer' + +const INVALID_TOOL_STATES = [ + 'input-streaming', + 'input-available', + 'approval-requested', + 'output-error', +] + +/** + * Notebook tool errors are opted into model visibility explicitly (NotebookToolError, + * encoded in generate-v4.ts's onError). A successful parse against the schema — including + * its literal `tag` — is proof enough that this was a NotebookToolError; every other + * output-error stays hidden, same as before. Rewrites errorText back to the plain message + * so the model sees prose, not JSON. + */ +function exposedNotebookErrorPart( + part: ToolUIPart | DynamicToolUIPart +): ToolUIPart | DynamicToolUIPart | null { + if (part.state !== 'output-error') return null + const decoded = decodeNotebookToolError(part.errorText) + if (!decoded?.exposeToAssistant) return null + return { ...part, errorText: decoded.message } +} + +/** Trims history to the last 7 messages and strips tool parts the model shouldn't see. */ +export function prepareMessagesForModel(rawMessages: UIMessage[], aiOptInLevel: AiOptInLevel) { + return (rawMessages || []).slice(-7).map((msg) => { + if (msg && msg.role === 'assistant' && 'results' in msg) { + const cleanedMsg = { ...msg } + delete cleanedMsg.results + return cleanedMsg + } + if (msg && msg.role === 'assistant' && msg.parts) { + const cleanedParts = msg.parts.flatMap((part) => { + if (!isToolUIPart(part)) return [part] + if (!INVALID_TOOL_STATES.includes(part.state)) + return [sanitizeMessagePart(part, aiOptInLevel)] + const exposed = exposedNotebookErrorPart(part) + return exposed ? [sanitizeMessagePart(exposed, aiOptInLevel)] : [] + }) + return { ...msg, parts: cleanedParts } + } + return msg + }) +} diff --git a/apps/studio/lib/ai/tools/notebook-tools.test.ts b/apps/studio/lib/ai/tools/notebook-tools.test.ts index 2650d326ca0..fb692afbf20 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.test.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.test.ts @@ -2,7 +2,12 @@ import { components } from 'api-types' import { HttpResponse } from 'msw' import { describe, expect, it } from 'vitest' -import { getNotebookTools } from './notebook-tools' +import { + decodeNotebookToolError, + encodeNotebookToolError, + getNotebookTools, + NotebookToolError, +} from './notebook-tools' import type { AgentNotebook } from '@/data/content/notebooks/notebook-schema' import { addAPIMock, type APIErrorBody } from '@/tests/lib/msw' @@ -432,40 +437,80 @@ describe('ai/tools/notebook-tools', () => { expect(result).toEqual({ id: 'notebook-1', name: 'Signup funnel' }) }) - it('should throw a descriptive error instead of PUTting when an operation targets an unknown cell id', async () => { + it('should throw a descriptive, assistant-exposable error instead of PUTting when an operation targets an unknown cell id', async () => { mockGetNotebook() const tools = getNotebookTools({ projectRef: 'test-project' }) if (!tools.update_notebook.execute) throw new Error('execute is undefined') - await expect( - tools.update_notebook.execute( - { - id: 'notebook-1', - expected_updated_at: '2026-01-01T00:00:00.000Z', - operations: [{ _tag: 'delete_cell', cell_id: 'missing-cell' }], - }, - { toolCallId: 'test', messages: [], context: {} } - ) - ).rejects.toThrow('No cell with id "missing-cell"') + const execute = tools.update_notebook.execute( + { + id: 'notebook-1', + expected_updated_at: '2026-01-01T00:00:00.000Z', + operations: [{ _tag: 'delete_cell', cell_id: 'missing-cell' }], + }, + { toolCallId: 'test', messages: [], context: {} } + ) + + await expect(execute).rejects.toThrow('No cell with id "missing-cell"') + await expect(execute).rejects.toBeInstanceOf(NotebookToolError) + await expect(execute).rejects.toMatchObject({ metadata: { exposeToAssistant: true } }) }) - it('should throw instead of PUTting when the notebook changed since expected_updated_at', async () => { + it('should throw an assistant-exposable error instead of PUTting when the notebook changed since expected_updated_at', async () => { mockGetNotebook() const tools = getNotebookTools({ projectRef: 'test-project' }) if (!tools.update_notebook.execute) throw new Error('execute is undefined') - await expect( - tools.update_notebook.execute( - { - id: 'notebook-1', - expected_updated_at: '2025-12-31T00:00:00.000Z', - operations: [{ _tag: 'delete_cell', cell_id: 'cell-3' }], - }, - { toolCallId: 'test', messages: [], context: {} } - ) - ).rejects.toThrow(/changed since expected_updated_at/) + const execute = tools.update_notebook.execute( + { + id: 'notebook-1', + expected_updated_at: '2025-12-31T00:00:00.000Z', + operations: [{ _tag: 'delete_cell', cell_id: 'cell-3' }], + }, + { toolCallId: 'test', messages: [], context: {} } + ) + + await expect(execute).rejects.toThrow(/changed since expected_updated_at/) + await expect(execute).rejects.toBeInstanceOf(NotebookToolError) + await expect(execute).rejects.toMatchObject({ metadata: { exposeToAssistant: true } }) + }) + }) + + describe('encodeNotebookToolError / decodeNotebookToolError', () => { + it('round-trips a NotebookToolError through JSON', () => { + const error = new NotebookToolError('Notebook changed since expected_updated_at', { + exposeToAssistant: true, + }) + + const encoded = encodeNotebookToolError(error) + expect(encoded).not.toBeNull() + + const decoded = decodeNotebookToolError(encoded!) + expect(decoded).toEqual({ + exposeToAssistant: true, + tag: 'notebook_tool_error', + message: 'Notebook changed since expected_updated_at', + }) + }) + + it('does not encode a plain Error', () => { + expect(encodeNotebookToolError(new Error('boom'))).toBeNull() + }) + + it('does not decode a plain error message', () => { + expect(decodeNotebookToolError('boom')).toBeNull() + }) + + it('does not decode unrelated JSON', () => { + expect(decodeNotebookToolError(JSON.stringify({ foo: 'bar' }))).toBeNull() + }) + + it('does not decode JSON that merely looks like a NotebookToolError but lacks the tag', () => { + expect( + decodeNotebookToolError(JSON.stringify({ exposeToAssistant: true, message: 'boom' })) + ).toBeNull() }) }) }) diff --git a/apps/studio/lib/ai/tools/notebook-tools.ts b/apps/studio/lib/ai/tools/notebook-tools.ts index 1ec84bdb88d..2c4498d459e 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.ts @@ -24,6 +24,54 @@ export type NotebookToolsContext = { authorization?: string } +const notebookToolErrorMetadataSchema = z.object({ + exposeToAssistant: z.boolean(), +}) + +export type NotebookToolErrorMetadata = z.infer + +const notebookToolErrorSchema = notebookToolErrorMetadataSchema.extend({ + tag: z.literal('notebook_tool_error'), + message: z.string(), +}) + +export type EncodedNotebookToolError = z.infer + +/** Thrown by update_notebook for failures the assistant can act on by retrying. */ +export class NotebookToolError extends Error { + readonly metadata: NotebookToolErrorMetadata + + constructor(message: string, metadata: NotebookToolErrorMetadata) { + super(message) + this.name = 'NotebookToolError' + this.metadata = notebookToolErrorMetadataSchema.parse(metadata) + } +} + +/** Called from the stream's `onError` (generate-v4.ts), which still has the live Error. */ +export function encodeNotebookToolError(error: unknown): string | null { + if (!(error instanceof NotebookToolError)) return null + return JSON.stringify( + notebookToolErrorSchema.parse({ + ...error.metadata, + tag: 'notebook_tool_error', + message: error.message, + }) + ) +} + +/** Called from the history filter (generate-assistant-response.ts) on a persisted errorText. */ +export function decodeNotebookToolError(errorText: string): EncodedNotebookToolError | null { + let parsed: unknown + try { + parsed = JSON.parse(errorText) + } catch { + return null + } + const result = notebookToolErrorSchema.safeParse(parsed) + return result.success ? result.data : null +} + export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { const { projectRef, authorization } = ctx const authHeaders = authorization ? { Authorization: authorization } : undefined @@ -153,8 +201,9 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { const notebook = await getNotebook({ projectRef, id }, undefined, authHeaders) if (notebook.updated_at !== expected_updated_at) { - throw new Error( - `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue update_notebook against the current content.` + throw new NotebookToolError( + `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue update_notebook against the current content.`, + { exposeToAssistant: true } ) } @@ -165,7 +214,9 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { const result = applyNotebookOperations(wireNotebook, operations) if (!result.success) { - throw new Error(describeNotebookOperationError(result.error)) + throw new NotebookToolError(describeNotebookOperationError(result.error), { + exposeToAssistant: true, + }) } // Same promotion as create_notebook above, inlined here for the same auditability diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index d8900ab88db..1b056af7c64 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -26,6 +26,7 @@ import { type AssistantModelId, } from '@/lib/ai/model.utils' import { getTools } from '@/lib/ai/tools' +import { encodeNotebookToolError } from '@/lib/ai/tools/notebook-tools' import { apiWrapper } from '@/lib/api/apiWrapper' import { executeQuery } from '@/lib/api/self-hosted/query' import { getURL } from '@/lib/helpers' @@ -250,6 +251,9 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw onError: (error) => { console.error('Assistant stream error:', error) + const encoded = encodeNotebookToolError(error) + if (encoded !== null) return encoded + if (error == null) { return 'unknown error' }