diff --git a/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx b/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx
index ceccee89c71..21f9af24ffa 100644
--- a/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx
+++ b/apps/studio/components/interfaces/Explorer/ExplorerNotebookTabCoordinator.tsx
@@ -54,7 +54,7 @@ export const ExplorerNotebookTabCoordinator = () => {
         const notebookId = tab.metadata?.notebookId
         if (!ref || !notebookId) return
 
-        evictNotebookFromCaches({ queryClient, projectRef: ref, id: notebookId, mode: 'remove' })
+        evictNotebookFromCaches({ queryClient, projectRef: ref, id: notebookId })
       },
       confirmClose: (notebookTabs) => {
         const dirtyCount = notebookTabs.filter((tab) => {
diff --git a/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx
new file mode 100644
index 00000000000..6762874a9b9
--- /dev/null
+++ b/apps/studio/components/interfaces/Explorer/__tests__/ExplorerNotebookTab.assistant-cache-invalidation.test.tsx
@@ -0,0 +1,176 @@
+import { QueryClient } from '@tanstack/react-query'
+import { screen } from '@testing-library/react'
+import { HttpResponse } from 'msw'
+import { describe, expect, it, vi } from 'vitest'
+
+import { ExplorerNotebookTab } from '../ExplorerNotebookTab'
+import type { components } from '@/data/api'
+import { contentKeys } from '@/data/content/keys'
+import {
+  applyNotebookCacheEffects,
+  collectNotebookCacheEffects,
+} from '@/lib/ai/notebook-cache-invalidation'
+import { createAssistantMessageWithUpdateNotebookTool } from '@/lib/ai/test-fixtures'
+import { notebooksState } from '@/state/notebooks/notebooks-state'
+import type { Notebook } from '@/state/notebooks/types'
+import { createTabsState, TabsStateContext } from '@/state/tabs'
+import { customRender } from '@/tests/lib/custom-render'
+import { addAPIMock } from '@/tests/lib/msw'
+import { setupSqlEditorMocks } from '@/tests/lib/sql-editor-test-utils'
+
+const PROJECT_REF = 'default'
+const NOTEBOOK_ID = 'notebook-assistant-cache-test'
+
+vi.mock('common', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('common')>()
+  return {
+    ...actual,
+    IS_PLATFORM: true,
+    useParams: () => ({ ref: 'default', id: 'notebook-assistant-cache-test' }),
+    useFlag: () => false,
+  }
+})
+
+vi.mock('@/components/ui/CodeEditor/CodeEditor', () => ({
+  CodeEditor: ({ value }: { value: string }) => (
+    <textarea aria-label="SQL editor" value={value} readOnly />
+  ),
+}))
+
+vi.mock('../QueryEditor/QuerySourceMenu', () => ({
+  QuerySourceMenu: () => null,
+}))
+
+const seedNotebook = () => {
+  delete notebooksState.notebooks[NOTEBOOK_ID]
+  const notebook: Notebook = {
+    id: NOTEBOOK_ID,
+    type: 'notebook',
+    name: 'Test notebook',
+    visibility: 'project',
+    favorite: false,
+    owner_id: 1,
+    project_id: 1,
+    content: {
+      schema_version: 1,
+      cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'Original content' }],
+    },
+  }
+  notebooksState.setNotebook({ projectRef: PROJECT_REF, notebook })
+}
+
+const renderNotebookTab = (queryClient: QueryClient) =>
+  customRender(
+    <TabsStateContext.Provider value={createTabsState(PROJECT_REF)}>
+      <ExplorerNotebookTab />
+    </TabsStateContext.Provider>,
+    { queryClient }
+  )
+
+describe('ExplorerNotebookTab — assistant cache invalidation', () => {
+  it('refetches and renders the updated cells after an assistant update_notebook tool call completes', async () => {
+    setupSqlEditorMocks()
+    seedNotebook()
+
+    const queryClient = new QueryClient()
+    queryClient.setQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID), { id: NOTEBOOK_ID })
+
+    addAPIMock({
+      method: 'get',
+      path: '/platform/projects/:ref/content/item/:id',
+      response: () =>
+        HttpResponse.json<components['schemas']['GetUserContentByIdResponse']>({
+          id: NOTEBOOK_ID,
+          type: 'notebook',
+          name: 'Test notebook',
+          description: '',
+          favorite: false,
+          folder_id: null,
+          inserted_at: '2024-01-01T00:00:00.000Z',
+          updated_at: '2024-01-02T00:00:00.000Z',
+          visibility: 'project',
+          owner_id: 1,
+          project_id: 1,
+          content: {
+            schema_version: 1,
+            cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'Updated by assistant' }],
+          },
+        }),
+    })
+
+    renderNotebookTab(queryClient)
+
+    expect(await screen.findByText('Original content')).toBeInTheDocument()
+
+    const message = createAssistantMessageWithUpdateNotebookTool({
+      id: NOTEBOOK_ID,
+      name: 'Test notebook',
+    })
+    const effects = collectNotebookCacheEffects([message], new Set())
+    await applyNotebookCacheEffects({ queryClient, projectRef: PROJECT_REF, effects })
+
+    expect(await screen.findByText('Updated by assistant')).toBeInTheDocument()
+    expect(screen.queryByText('Original content')).not.toBeInTheDocument()
+  })
+
+  it('shows the updated cells on remount, when the query cache still holds the pre-eviction notebook', async () => {
+    setupSqlEditorMocks()
+    seedNotebook()
+
+    const queryClient = new QueryClient()
+    queryClient.setQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID), {
+      id: NOTEBOOK_ID,
+      type: 'notebook',
+      name: 'Test notebook',
+      description: '',
+      favorite: false,
+      visibility: 'project',
+      owner_id: 1,
+      project_id: 1,
+      inserted_at: '2024-01-01T00:00:00.000Z',
+      updated_at: '2024-01-01T00:00:00.000Z',
+      content: {
+        schema_version: 1,
+        cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'Original content' }],
+      },
+    })
+
+    addAPIMock({
+      method: 'get',
+      path: '/platform/projects/:ref/content/item/:id',
+      response: () =>
+        HttpResponse.json<components['schemas']['GetUserContentByIdResponse']>({
+          id: NOTEBOOK_ID,
+          type: 'notebook',
+          name: 'Test notebook',
+          description: '',
+          favorite: false,
+          folder_id: null,
+          inserted_at: '2024-01-01T00:00:00.000Z',
+          updated_at: '2024-01-02T00:00:00.000Z',
+          visibility: 'project',
+          owner_id: 1,
+          project_id: 1,
+          content: {
+            schema_version: 1,
+            cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'Updated by assistant' }],
+          },
+        }),
+    })
+
+    // While the tab is unmounted (navigated away to a separate Explorer chat tab), the
+    // assistant updates the notebook.
+    const message = createAssistantMessageWithUpdateNotebookTool({
+      id: NOTEBOOK_ID,
+      name: 'Test notebook',
+    })
+    const effects = collectNotebookCacheEffects([message], new Set())
+    await applyNotebookCacheEffects({ queryClient, projectRef: PROJECT_REF, effects })
+
+    // Simulate navigating back: the notebook tab mounts fresh.
+    renderNotebookTab(queryClient)
+
+    expect(await screen.findByText('Updated by assistant')).toBeInTheDocument()
+    expect(screen.queryByText('Original content')).not.toBeInTheDocument()
+  })
+})
diff --git a/apps/studio/data/content/notebooks/notebook-cache.test.ts b/apps/studio/data/content/notebooks/notebook-cache.test.ts
index 6e95ffbdd84..875f5c23103 100644
--- a/apps/studio/data/content/notebooks/notebook-cache.test.ts
+++ b/apps/studio/data/content/notebooks/notebook-cache.test.ts
@@ -38,7 +38,7 @@ afterEach(() => {
 })
 
 describe('evictNotebookFromCaches', () => {
-  it('removes a saved notebook from the store and invalidates its cache entry in "refresh" mode', async () => {
+  it('removes a saved notebook from the store and drops its cache entry', async () => {
     seedNotebook('saved')
     const queryClient = new QueryClient()
     seedQueryData(queryClient)
@@ -47,26 +47,6 @@ describe('evictNotebookFromCaches', () => {
       queryClient,
       projectRef: PROJECT_REF,
       id: NOTEBOOK_ID,
-      mode: 'refresh',
-    })
-
-    expect(evicted).toBe(true)
-    expect(notebooksState.notebooks[NOTEBOOK_ID]).toBeUndefined()
-    expect(
-      queryClient.getQueryState(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID))?.isInvalidated
-    ).toBe(true)
-  })
-
-  it('removes a saved notebook from the store and drops its cache entry in "remove" mode', async () => {
-    seedNotebook('saved')
-    const queryClient = new QueryClient()
-    seedQueryData(queryClient)
-
-    const evicted = await evictNotebookFromCaches({
-      queryClient,
-      projectRef: PROJECT_REF,
-      id: NOTEBOOK_ID,
-      mode: 'remove',
     })
 
     expect(evicted).toBe(true)
@@ -83,7 +63,6 @@ describe('evictNotebookFromCaches', () => {
       queryClient,
       projectRef: PROJECT_REF,
       id: NOTEBOOK_ID,
-      mode: 'remove',
     })
 
     expect(evicted).toBe(false)
@@ -101,7 +80,6 @@ describe('evictNotebookFromCaches', () => {
       queryClient,
       projectRef: PROJECT_REF,
       id: NOTEBOOK_ID,
-      mode: 'remove',
     })
 
     expect(evicted).toBe(false)
diff --git a/apps/studio/data/content/notebooks/notebook-cache.ts b/apps/studio/data/content/notebooks/notebook-cache.ts
index 3ad97128b15..be8db93ff53 100644
--- a/apps/studio/data/content/notebooks/notebook-cache.ts
+++ b/apps/studio/data/content/notebooks/notebook-cache.ts
@@ -6,8 +6,6 @@ import { notebooksState } from '@/state/notebooks/notebooks-state'
 import type { StateNotebook } from '@/state/notebooks/types'
 import { hasUnsavedChanges } from '@/state/sql-editor/sql-editor-lifecycle'
 
-export type NotebookCacheEvictionMode = 'refresh' | 'remove'
-
 /**
  * Whether a notebook has edits worth discarding on close: anything not yet
  * saved, except a never-persisted notebook that's still empty (nothing to
@@ -28,6 +26,12 @@ export function hasDiscardableChanges(
  * Applies to a persisted notebook (always safe to refetch) and to a dirty
  * unsaved notebook (safe once the caller has confirmed discarding it).
  *
+ * Removes the query entry outright rather than invalidating it: a mounted
+ * `useNotebookQuery` observer would otherwise read the stale cached value
+ * synchronously, before its refetch lands, and `notebooksState.setNotebook`'s
+ * merge guard would treat that stale merge as already-loaded and drop the
+ * real update.
+ *
  * @returns A boolean indicating whether the notebook was successfully evicted
  *          from the cache.
  */
@@ -35,25 +39,17 @@ export async function evictNotebookFromCaches({
   queryClient,
   projectRef,
   id,
-  mode,
 }: {
   queryClient: QueryClient
   projectRef: string
   id: string
-  mode: NotebookCacheEvictionMode
 }): Promise<boolean> {
   const stateNotebook = notebooksState.notebooks[id]
   const canEvict = stateNotebook?.status === 'saved' || hasDiscardableChanges(stateNotebook)
   if (!canEvict) return false
 
   notebooksState.removeNotebook({ id })
-
-  const queryKey = contentKeys.resource(projectRef, id)
-  if (mode === 'remove') {
-    queryClient.removeQueries({ queryKey })
-  } else {
-    await queryClient.invalidateQueries({ queryKey })
-  }
+  queryClient.removeQueries({ queryKey: contentKeys.resource(projectRef, id) })
 
   return true
 }
diff --git a/apps/studio/lib/ai/notebook-cache-invalidation.test.ts b/apps/studio/lib/ai/notebook-cache-invalidation.test.ts
new file mode 100644
index 00000000000..5ece69aa63f
--- /dev/null
+++ b/apps/studio/lib/ai/notebook-cache-invalidation.test.ts
@@ -0,0 +1,182 @@
+import { QueryClient } from '@tanstack/react-query'
+import type { ToolUIPart, UIMessage } from 'ai'
+import { afterEach, describe, expect, it } from 'vitest'
+
+import {
+  applyNotebookCacheEffects,
+  collectNotebookCacheEffects,
+} from './notebook-cache-invalidation'
+import {
+  createAssistantMessageWithCreateNotebookTool,
+  createAssistantMessageWithUpdateNotebookTool,
+  createAssistantTextMessage,
+  createUserMessage,
+} from './test-fixtures'
+import { contentKeys } from '@/data/content/keys'
+import { notebooksState } from '@/state/notebooks/notebooks-state'
+import type { Notebook } from '@/state/notebooks/types'
+
+const PROJECT_REF = 'default'
+
+function toolPart(overrides: Partial<ToolUIPart>): ToolUIPart {
+  return {
+    type: 'tool-update_notebook',
+    toolCallId: 'call-1',
+    state: 'output-available',
+    input: {},
+    output: { id: 'notebook-1', name: 'Signup funnel' },
+    ...overrides,
+  } as ToolUIPart
+}
+
+function assistantMessage(parts: UIMessage['parts'], id = 'assistant-1'): UIMessage {
+  return { id, role: 'assistant', parts }
+}
+
+describe('collectNotebookCacheEffects', () => {
+  it('collects a create_notebook output-available part', () => {
+    const messages = [createAssistantMessageWithCreateNotebookTool()]
+
+    const effects = collectNotebookCacheEffects(messages, new Set())
+
+    expect(effects).toEqual([{ _tag: 'upserted', toolCallId: 'call-notebook-1', id: 'notebook-1' }])
+  })
+
+  it('collects an update_notebook output-available part', () => {
+    const messages = [createAssistantMessageWithUpdateNotebookTool()]
+
+    const effects = collectNotebookCacheEffects(messages, new Set())
+
+    expect(effects).toEqual([{ _tag: 'upserted', toolCallId: 'call-notebook-1', id: 'notebook-1' }])
+  })
+
+  it('ignores non-terminal tool states', () => {
+    const messages = [
+      assistantMessage([
+        toolPart({ type: 'tool-create_notebook', state: 'input-streaming' }),
+        toolPart({ type: 'tool-update_notebook', state: 'approval-requested' }),
+        toolPart({ type: 'tool-update_notebook', state: 'output-error' }),
+      ]),
+    ]
+
+    expect(collectNotebookCacheEffects(messages, new Set())).toEqual([])
+  })
+
+  it('ignores malformed output', () => {
+    const messages = [
+      assistantMessage([toolPart({ output: { unexpected: true } })]),
+      assistantMessage([toolPart({ toolCallId: 'call-2', output: undefined })]),
+    ]
+
+    expect(collectNotebookCacheEffects(messages, new Set())).toEqual([])
+  })
+
+  it('ignores unrelated message types and tool parts', () => {
+    const messages = [
+      createUserMessage('update my notebook'),
+      createAssistantTextMessage('Sure, updating it now.'),
+    ]
+
+    expect(collectNotebookCacheEffects(messages, new Set())).toEqual([])
+  })
+
+  it('collects output on an earlier message even when a later message has already finished', () => {
+    const messages = [
+      assistantMessage([toolPart({ toolCallId: 'call-earlier' })], 'assistant-1'),
+      createAssistantTextMessage('All done!', 'assistant-2'),
+    ]
+
+    const effects = collectNotebookCacheEffects(messages, new Set())
+
+    expect(effects).toEqual([{ _tag: 'upserted', toolCallId: 'call-earlier', id: 'notebook-1' }])
+  })
+
+  it('dedupes against the processed set', () => {
+    const messages = [assistantMessage([toolPart({ toolCallId: 'call-1' })])]
+
+    expect(collectNotebookCacheEffects(messages, new Set(['call-1']))).toEqual([])
+  })
+})
+
+describe('applyNotebookCacheEffects', () => {
+  const NOTEBOOK: Notebook = {
+    id: 'notebook-1',
+    type: 'notebook',
+    name: 'Test notebook',
+    visibility: 'project',
+    favorite: false,
+    owner_id: 1,
+    project_id: 1,
+    content: { schema_version: 1, cells: [] },
+  }
+
+  afterEach(() => {
+    delete notebooksState.notebooks[NOTEBOOK.id]
+    notebooksState.needsSaving.clear()
+  })
+
+  it('invalidates the nav list and evicts an upserted, saved notebook from the cache', async () => {
+    notebooksState.setNotebook({ projectRef: PROJECT_REF, notebook: NOTEBOOK })
+    const queryClient = new QueryClient()
+    queryClient.setQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK.id), { id: NOTEBOOK.id })
+    queryClient.setQueryData(contentKeys.allContentLists(PROJECT_REF), [])
+    queryClient.setQueryData(contentKeys.infiniteList(PROJECT_REF), {})
+
+    await applyNotebookCacheEffects({
+      queryClient,
+      projectRef: PROJECT_REF,
+      effects: [{ _tag: 'upserted', toolCallId: 'call-1', id: NOTEBOOK.id }],
+    })
+
+    expect(queryClient.getQueryState(contentKeys.allContentLists(PROJECT_REF))?.isInvalidated).toBe(
+      true
+    )
+    expect(queryClient.getQueryState(contentKeys.infiniteList(PROJECT_REF))?.isInvalidated).toBe(
+      true
+    )
+    expect(queryClient.getQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK.id))).toBeUndefined()
+    expect(notebooksState.notebooks[NOTEBOOK.id]).toBeUndefined()
+  })
+
+  it('leaves a dirty (unsaved) notebook in the store untouched', async () => {
+    notebooksState.addNotebook({ projectRef: PROJECT_REF, notebook: NOTEBOOK })
+    const queryClient = new QueryClient()
+
+    await applyNotebookCacheEffects({
+      queryClient,
+      projectRef: PROJECT_REF,
+      effects: [{ _tag: 'upserted', toolCallId: 'call-1', id: NOTEBOOK.id }],
+    })
+
+    expect(notebooksState.notebooks[NOTEBOOK.id]).toBeDefined()
+  })
+
+  it('leaves an edited (unsaved, non-empty) notebook untouched', async () => {
+    notebooksState.setNotebook({ projectRef: PROJECT_REF, notebook: NOTEBOOK })
+    notebooksState.updateCells({
+      id: NOTEBOOK.id,
+      cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'local edit' }],
+    })
+    expect(notebooksState.notebooks[NOTEBOOK.id].status).toBe('unsaved')
+    const queryClient = new QueryClient()
+
+    await applyNotebookCacheEffects({
+      queryClient,
+      projectRef: PROJECT_REF,
+      effects: [{ _tag: 'upserted', toolCallId: 'call-1', id: NOTEBOOK.id }],
+    })
+
+    expect(notebooksState.notebooks[NOTEBOOK.id]).toBeDefined()
+  })
+
+  it('no-ops entirely for an empty effects list', async () => {
+    const queryClient = new QueryClient()
+    queryClient.setQueryData(contentKeys.allContentLists(PROJECT_REF), [])
+
+    await applyNotebookCacheEffects({ queryClient, projectRef: PROJECT_REF, effects: [] })
+
+    expect(queryClient.getQueryState(contentKeys.allContentLists(PROJECT_REF))?.isInvalidated).toBe(
+      false
+    )
+  })
+})
diff --git a/apps/studio/lib/ai/notebook-cache-invalidation.ts b/apps/studio/lib/ai/notebook-cache-invalidation.ts
new file mode 100644
index 00000000000..8e990b13dee
--- /dev/null
+++ b/apps/studio/lib/ai/notebook-cache-invalidation.ts
@@ -0,0 +1,67 @@
+import type { QueryClient } from '@tanstack/react-query'
+import type { ToolUIPart, UIMessage } from 'ai'
+
+import { notebookToolOutputSchema } from '@/components/ui/AIAssistantPanel/Message.utils'
+import { contentKeys } from '@/data/content/keys'
+import { evictNotebookFromCaches } from '@/data/content/notebooks/notebook-cache'
+import { notebooksState } from '@/state/notebooks/notebooks-state'
+
+export type NotebookCacheEffect =
+  | { _tag: 'upserted'; toolCallId: string; id: string }
+  | { _tag: 'deleted'; toolCallId: string; id: string }
+
+const NOTEBOOK_UPSERT_TOOL_TYPES = new Set(['tool-create_notebook', 'tool-update_notebook'])
+
+function isNotebookUpsertPart(part: UIMessage['parts'][number]): part is ToolUIPart {
+  return NOTEBOOK_UPSERT_TOOL_TYPES.has(part.type)
+}
+
+export function collectNotebookCacheEffects(
+  messages: Array<UIMessage>,
+  processed: ReadonlySet<string>
+): Array<NotebookCacheEffect> {
+  const effects: Array<NotebookCacheEffect> = []
+
+  for (const message of messages) {
+    if (message.role !== 'assistant') continue
+
+    for (const part of message.parts ?? []) {
+      if (!isNotebookUpsertPart(part)) continue
+      if (part.state !== 'output-available') continue
+      if (processed.has(part.toolCallId)) continue
+
+      const result = notebookToolOutputSchema.safeParse(part.output)
+      if (!result.success) continue
+
+      effects.push({ _tag: 'upserted', toolCallId: part.toolCallId, id: result.data.id })
+    }
+  }
+
+  return effects
+}
+
+export async function applyNotebookCacheEffects({
+  queryClient,
+  projectRef,
+  effects,
+}: {
+  queryClient: QueryClient
+  projectRef: string
+  effects: Array<NotebookCacheEffect>
+}): Promise<void> {
+  if (effects.length === 0) return
+
+  await Promise.all([
+    queryClient.invalidateQueries({ queryKey: contentKeys.allContentLists(projectRef) }),
+    queryClient.invalidateQueries({ queryKey: contentKeys.infiniteList(projectRef) }),
+  ])
+
+  await Promise.all(
+    effects.map((effect) => {
+      const stateNotebook = notebooksState.notebooks[effect.id]
+      if (stateNotebook && stateNotebook.status !== 'saved') return
+
+      return evictNotebookFromCaches({ queryClient, projectRef, id: effect.id })
+    })
+  )
+}
diff --git a/apps/studio/lib/ai/test-fixtures.ts b/apps/studio/lib/ai/test-fixtures.ts
index 83bcfbc3043..79500ae4d98 100644
--- a/apps/studio/lib/ai/test-fixtures.ts
+++ b/apps/studio/lib/ai/test-fixtures.ts
@@ -73,6 +73,28 @@ export function createAssistantMessageWithUpdateNotebookTool(
   }
 }
 
+export function createAssistantMessageWithCreateNotebookTool(
+  output: Record<string, unknown> = {
+    id: 'notebook-1',
+    name: 'Signup funnel',
+  },
+  id = 'assistant-notebook-msg-1'
+): UIMessage {
+  return {
+    id,
+    role: 'assistant',
+    parts: [
+      {
+        type: 'tool-create_notebook',
+        state: 'output-available',
+        toolCallId: 'call-notebook-1',
+        input: { name: 'Signup funnel', content: { schema_version: 1, cells: [] } },
+        output,
+      } satisfies ToolUIPart,
+    ],
+  }
+}
+
 export function createAssistantMessageWithMultipleTools(
   id = 'assistant-multi-tool-msg-1'
 ): UIMessage {
diff --git a/apps/studio/state/ai-assistant-state.notebook-cache-invalidation.test.ts b/apps/studio/state/ai-assistant-state.notebook-cache-invalidation.test.ts
new file mode 100644
index 00000000000..8126f6bf928
--- /dev/null
+++ b/apps/studio/state/ai-assistant-state.notebook-cache-invalidation.test.ts
@@ -0,0 +1,181 @@
+import type { UIMessageChunk } from 'ai'
+import { convertArrayToReadableStream } from 'ai/test'
+import { afterEach, describe, expect, it, vi } from 'vitest'
+
+import { contentKeys } from '@/data/content/keys'
+import { getQueryClient } from '@/data/query-client'
+import { notebooksState } from '@/state/notebooks/notebooks-state'
+import type { Notebook } from '@/state/notebooks/types'
+
+const testContext = vi.hoisted(() => ({
+  queuedStreams: [] as Array<Array<UIMessageChunk>>,
+  onSend: undefined as (() => void) | undefined,
+}))
+
+vi.mock('ai', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('ai')>()
+  return {
+    ...actual,
+    DefaultChatTransport: class {
+      constructor(private options: any = {}) {}
+      async sendMessages(opts: any) {
+        await this.options.prepareSendMessagesRequest?.({
+          api: this.options.api,
+          id: opts.chatId,
+          messages: opts.messages,
+          body: { ...this.options.body, ...opts.body },
+          headers: {},
+          credentials: undefined,
+          requestMetadata: opts.metadata,
+          trigger: opts.trigger,
+          messageId: opts.messageId,
+        })
+        testContext.onSend?.()
+        const chunks = testContext.queuedStreams.shift()
+        if (!chunks) throw new Error('No queued stream for this sendMessages call')
+        return convertArrayToReadableStream(chunks)
+      }
+      async reconnectToStream() {
+        return null
+      }
+    },
+  }
+})
+
+const { createAiAssistantState } = await import('./ai-assistant-state')
+
+const PROJECT_REF = 'default'
+const NOTEBOOK_ID = 'notebook-onfinish-test'
+
+const seedNotebook = () => {
+  delete notebooksState.notebooks[NOTEBOOK_ID]
+  const notebook: Notebook = {
+    id: NOTEBOOK_ID,
+    type: 'notebook',
+    name: 'Test notebook',
+    visibility: 'project',
+    favorite: false,
+    owner_id: 1,
+    project_id: 1,
+    content: {
+      schema_version: 1,
+      cells: [{ _tag: 'markdown_cell', _id: 'cell-1', text: 'Original' }],
+    },
+  }
+  notebooksState.setNotebook({ projectRef: PROJECT_REF, notebook })
+}
+
+afterEach(() => {
+  delete notebooksState.notebooks[NOTEBOOK_ID]
+  notebooksState.needsSaving.clear()
+  testContext.queuedStreams = []
+  testContext.onSend = undefined
+})
+
+describe('createChatInstance onFinish — notebook cache invalidation via a real Chat/stream', () => {
+  it('evicts the notebook from caches once update_notebook reaches output-available, after an approval round trip', async () => {
+    seedNotebook()
+
+    const queryClient = getQueryClient()
+    queryClient.setQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID), { id: NOTEBOOK_ID })
+
+    const state = createAiAssistantState()
+    state.setContext({ projectRef: PROJECT_REF })
+    const chatId = state.createChat({ name: 'Delete a cell' })
+    const chatInstance = state.chatInstances[chatId]
+
+    // First stream: assistant proposes update_notebook, which needs approval.
+    testContext.queuedStreams.push([
+      { type: 'start' },
+      {
+        type: 'tool-input-available',
+        toolCallId: 'call-1',
+        toolName: 'update_notebook',
+        input: {
+          id: NOTEBOOK_ID,
+          expected_updated_at: '2024-01-01T00:00:00.000Z',
+          operations: [{ _tag: 'delete_cell', cell_id: 'cell-1' }],
+        },
+      },
+      { type: 'tool-approval-request', approvalId: 'approval-1', toolCallId: 'call-1' },
+      { type: 'finish' },
+    ])
+
+    await chatInstance.sendMessage({ text: 'Delete the first cell' })
+
+    // Still pending approval — must not evict yet.
+    expect(notebooksState.notebooks[NOTEBOOK_ID]).toBeDefined()
+
+    // Second stream: after approval, the tool actually executes and returns its result.
+    testContext.queuedStreams.push([
+      { type: 'start' },
+      {
+        type: 'tool-output-available',
+        toolCallId: 'call-1',
+        output: { id: NOTEBOOK_ID, name: 'Test notebook' },
+      },
+      { type: 'finish' },
+    ])
+
+    await chatInstance.addToolApprovalResponse({ id: 'approval-1', approved: true })
+
+    await vi.waitFor(() => {
+      expect(notebooksState.notebooks[NOTEBOOK_ID]).toBeUndefined()
+    })
+    expect(queryClient.getQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID))).toBeUndefined()
+  })
+
+  it('binds cache effects to the project the request was sent under, not whatever project is active when the stream finishes', async () => {
+    seedNotebook()
+    const OTHER_PROJECT_REF = 'other-project'
+
+    const queryClient = getQueryClient()
+    queryClient.setQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID), { id: NOTEBOOK_ID })
+
+    const state = createAiAssistantState()
+    state.setContext({ projectRef: PROJECT_REF })
+    const chatId = state.createChat({ name: 'Delete a cell' })
+    const chatInstance = state.chatInstances[chatId]
+
+    testContext.queuedStreams.push([
+      { type: 'start' },
+      {
+        type: 'tool-input-available',
+        toolCallId: 'call-1',
+        toolName: 'update_notebook',
+        input: {
+          id: NOTEBOOK_ID,
+          expected_updated_at: '2024-01-01T00:00:00.000Z',
+          operations: [{ _tag: 'delete_cell', cell_id: 'cell-1' }],
+        },
+      },
+      { type: 'tool-approval-request', approvalId: 'approval-1', toolCallId: 'call-1' },
+      { type: 'finish' },
+    ])
+    await chatInstance.sendMessage({ text: 'Delete the first cell' })
+
+    testContext.queuedStreams.push([
+      { type: 'start' },
+      {
+        type: 'tool-output-available',
+        toolCallId: 'call-1',
+        output: { id: NOTEBOOK_ID, name: 'Test notebook' },
+      },
+      { type: 'finish' },
+    ])
+
+    // Simulate navigating to a different project while the approval round-trip is in
+    // flight: after the request was sent (with the origin project ref in its body), but
+    // before its stream resolves and onFinish runs.
+    testContext.onSend = () => state.setContext({ projectRef: OTHER_PROJECT_REF })
+
+    await chatInstance.addToolApprovalResponse({ id: 'approval-1', approved: true })
+
+    await vi.waitFor(() => {
+      expect(notebooksState.notebooks[NOTEBOOK_ID]).toBeUndefined()
+    })
+    // Must evict the origin project's cache entry — the one the write actually happened
+    // in — not whichever project happened to be active once the stream finished.
+    expect(queryClient.getQueryData(contentKeys.resource(PROJECT_REF, NOTEBOOK_ID))).toBeUndefined()
+  })
+})
diff --git a/apps/studio/state/ai-assistant-state.tsx b/apps/studio/state/ai-assistant-state.tsx
index 8d0ee276b2f..b86c82e3a11 100644
--- a/apps/studio/state/ai-assistant-state.tsx
+++ b/apps/studio/state/ai-assistant-state.tsx
@@ -17,10 +17,15 @@ import { proxy, ref, snapshot, subscribe, useSnapshot } from 'valtio'
 import type { SqlSnippetSource } from '@/components/interfaces/SQLEditor/querySource'
 import type { AiSupportStatus } from '@/data/feedback/ai-chat-front-sync'
 import { constructHeaders } from '@/data/fetchers'
+import { getQueryClient } from '@/data/query-client'
 import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
 import { prepareMessagesForAPI } from '@/lib/ai/message-utils'
 import { isKnownAssistantModelId } from '@/lib/ai/model.utils'
 import type { AssistantModelId } from '@/lib/ai/model.utils'
+import {
+  applyNotebookCacheEffects,
+  collectNotebookCacheEffects,
+} from '@/lib/ai/notebook-cache-invalidation'
 import { BASE_PATH, IS_PLATFORM } from '@/lib/constants'
 
 type SuggestionsType = {
@@ -277,6 +282,19 @@ function createChatInstance(
   state: AiAssistantState,
   options: { id: string; initialMessages: MessageType[] }
 ) {
+  // Seeded so effects already reflected in persisted history aren't replayed on the first
+  // onFinish after a reload.
+  const processedNotebookToolCallIds = new Set<string>(
+    collectNotebookCacheEffects(options.initialMessages, new Set()).map(
+      (effect) => effect.toolCallId
+    )
+  )
+
+  // The project a pending request's tool calls actually ran against — captured when the
+  // request is sent, not re-read from (mutable) state.context in onFinish, since the user
+  // can switch projects while the request is still in flight.
+  let requestProjectRef: string | undefined
+
   return new Chat<MessageType>({
     id: options.id,
     messages: options.initialMessages.map((message) => sanitizeForCloning(message)),
@@ -299,6 +317,8 @@ function createChatInstance(
         // Get the chat specific to this request to ensure we have the correct name
         const chat = state.chats[options.id]
 
+        requestProjectRef = state.context.projectRef
+
         return {
           ...opts,
           body: {
@@ -368,6 +388,15 @@ function createChatInstance(
             .then(({ syncSupportChatToFront }) => syncSupportChatToFront(options.id, state))
             .catch(() => {})
         }
+
+        const projectRef = requestProjectRef
+        if (projectRef) {
+          const effects = collectNotebookCacheEffects(messages, processedNotebookToolCallIds)
+          effects.forEach((effect) => processedNotebookToolCallIds.add(effect.toolCallId))
+          if (effects.length > 0) {
+            void applyNotebookCacheEffects({ queryClient: getQueryClient(), projectRef, effects })
+          }
+        }
       }
     },
   })