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 index 6762874a9b9..497d111046f 100644 --- 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 @@ -10,12 +10,15 @@ import { applyNotebookCacheEffects, collectNotebookCacheEffects, } from '@/lib/ai/notebook-cache-invalidation' -import { createAssistantMessageWithUpdateNotebookTool } from '@/lib/ai/test-fixtures' +import { + createAssistantMessageWithDeleteNotebookTool, + 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 { addAPIMock, type APIErrorBody } from '@/tests/lib/msw' import { setupSqlEditorMocks } from '@/tests/lib/sql-editor-test-utils' const PROJECT_REF = 'default' @@ -173,4 +176,33 @@ describe('ExplorerNotebookTab — assistant cache invalidation', () => { expect(await screen.findByText('Updated by assistant')).toBeInTheDocument() expect(screen.queryByText('Original content')).not.toBeInTheDocument() }) + + it('shows the not-found state after an assistant delete_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({ message: 'Notebook not found' }, { status: 404 }), + }) + + renderNotebookTab(queryClient) + + expect(await screen.findByText('Original content')).toBeInTheDocument() + + const message = createAssistantMessageWithDeleteNotebookTool({ + id: NOTEBOOK_ID, + name: 'Test notebook', + }) + const effects = collectNotebookCacheEffects([message], new Set()) + await applyNotebookCacheEffects({ queryClient, projectRef: PROJECT_REF, effects }) + + expect(await screen.findByText('Notebook not found')).toBeInTheDocument() + expect(screen.queryByText('Original content')).not.toBeInTheDocument() + }) }) diff --git a/apps/studio/lib/ai/notebook-cache-invalidation.test.ts b/apps/studio/lib/ai/notebook-cache-invalidation.test.ts index 5ece69aa63f..1b51c14642e 100644 --- a/apps/studio/lib/ai/notebook-cache-invalidation.test.ts +++ b/apps/studio/lib/ai/notebook-cache-invalidation.test.ts @@ -8,6 +8,7 @@ import { } from './notebook-cache-invalidation' import { createAssistantMessageWithCreateNotebookTool, + createAssistantMessageWithDeleteNotebookTool, createAssistantMessageWithUpdateNotebookTool, createAssistantTextMessage, createUserMessage, @@ -50,6 +51,14 @@ describe('collectNotebookCacheEffects', () => { expect(effects).toEqual([{ _tag: 'upserted', toolCallId: 'call-notebook-1', id: 'notebook-1' }]) }) + it('collects a delete_notebook output-available part', () => { + const messages = [createAssistantMessageWithDeleteNotebookTool()] + + const effects = collectNotebookCacheEffects(messages, new Set()) + + expect(effects).toEqual([{ _tag: 'deleted', toolCallId: 'call-notebook-1', id: 'notebook-1' }]) + }) + it('ignores non-terminal tool states', () => { const messages = [ assistantMessage([ @@ -138,6 +147,29 @@ describe('applyNotebookCacheEffects', () => { expect(notebooksState.notebooks[NOTEBOOK.id]).toBeUndefined() }) + it('invalidates the nav list and evicts a deleted 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: 'deleted', 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() diff --git a/apps/studio/lib/ai/notebook-cache-invalidation.ts b/apps/studio/lib/ai/notebook-cache-invalidation.ts index 8e990b13dee..cc40e3ba7db 100644 --- a/apps/studio/lib/ai/notebook-cache-invalidation.ts +++ b/apps/studio/lib/ai/notebook-cache-invalidation.ts @@ -10,10 +10,14 @@ 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']) +const NOTEBOOK_MUTATION_TOOL_TYPES = new Set([ + 'tool-create_notebook', + 'tool-update_notebook', + 'tool-delete_notebook', +]) -function isNotebookUpsertPart(part: UIMessage['parts'][number]): part is ToolUIPart { - return NOTEBOOK_UPSERT_TOOL_TYPES.has(part.type) +function isNotebookMutationPart(part: UIMessage['parts'][number]): part is ToolUIPart { + return NOTEBOOK_MUTATION_TOOL_TYPES.has(part.type) } export function collectNotebookCacheEffects( @@ -26,14 +30,18 @@ export function collectNotebookCacheEffects( if (message.role !== 'assistant') continue for (const part of message.parts ?? []) { - if (!isNotebookUpsertPart(part)) continue + if (!isNotebookMutationPart(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 }) + effects.push({ + _tag: part.type === 'tool-delete_notebook' ? 'deleted' : 'upserted', + toolCallId: part.toolCallId, + id: result.data.id, + }) } } diff --git a/apps/studio/lib/ai/test-fixtures.ts b/apps/studio/lib/ai/test-fixtures.ts index 79500ae4d98..46db994a2a7 100644 --- a/apps/studio/lib/ai/test-fixtures.ts +++ b/apps/studio/lib/ai/test-fixtures.ts @@ -95,6 +95,28 @@ export function createAssistantMessageWithCreateNotebookTool( } } +export function createAssistantMessageWithDeleteNotebookTool( + output: Record = { + id: 'notebook-1', + name: 'Signup funnel', + }, + id = 'assistant-notebook-msg-1' +): UIMessage { + return { + id, + role: 'assistant', + parts: [ + { + type: 'tool-delete_notebook', + state: 'output-available', + toolCallId: 'call-notebook-1', + input: { id: 'notebook-1' }, + output, + } satisfies ToolUIPart, + ], + } +} + export function createAssistantMessageWithMultipleTools( id = 'assistant-multi-tool-msg-1' ): UIMessage {