fix(studio): auto-retry notebook updates on stale/invalid conflicts (#49323)

## 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)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

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

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

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(
-      <NotebookProposalRenderer
-        mode="update"
-        state="approval-requested"
-        confirmState="approval-requested"
-        input={{
-          id: NOTEBOOK_ID,
-          expected_updated_at: '2024-01-01T00:00:00.000Z',
-          operations: [{ _tag: 'delete_cell', cell_id: 'cell-1' }],
-        }}
-        output={undefined}
-        onApprove={onApprove}
-        onDeny={vi.fn()}
-      />
-    )
-
-    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()

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 (
-      <NotebookConfirm
-        mode="update"
-        confirmState={confirmState}
-        confirmLabel="Refresh"
-        confirmLabelLoading="Refreshing..."
-        extraLoading={isFetching}
-        onDeny={onDeny}
-        onApprove={() => refetch()}
-      >
-        <div className="p-3">
-          <Admonition
-            type="warning"
-            title="This notebook changed since the assistant planned this update"
-            description={`"${notebook.name}" was updated after the assistant read it. Refresh to see the latest version before deciding.`}
-          />
-        </div>
-      </NotebookConfirm>
-    )
-  }
-
   const diff = deriveNotebookDiff(toWireNotebook(notebook.content), parsedInput.data.operations)
 
   if (!diff.success) {

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

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>): 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<ToolUIPart>),
+      ]),
+    ]
+
+    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' }])
+  })
+})

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

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

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<typeof notebookToolErrorMetadataSchema>
+
+const notebookToolErrorSchema = notebookToolErrorMetadataSchema.extend({
+  tag: z.literal('notebook_tool_error'),
+  message: z.string(),
+})
+
+export type EncodedNotebookToolError = z.infer<typeof notebookToolErrorSchema>
+
+/** 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

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