From 804f4b5acb47d6eac18ee36c46cd3ffd4b73b167 Mon Sep 17 00:00:00 2001 From: keval shah Date: Wed, 4 Mar 2026 23:56:58 -0500 Subject: [PATCH] fix(ui): ensure session deletion properly clears backend and cache (#12043) * fix(ui): ensure session deletion properly clears backend and cache * fix: resolved PR comments and add new regression test * fix: resolved PR comments and add new regression test * [autofix.ci] apply automated fixes --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../hooks/use-edit-session-info.ts | 4 +- .../chat-header/hooks/use-get-add-sessions.ts | 1 + .../chat-messages/hooks/use-chat-history.ts | 53 ++++-- .../chat-view/utils/session-filter.ts | 31 ++++ .../queries/messages/use-delete-sessions.ts | 87 +++++++-- .../IOModal/components/session-view.tsx | 4 +- src/frontend/src/types/messages/index.ts | 14 +- src/frontend/src/types/messages/session.ts | 17 ++ .../session-deletion-data-leakage.spec.ts | 168 ++++++++++++++++++ 9 files changed, 340 insertions(+), 39 deletions(-) create mode 100644 src/frontend/src/components/core/playgroundComponent/chat-view/utils/session-filter.ts create mode 100644 src/frontend/src/types/messages/session.ts create mode 100644 src/frontend/tests/core/regression/session-deletion-data-leakage.spec.ts diff --git a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-edit-session-info.ts b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-edit-session-info.ts index e7d6dd9b9e..31e075ca24 100644 --- a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-edit-session-info.ts +++ b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-edit-session-info.ts @@ -37,8 +37,8 @@ export const useEditSessionInfo = ({ const { mutate: deleteSession } = useDeleteSession(); const handleDelete = (sessionId: string) => { - if (sessionId && dbSessions.includes(sessionId)) { - deleteSession({ sessionId: sessionId }); + if (sessionId && flowId) { + deleteSession({ sessionId: sessionId, flowId: flowId }); } if (flowId && sessionId === selectedSession) { setSelectedSession(flowId); diff --git a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-get-add-sessions.ts b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-get-add-sessions.ts index 75e44a9000..c13b3b264a 100644 --- a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-get-add-sessions.ts +++ b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-header/hooks/use-get-add-sessions.ts @@ -138,6 +138,7 @@ export const useGetAddSessions: UseGetAddSessionsReturnType = ({ }; const removeLocalSession = (sessionId: string) => { + // Update state - the useEffect on line 67-77 will sync to sessionStorage setLocalSessions((prev) => { const updated = new Set(prev); updated.delete(sessionId); diff --git a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/hooks/use-chat-history.ts b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/hooks/use-chat-history.ts index 8695523ae0..cc98c05558 100644 --- a/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/hooks/use-chat-history.ts +++ b/src/frontend/src/components/core/playgroundComponent/chat-view/chat-messages/hooks/use-chat-history.ts @@ -4,6 +4,7 @@ import { useGetFlowId } from "@/components/core/playgroundComponent/hooks/use-ge import { useGetMessagesQuery } from "@/controllers/API/queries/messages"; import type { ChatMessageType } from "@/types/chat"; import type { Message } from "@/types/messages"; +import { isMessageForSession } from "../../utils/session-filter"; import sortSenderMessages from "../utils/sort-sender-messages"; export const useChatHistory = (visibleSession: string | null) => { @@ -45,17 +46,20 @@ export const useChatHistory = (visibleSession: string | null) => { if (queryData && typeof queryData === "object" && "rows" in queryData) { const rowsData = queryData.rows as { data?: Message[] } | undefined; if (rowsData && typeof rowsData === "object" && "data" in rowsData) { - const backendMessages = rowsData.data || []; + const backendMessages = (rowsData.data || []).filter((msg: Message) => + isMessageForSession(msg, currentFlowId, visibleSession), + ); + const existingCache = queryClient.getQueryData(sessionCacheKey) || []; - // Only initialize if cache is empty and we have backend messages + // Only initialize if cache is empty and we have backend messages for this session if (existingCache.length === 0 && backendMessages.length > 0) { queryClient.setQueryData(sessionCacheKey, backendMessages); } } } - }, [queryData, queryClient, sessionCacheKey]); + }, [queryData, queryClient, sessionCacheKey, currentFlowId, visibleSession]); // Use session cache as the single source of truth // updateMessage and addUserMessage handle all updates (placeholders, streaming, etc.) @@ -65,21 +69,10 @@ export const useChatHistory = (visibleSession: string | null) => { const chatHistory = useMemo(() => { // Filter messages for current session const filteredMessages: ChatMessageType[] = messages - .filter((message: Message) => { - const isCurrentFlow = message.flow_id === currentFlowId; - // If visibleSession is the flow_id, it means we are in the default session - // In the default session, we show messages that have the same session_id as the flow_id - // OR messages that have NO session_id (legacy behavior) - if (visibleSession === currentFlowId) { - const matches = - isCurrentFlow && - (message.session_id === visibleSession || !message.session_id); - return matches; - } - const matches = isCurrentFlow && message.session_id === visibleSession; - return matches; - }) - .map((message: Message) => { + .filter((message: Message) => + isMessageForSession(message, currentFlowId, visibleSession), + ) + .map((message: Message): ChatMessageType => { let files = message.files; // Handle the "[]" case, empty string, or already parsed array if (Array.isArray(files)) { @@ -96,6 +89,28 @@ export const useChatHistory = (visibleSession: string | null) => { } const messageText = message.text || ""; + // Convert Message.properties to ChatMessageType.properties (PropertiesType) + // Properties are now properly typed in Message, no cast needed + let properties: ChatMessageType["properties"] = undefined; + if (message.properties?.source?.id) { + properties = { + source: { + id: message.properties.source.id, + display_name: message.properties.source.display_name || "", + source: message.properties.source.source || "", + }, + state: message.properties.state, + icon: message.properties.icon, + background_color: message.properties.background_color, + text_color: message.properties.text_color, + targets: message.properties.targets, + edited: message.properties.edited, + allow_markdown: message.properties.allow_markdown, + positive_feedback: message.properties.positive_feedback, + build_duration: message.properties.build_duration, + }; + } + return { isSend: message.sender === "User", message: messageText, @@ -110,7 +125,7 @@ export const useChatHistory = (visibleSession: string | null) => { text_color: message.text_color, content_blocks: message.content_blocks, category: message.category, - properties: message.properties, + properties: properties, }; }); diff --git a/src/frontend/src/components/core/playgroundComponent/chat-view/utils/session-filter.ts b/src/frontend/src/components/core/playgroundComponent/chat-view/utils/session-filter.ts new file mode 100644 index 0000000000..7b773eefae --- /dev/null +++ b/src/frontend/src/components/core/playgroundComponent/chat-view/utils/session-filter.ts @@ -0,0 +1,31 @@ +import type { Message } from "@/types/messages"; + +/** + * Determines if a message belongs to a specific session. + * + * Why this logic exists: + * - Default session (sessionId === flowId): Shows messages with matching session_id OR no session_id (legacy) + * - Named sessions: Only shows messages with exact session_id match + * - This prevents cross-session data leakage after deletions + * + * @param msg - The message to check + * @param flowId - The current flow ID + * @param sessionId - The session ID to filter for (null means no filtering) + * @returns true if the message belongs to the session + */ +export function isMessageForSession( + msg: Message, + flowId: string, + sessionId: string | null, +): boolean { + if (!sessionId) return false; + + const isCurrentFlow = msg.flow_id === flowId; + + if (sessionId === flowId) { + // Default session: include messages with matching session_id or no session_id (legacy behavior) + return isCurrentFlow && (msg.session_id === sessionId || !msg.session_id); + } + + return isCurrentFlow && msg.session_id === sessionId; +} diff --git a/src/frontend/src/controllers/API/queries/messages/use-delete-sessions.ts b/src/frontend/src/controllers/API/queries/messages/use-delete-sessions.ts index bf18e53d8f..ff5d0b510e 100644 --- a/src/frontend/src/controllers/API/queries/messages/use-delete-sessions.ts +++ b/src/frontend/src/controllers/API/queries/messages/use-delete-sessions.ts @@ -1,22 +1,32 @@ import type { UseMutationResult } from "@tanstack/react-query"; -import type { useMutationFunctionType } from "@/types/api"; +import type { + DeleteSessionError, + DeleteSessionParams, + DeleteSessionResponse, +} from "@/types/messages/session"; import { api } from "../../api"; import { getURL } from "../../helpers/constants"; import { UseRequestProcessor } from "../../services/request-processor"; -interface DeleteSessionParams { - sessionId: string; -} - -export const useDeleteSession: useMutationFunctionType< - undefined, - DeleteSessionParams -> = (options?) => { +export const useDeleteSession = (options?: { + onSuccess?: ( + data: DeleteSessionResponse, + variables: DeleteSessionParams, + context: unknown, + ) => void; + onSettled?: ( + data: DeleteSessionResponse | undefined, + error: DeleteSessionError | null, + variables: DeleteSessionParams, + context: unknown, + ) => void; + onError?: (error: DeleteSessionError) => void; +}) => { const { mutate, queryClient } = UseRequestProcessor(); const deleteSession = async ({ sessionId, - }: DeleteSessionParams): Promise => { + }: DeleteSessionParams): Promise => { const response = await api.delete( `${getURL("MESSAGES")}/session/${sessionId}`, ); @@ -24,16 +34,65 @@ export const useDeleteSession: useMutationFunctionType< }; const mutation: UseMutationResult< - DeleteSessionParams, - any, + DeleteSessionResponse, + DeleteSessionError, DeleteSessionParams > = mutate(["useDeleteSession"], deleteSession, { ...options, - onSettled: (data, error, variables, context) => { + onSuccess: (data, variables, context, ...rest) => { + // Cast needed because UseRequestProcessor's mutate doesn't properly infer callback types + const vars = variables as unknown as DeleteSessionParams; + + // Remove all message queries for this session immediately to prevent stale data + if (vars.flowId) { + // Remove session-specific queries + queryClient.removeQueries({ + queryKey: [ + "useGetMessagesQuery", + { id: vars.flowId, session_id: vars.sessionId }, + ], + }); + + // Also remove any queries that might have the session_id in params (e.g., Message Logs) + queryClient.removeQueries({ + predicate: (query) => { + const queryKey = query.queryKey; + if ( + Array.isArray(queryKey) && + queryKey[0] === "useGetMessagesQuery" + ) { + const params = queryKey[1] as Record; + if (params?.params && typeof params.params === "object") { + const nestedParams = params.params as Record; + if (nestedParams.session_id === vars.sessionId) { + return true; + } + } + } + return false; + }, + }); + } + options?.onSuccess?.(data, vars, context); + }, + onSettled: (data, error, variables, context, ...rest) => { + // Cast needed because UseRequestProcessor's mutate doesn't properly infer callback types + const vars = variables as unknown as DeleteSessionParams; + + // Invalidate sessions list to refresh the sidebar queryClient.invalidateQueries({ queryKey: ["useGetSessionsFromFlowQuery"], }); - options?.onSettled?.(data, error, variables, context); + + // Invalidate all message queries to ensure fresh data everywhere + if (vars.flowId) { + queryClient.invalidateQueries({ + queryKey: ["useGetMessagesQuery"], + refetchType: "none", // Prevent automatic refetching to avoid race conditions + }); + } + + options?.onSettled?.(data, error, vars, context); }, }); diff --git a/src/frontend/src/modals/IOModal/components/session-view.tsx b/src/frontend/src/modals/IOModal/components/session-view.tsx index ef9e5f853f..1910b05ce9 100644 --- a/src/frontend/src/modals/IOModal/components/session-view.tsx +++ b/src/frontend/src/modals/IOModal/components/session-view.tsx @@ -56,9 +56,7 @@ export default function SessionView({ const rowsData = queryData.rows as { data?: any[] } | undefined; if (rowsData && typeof rowsData === "object" && "data" in rowsData) { const fetchedMessages = rowsData.data || []; - if (fetchedMessages.length > 0) { - setMessages(fetchedMessages); - } + setMessages(fetchedMessages); } } }, [queryData, setMessages]); diff --git a/src/frontend/src/types/messages/index.ts b/src/frontend/src/types/messages/index.ts index 11b58a46b2..03931a92b2 100644 --- a/src/frontend/src/types/messages/index.ts +++ b/src/frontend/src/types/messages/index.ts @@ -15,7 +15,19 @@ type Message = { category?: string; properties?: { state?: "partial" | "complete"; - source?: { id?: string }; + source?: { + id?: string; + display_name?: string; + source?: string; + }; + icon?: string; + background_color?: string; + text_color?: string; + targets?: string[]; + edited?: boolean; + allow_markdown?: boolean; + positive_feedback?: boolean | null; + build_duration?: number | null; [key: string]: unknown; }; content_blocks?: ContentBlock[]; diff --git a/src/frontend/src/types/messages/session.ts b/src/frontend/src/types/messages/session.ts new file mode 100644 index 0000000000..24696e4b92 --- /dev/null +++ b/src/frontend/src/types/messages/session.ts @@ -0,0 +1,17 @@ +export interface DeleteSessionParams { + sessionId: string; + flowId?: string; +} + +export interface DeleteSessionResponse { + message: string; +} + +export interface DeleteSessionError { + response?: { + data?: { + detail?: string; + }; + }; + message?: string; +} diff --git a/src/frontend/tests/core/regression/session-deletion-data-leakage.spec.ts b/src/frontend/tests/core/regression/session-deletion-data-leakage.spec.ts new file mode 100644 index 0000000000..a1171a2fa7 --- /dev/null +++ b/src/frontend/tests/core/regression/session-deletion-data-leakage.spec.ts @@ -0,0 +1,168 @@ +import { expect, test } from "../../fixtures"; +import { awaitBootstrapTest } from "../../utils/await-bootstrap-test"; +import { initialGPTsetup } from "../../utils/initialGPTsetup"; +import type { Page } from "@playwright/test"; + +test.describe("Session Deletion Data Leakage Fix", () => { + // Helper to send a message in the playground + async function sendMessage(page: Page, message: string) { + await page.waitForSelector('[data-testid="input-chat-playground"]', { + timeout: 10000, + }); + await page.getByTestId("input-chat-playground").last().fill(message); + await page.getByTestId("button-send").last().click(); + await page.waitForTimeout(2000); // Wait for message to be processed + } + + // Helper to create a new session + async function createNewSession(page: Page) { + await page.getByTestId("new-chat").click(); + await page.waitForTimeout(1000); // Wait for session to be created + } + + // Helper to delete a session via the more menu + async function deleteSession(page: Page, sessionName: string) { + // Find all session selectors + const sessionSelectors = await page.getByTestId("session-selector").all(); + + // Find the one with exact matching text + for (const selector of sessionSelectors) { + const text = await selector.textContent(); + // Use exact match to avoid matching "Default Session" when looking for "New Session 0" + if (text?.trim() === sessionName) { + // Hover to make the more button visible + await selector.hover(); + await page.waitForTimeout(500); // Wait for hover effects + + // Click the more options button + const moreButton = selector.locator('[aria-label="More options"]'); + await moreButton.click({ timeout: 5000 }); + + // Wait for the menu to open + await page.waitForTimeout(500); + + // Wait for delete option to be visible and click it + await page + .getByTestId("delete-session-option") + .waitFor({ state: "visible", timeout: 5000 }); + await page.getByTestId("delete-session-option").click(); + await page.waitForTimeout(1000); // Wait for deletion to complete + break; + } + } + } + + // Helper to get message count in the current view + async function getMessageCount(page: Page): Promise { + const messages = await page + .locator('[data-testid="div-chat-message"]') + .all(); + return messages.length; + } + + // Helper to check if a message exists + async function messageExists(page: Page, text: string): Promise { + const message = page.getByText(text, { exact: false }); + return await message.isVisible().catch(() => false); + } + + test( + "should prevent data leakage when default session is deleted and recreated", + { tag: ["@release", "@regression"] }, + async ({ page }) => { + test.skip( + !process?.env?.OPENAI_API_KEY, + "OPENAI_API_KEY required to run this test", + ); + + await awaitBootstrapTest(page); + + // Load a starter project + await page.getByTestId("side_nav_options_all-templates").click(); + await page.getByRole("heading", { name: "Basic Prompting" }).click(); + await initialGPTsetup(page); + + // Open playground + await page + .getByRole("button", { name: "Playground", exact: true }) + .click(); + await page.waitForTimeout(2000); + + // Send message in default session + const originalMessage = `Original message ${Date.now()}`; + await sendMessage(page, originalMessage); + await page.waitForTimeout(2000); + + // Verify message appears + expect(await messageExists(page, originalMessage)).toBeTruthy(); + + // Delete the default session + await deleteSession(page, "Default Session"); + await page.waitForTimeout(1000); + + // Verify the old message does NOT appear after deletion + expect(await messageExists(page, originalMessage)).toBeFalsy(); + + // Send a different message (this will be in a new/recreated default session) + const newMessage = `New message ${Date.now()}`; + await sendMessage(page, newMessage); + await page.waitForTimeout(2000); + + // Verify only the new message appears + expect(await messageExists(page, newMessage)).toBeTruthy(); + expect(await messageExists(page, originalMessage)).toBeFalsy(); + + // Verify message count is correct (should only have the new message) + const messageCount = await getMessageCount(page); + expect(messageCount).toBe(1); + }, + ); + + test( + "should clear LLM context when session is deleted", + { tag: ["@release", "@regression"] }, + async ({ page }) => { + test.skip( + !process?.env?.OPENAI_API_KEY, + "OPENAI_API_KEY required to run this test", + ); + + await awaitBootstrapTest(page); + + // Load a starter project with memory + await page.getByTestId("side_nav_options_all-templates").click(); + await page.getByRole("heading", { name: "Basic Prompting" }).click(); + await initialGPTsetup(page); + + // Open playground + await page + .getByRole("button", { name: "Playground", exact: true }) + .click(); + await page.waitForTimeout(2000); + + // Send a message with specific information in default session + await sendMessage(page, "My name is Victor"); + await page.waitForTimeout(3000); // Wait for AI response + + // Delete the default session to clear context + await deleteSession(page, "Default Session"); + await page.waitForTimeout(1000); + + // The playground should now show an empty state or create a new default session + // Ask a question that would require the deleted context + await sendMessage(page, "What is my name?"); + await page.waitForTimeout(3000); // Wait for AI response + + // Get the response text + const messages = await page + .locator('[data-testid="div-chat-message"]') + .all(); + const lastMessage = messages[messages.length - 1]; + const responseText = await lastMessage.textContent(); + + // Verify the AI does NOT remember "Victor" from the deleted session + // The response should indicate it doesn't know the name + expect(responseText?.toLowerCase()).not.toContain("victor"); + }, + ); +});