diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx index ada70dee893..59ed3744931 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx @@ -1,4 +1,7 @@ -import type { Meta, StoryObj } from "@storybook/react-vite"; +import type { Decorator, Meta, StoryObj } from "@storybook/react-vite"; +import { delay } from "msw"; +import { useState } from "react"; +import { QueryClient, QueryClientProvider } from "react-query"; import { expect, fn, @@ -19,11 +22,12 @@ import { MockWorkspace, } from "#/testHelpers/entities"; import { withDashboardProvider } from "#/testHelpers/storybook"; +import { persistedAttachmentsStorageKey } from "../hooks/useFileAttachments"; import { getReasoningEffortForModel, saveReasoningEffortForModel, } from "../utils/reasoningEffort"; -import { AgentCreateForm } from "./AgentCreateForm"; +import { AgentCreateForm, emptyInputStorageKey } from "./AgentCreateForm"; const permittedOrgsKey = permittedOrganizationsKey({ object: { resource_type: "chat", owner_id: "me" }, @@ -131,19 +135,27 @@ type Story = StoryObj; const defaultArgs = meta.args; -const mockPermittedOrganizations = (permissions: Record) => { +const mockPermittedOrganizations = ( + permissions: Record, + delayMs = 0, +) => { spyOn(API, "getOrganizations").mockResolvedValue([ MockDefaultOrganization, MockOrganization2, ]); - spyOn(API, "checkAuthorization").mockResolvedValue(permissions); + spyOn(API, "checkAuthorization").mockImplementation(async () => { + if (delayMs > 0) { + await delay(delayMs); + } + return permissions; + }); }; export const Default: Story = {}; const submitMessage = async (canvasElement: HTMLElement, message: string) => { const canvas = within(canvasElement); - const input = canvas.getByTestId("chat-message-input"); + const input = canvas.getByRole("textbox", { name: "Chat message" }); await userEvent.click(input); await userEvent.keyboard(message); await userEvent.click(canvas.getByRole("button", { name: "Send" })); @@ -895,24 +907,496 @@ export const WithOrganizationPicker: Story = { queries: [ { key: permittedOrgsKey, - data: [MockDefaultOrganization, MockOrganization2], + data: [MockOrganization2, MockDefaultOrganization], }, ], }, play: async ({ canvasElement }) => { const canvas = within(canvasElement); - // Verify the org picker rendered (component didn't crash). - await waitFor(() => { - expect(canvas.getByTestId("compact-org-selector")).toBeInTheDocument(); + const organizationPicker = canvas.getByRole("button", { + name: "Organization: My Organization", }); - // Type into the chat input to trigger re-renders. If the - // permittedOrgs fallback is referentially unstable, this - // causes a render cascade that hits React's update limit. - const input = canvas.getByTestId("chat-message-input"); + await expect(organizationPicker).toBeVisible(); + + const input = canvas.getByRole("textbox", { name: "Chat message" }); await userEvent.click(input); await userEvent.keyboard("hello world"); - // The org picker should still be present after typing. - expect(canvas.getByTestId("compact-org-selector")).toBeInTheDocument(); + await expect( + canvas.getByRole("button", { + name: "Organization: My Organization", + }), + ).toBeVisible(); + }, +}; + +export const RestrictedMultiOrganizationUser: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + beforeEach: () => { + spyOn(API, "getOrganizations").mockResolvedValue([ + MockDefaultOrganization, + MockOrganization2, + ]); + // Model agents-access: "me" supplies the owner for member-scoped chat:create. + spyOn(API, "checkAuthorization").mockImplementation(async ({ checks }) => + Object.fromEntries( + Object.entries(checks).map(([id, check]) => [ + id, + check.object.owner_id === "me" && + check.object.organization_id === MockOrganization2.id, + ]), + ), + ); + }, + args: { + ...defaultArgs, + onCreateChat: fn().mockResolvedValue(undefined), + }, + play: async ({ canvasElement, args }) => { + await submitMessage(canvasElement, "test message"); + await waitFor(() => { + expect(args.onCreateChat).toHaveBeenCalledWith( + expect.objectContaining({ + organizationId: MockOrganization2.id, + }), + ); + }); + }, +}; + +export const RestrictedUserKeepsPersistedWorkspace: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + args: { + ...defaultArgs, + onCreateChat: fn().mockResolvedValue(undefined), + workspaceOptions: [ + { + ...MockWorkspace, + id: "ws-permitted-org", + name: "permitted-workspace", + organization_id: MockOrganization2.id, + }, + ], + workspaceCount: 1, + }, + beforeEach: () => { + localStorage.setItem("agents.selected-workspace-id", "ws-permitted-org"); + mockPermittedOrganizations({ + [MockDefaultOrganization.id]: false, + [MockOrganization2.id]: true, + }); + }, + play: async ({ canvasElement, args }) => { + await submitMessage(canvasElement, "test message"); + await waitFor(() => { + expect(args.onCreateChat).toHaveBeenCalledWith( + expect.objectContaining({ + organizationId: MockOrganization2.id, + workspaceId: "ws-permitted-org", + }), + ); + }); + }, +}; + +export const RestrictedUserKeepsPersistedAttachments: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + beforeEach: () => { + localStorage.clear(); + localStorage.setItem( + "agents.persisted-attachments", + JSON.stringify([ + { + fileId: "file-permitted-org", + fileName: "notes.txt", + fileType: "text/plain", + lastModified: 1700000000000, + organizationId: MockOrganization2.id, + }, + ]), + ); + mockPermittedOrganizations({ + [MockDefaultOrganization.id]: false, + [MockOrganization2.id]: true, + }); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => { + expect(canvas.getByLabelText("Remove notes.txt")).toBeInTheDocument(); + }); + const stored = localStorage.getItem("agents.persisted-attachments"); + expect(stored).toContain("file-permitted-org"); + }, +}; + +export const OrganizationAuthorizationFailure: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + beforeEach: () => { + localStorage.clear(); + localStorage.setItem(emptyInputStorageKey, "draft message"); + spyOn(API, "getOrganizations").mockResolvedValue([ + MockDefaultOrganization, + MockOrganization2, + ]); + spyOn(API, "checkAuthorization").mockRejectedValue( + new Error("authorization check failed"), + ); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await canvas.findAllByText(/authorization check failed/i); + expect(canvas.getByRole("button", { name: "Send" })).toBeDisabled(); + }, +}; + +export const LoadingWorkspacesBlocksSendUntilValidated: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + args: { + ...defaultArgs, + workspaceOptions: [], + isWorkspacesLoading: true, + }, + beforeEach: () => { + localStorage.setItem(emptyInputStorageKey, "draft message"); + localStorage.setItem("agents.selected-workspace-id", "ws-default-org"); + mockPermittedOrganizations({ + [MockDefaultOrganization.id]: true, + [MockOrganization2.id]: true, + }); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + // Wait for permissions to settle before checking workspace validation. + await canvas.findByRole("button", { + name: "Organization: My Organization", + }); + await expect(canvas.getByRole("button", { name: "Send" })).toBeDisabled(); + }, +}; + +export const DelayedOrganizationAuthorization: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + beforeEach: () => { + localStorage.setItem(emptyInputStorageKey, "draft message"); + mockPermittedOrganizations( + { + [MockDefaultOrganization.id]: true, + [MockOrganization2.id]: true, + }, + 1_500, + ); + }, + args: { + ...defaultArgs, + workspaceOptions: [ + { + ...MockWorkspace, + id: "ws-provisional", + name: "provisional-workspace", + organization_id: MockDefaultOrganization.id, + }, + ], + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const sendButton = canvas.getByRole("button", { name: "Send" }); + await expect(sendButton).toBeDisabled(); + await expect( + canvas.getByRole("button", { name: "More options" }), + ).toBeDisabled(); + // dispatchEvent returns false when a handler accepted the drop + // via preventDefault, giving a race-free accepted/ignored signal. + const dropFile = (name: string): boolean => { + const dataTransfer = new DataTransfer(); + dataTransfer.items.add(new File(["hello"], name, { type: "text/plain" })); + return canvas.getByTestId("chat-composer").dispatchEvent( + new DragEvent("drop", { + bubbles: true, + cancelable: true, + dataTransfer, + }), + ); + }; + // Pending authorization leaves attachments without a valid org, so drops + // must be ignored. + expect(dropFile("drop.txt")).toBe(true); + expect(canvas.queryByLabelText("Remove drop.txt")).not.toBeInTheDocument(); + // The pending option list is unfiltered, so the picker must stay hidden. + expect( + canvas.queryByRole("button", { name: /organization/i }), + ).not.toBeInTheDocument(); + await waitFor(() => expect(sendButton).toBeEnabled(), { timeout: 3_000 }); + await canvas.findByRole("button", { name: /organization/i }); + // Positive control: once settled the same drop is accepted and + // attaches, so the pending-state assertions exercised a real path. + expect(dropFile("after.txt")).toBe(false); + await waitFor(() => + expect(canvas.getByLabelText("Remove after.txt")).toBeInTheDocument(), + ); + }, +}; + +// Mutable permissions let play functions change authorization across refetches. +// The story-local QueryClient exposes those refetches; the preview client's +// instance is inaccessible and uses infinite stale time. +const revocablePermissions: Record = {}; +let revocableQueryClient: QueryClient | undefined; + +const withRevocableQueryClient: Decorator = (Story) => { + const [queryClient] = useState( + () => + new QueryClient({ + defaultOptions: { + queries: { + staleTime: Number.POSITIVE_INFINITY, + retry: false, + }, + }, + }), + ); + revocableQueryClient = queryClient; + return ( + + + + ); +}; + +const mockRevocablePermissions = (permissions: Record) => { + for (const key of Object.keys(revocablePermissions)) { + delete revocablePermissions[key]; + } + Object.assign(revocablePermissions, permissions); + spyOn(API, "getOrganizations").mockResolvedValue([ + MockDefaultOrganization, + MockOrganization2, + ]); + spyOn(API, "checkAuthorization").mockImplementation(async () => ({ + ...revocablePermissions, + })); +}; + +const revocableStoryContext = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, + decorators: [withRevocableQueryClient], +}; + +const allOrganizationsPermitted = { + [MockDefaultOrganization.id]: true, + [MockOrganization2.id]: true, +}; + +export const RevokedSelectionDoesNotResurrect: Story = { + ...revocableStoryContext, + beforeEach: () => { + mockRevocablePermissions(allOrganizationsPermitted); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const trigger = await canvas.findByRole("button", { + name: "Organization: My Organization", + }); + await userEvent.click(trigger); + await userEvent.click( + await screen.findByRole("option", { name: /My Organization 2/ }), + ); + await canvas.findByRole("button", { + name: "Organization: My Organization 2", + }); + + revocablePermissions[MockOrganization2.id] = false; + await revocableQueryClient?.invalidateQueries(); + await waitFor(() => + expect( + canvas.queryByRole("button", { name: /organization/i }), + ).not.toBeInTheDocument(), + ); + + revocablePermissions[MockOrganization2.id] = true; + await revocableQueryClient?.invalidateQueries(); + await canvas.findByRole("button", { + name: "Organization: My Organization", + }); + }, +}; + +export const RevokedOrgChangeClearsStoredWorkspace: Story = { + ...revocableStoryContext, + args: { + ...defaultArgs, + workspaceOptions: [ + { + ...MockWorkspace, + id: "ws-default-org", + name: "default-workspace", + organization_id: MockDefaultOrganization.id, + }, + ], + workspaceCount: 1, + }, + beforeEach: () => { + localStorage.setItem("agents.selected-workspace-id", "ws-default-org"); + mockRevocablePermissions(allOrganizationsPermitted); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => + expect( + canvas.getByLabelText("Remove workspace default-workspace"), + ).toBeInTheDocument(), + ); + + revocablePermissions[MockDefaultOrganization.id] = false; + await revocableQueryClient?.invalidateQueries(); + await waitFor(() => + expect( + canvas.queryByLabelText("Remove workspace default-workspace"), + ).not.toBeInTheDocument(), + ); + + revocablePermissions[MockDefaultOrganization.id] = true; + await revocableQueryClient?.invalidateQueries(); + await canvas.findByRole("button", { + name: "Organization: My Organization 2", + }); + expect( + canvas.queryByLabelText("Remove workspace default-workspace"), + ).not.toBeInTheDocument(); + expect(localStorage.getItem("agents.selected-workspace-id")).toBeNull(); + }, +}; + +export const EmptyPermittedSetPreservesStoredWorkspace: Story = { + ...revocableStoryContext, + args: { + ...defaultArgs, + workspaceOptions: [ + { + ...MockWorkspace, + id: "ws-org-2", + name: "org2-workspace", + organization_id: MockOrganization2.id, + }, + ], + }, + beforeEach: () => { + localStorage.setItem("agents.selected-workspace-id", "ws-org-2"); + mockRevocablePermissions({ + [MockDefaultOrganization.id]: false, + [MockOrganization2.id]: true, + }); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await canvas.findByLabelText("Remove workspace org2-workspace"); + + revocablePermissions[MockOrganization2.id] = false; + await revocableQueryClient?.invalidateQueries(); + await canvas.findByText(/don't have permission/i); + + revocablePermissions[MockOrganization2.id] = true; + await revocableQueryClient?.invalidateQueries(); + await canvas.findByLabelText("Remove workspace org2-workspace"); + expect(localStorage.getItem("agents.selected-workspace-id")).toBe( + "ws-org-2", + ); + }, +}; + +export const SingleOrgIgnoresStalePermittedCache: Story = { + parameters: { + showOrganizations: false, + organizations: [MockDefaultOrganization], + queries: [ + { + key: permittedOrgsKey, + data: [MockOrganization2], + }, + ], + }, + args: { + ...defaultArgs, + onCreateChat: fn().mockResolvedValue(undefined), + }, + play: async ({ canvasElement, args }) => { + await submitMessage(canvasElement, "test message"); + await waitFor(() => { + expect(args.onCreateChat).toHaveBeenCalledWith( + expect.objectContaining({ + organizationId: MockDefaultOrganization.id, + }), + ); + }); + }, +}; + +export const RevokedPendingOrgClosesConfirmDialog: Story = { + ...revocableStoryContext, + beforeEach: () => { + localStorage.clear(); + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([ + { + fileId: "file-default-org", + fileName: "notes.txt", + fileType: "text/plain", + lastModified: 1700000000000, + organizationId: MockDefaultOrganization.id, + }, + ]), + ); + mockRevocablePermissions(allOrganizationsPermitted); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const body = within(canvasElement.ownerDocument.body); + await waitFor(() => + expect(canvas.getByLabelText("Remove notes.txt")).toBeInTheDocument(), + ); + await userEvent.click( + await canvas.findByRole("button", { + name: "Organization: My Organization", + }), + ); + await userEvent.click( + await screen.findByRole("option", { name: /My Organization 2/ }), + ); + await body.findByText( + "Changing organization will remove your current attachments.", + ); + + revocablePermissions[MockOrganization2.id] = false; + await revocableQueryClient?.invalidateQueries(); + await waitFor(() => + expect( + body.queryByText( + "Changing organization will remove your current attachments.", + ), + ).not.toBeInTheDocument(), + ); + expect(canvas.getByLabelText("Remove notes.txt")).toBeInTheDocument(); }, }; @@ -1008,14 +1492,27 @@ export const PermittedOrgsResolvesToEmpty: Story = { parameters: { showOrganizations: true, organizations: [MockDefaultOrganization, MockOrganization2], - // Deliberately do not pre-seed permittedOrgsKey. Let the - // mocked API calls drive the async permission resolution. }, args: { ...defaultArgs, onCreateChat: fn().mockResolvedValue(undefined), }, beforeEach: () => { + localStorage.clear(); + // Another org's persisted attachment must survive a visit while + // the user has no chat permission anywhere. + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([ + { + fileId: "file-other-org", + fileName: "keep.txt", + fileType: "text/plain", + lastModified: 1000, + organizationId: MockOrganization2.id, + }, + ]), + ); mockPermittedOrganizations({ [MockDefaultOrganization.id]: false, [MockOrganization2.id]: false, @@ -1023,36 +1520,17 @@ export const PermittedOrgsResolvesToEmpty: Story = { }, play: async ({ canvasElement, args }) => { const canvas = within(canvasElement); - - // Wait for the permitted orgs query to resolve. The org picker - // should disappear since no org is permitted. await waitFor( () => { - expect( - canvas.queryByTestId("compact-org-selector"), - ).not.toBeInTheDocument(); + expect(canvas.getByText(/don't have permission/i)).toBeInTheDocument(); }, { timeout: 3000 }, ); - - // Type a message and submit the form. - const input = canvas.getByTestId("chat-message-input"); - await userEvent.click(input); - await userEvent.keyboard("test message"); - await userEvent.click(canvas.getByRole("button", { name: "Send" })); - - // Verify onCreateChat was called with a non-empty organizationId. - await waitFor(() => { - expect(args.onCreateChat).toHaveBeenCalled(); - }); - const options = (args.onCreateChat as ReturnType).mock - .calls[0]?.[0] as { organizationId: string } | undefined; - if (!options) { - throw new Error("Expected onCreateChat to receive options"); - } - expect(options.organizationId).not.toBe(""); - // It should fall back to the default org from the dashboard. - expect(options.organizationId).toBe(MockDefaultOrganization.id); + expect(canvas.getByRole("button", { name: "Send" })).toBeDisabled(); + expect(args.onCreateChat).not.toHaveBeenCalled(); + expect( + localStorage.getItem(persistedAttachmentsStorageKey) ?? "", + ).toContain("file-other-org"); }, }; diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx index 51032010dda..5b3f3822145 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx @@ -1,4 +1,4 @@ -import { type FC, useEffect, useEffectEvent, useRef, useState } from "react"; +import { type FC, useEffect, useRef, useState } from "react"; import { useQuery } from "react-query"; import { toast } from "sonner"; import { isApiError } from "#/api/errors"; @@ -271,42 +271,96 @@ export const AgentCreateForm: FC = ({ : undefined; const initialOrg = organizations.find((o) => o.is_default) ?? organizations[0]; + // effectiveWorkspaceId nulls a stored selection outside the effective org's + // filtered workspace list without deleting it. Preserve the stored value + // because the permitted-organizations query may resolve after mount and + // change the effective org. const [selectedWorkspaceId, setSelectedWorkspaceId] = useState( - () => { - const stored = localStorage.getItem(selectedWorkspaceIdStorageKey); - if (!stored) return null; - - // The stored value is kept optimistically until workspaces - // load. effectiveWorkspaceId (computed after render) drops - // it if it doesn't match the current org's workspaces. - if (workspaceOptions.length === 0) return stored; - - // Validate the stored workspace still exists and belongs - // to the initial org. Without this, a workspace from a - // previously selected org persists across sessions and - // gets submitted even though it's hidden from the picker. - const workspace = workspaceOptions.find((ws) => ws.id === stored); - if (!workspace) { - localStorage.removeItem(selectedWorkspaceIdStorageKey); - return null; - } - if ( - showOrganizations && - initialOrg && - workspace.organization_id !== initialOrg.id - ) { - localStorage.removeItem(selectedWorkspaceIdStorageKey); - return null; - } - return stored; - }, + () => localStorage.getItem(selectedWorkspaceIdStorageKey), ); const [selectedOrg, setSelectedOrg] = useState( - initialOrg ?? null, + null, ); const [pendingOrgChange, setPendingOrgChange] = useState(null); - const organizationId = selectedOrg?.id ?? ""; + const permittedOrgsQuery = useQuery({ + ...permittedOrganizations({ + // agents-access grants chat:create only at member scope. "me" is + // replaced with the caller ID so that permission can match. + object: { resource_type: "chat", owner_id: "me" }, + action: "create", + }), + enabled: showOrganizations, + }); + // Disabled queries retain cached data. When the dashboard hides organization + // selection, its organization list is authoritative so a removed org cannot + // remain selected for submission. + const permittedOrgs = showOrganizations + ? (permittedOrgsQuery.data ?? organizations) + : organizations; + // Treat the dashboard org as provisional until permissions resolve so + // sends and persisted attachments cannot use an unpermitted org. + const orgSelectionSettled = + !showOrganizations || permittedOrgsQuery.data !== undefined; + // Prevent effectiveOrg's dashboard fallback from bypassing an empty + // permitted set. + const noPermittedOrgs = + showOrganizations && permittedOrgsQuery.data?.length === 0; + const selectedOrgIsPermitted = + selectedOrg !== null && + permittedOrgs.some((org) => org.id === selectedOrg.id); + // Clear invalid selections during render so re-permission cannot silently + // restore them and switch attachment state. effectiveOrg already ignores + // the invalid selection in this render. + if (selectedOrg && orgSelectionSettled && !selectedOrgIsPermitted) { + setSelectedOrg(null); + } + // Same rule for a pending change awaiting confirmation: closing the + // dialog prevents confirming into an org that was just revoked. + if ( + pendingOrgChange && + orgSelectionSettled && + !permittedOrgs.some((org) => org.id === pendingOrgChange.id) + ) { + setPendingOrgChange(null); + } + const effectiveOrg = + selectedOrg && selectedOrgIsPermitted + ? selectedOrg + : (permittedOrgs.find((org) => org.is_default) ?? + permittedOrgs[0] ?? + initialOrg ?? + null); + const organizationId = effectiveOrg?.id ?? ""; + // Adopt a permitted fallback so later refetches cannot switch the form to a + // re-permitted default. The permission guard also avoids a render loop. + if ( + orgSelectionSettled && + !selectedOrg && + effectiveOrg && + permittedOrgs.some((org) => org.id === effectiveOrg.id) + ) { + setSelectedOrg(effectiveOrg); + } + // Clear a workspace after a settled org change, before its localStorage value + // is cleared post-commit. An empty permission set has no selectable org, so + // preserve the workspace until its org is re-permitted. + const [lastSettledOrgId, setLastSettledOrgId] = useState(null); + if ( + orgSelectionSettled && + !noPermittedOrgs && + organizationId !== lastSettledOrgId + ) { + setLastSettledOrgId(organizationId); + if (lastSettledOrgId !== null) { + setSelectedWorkspaceId(null); + } + } + useEffect(() => { + if (selectedWorkspaceId === null) { + localStorage.removeItem(selectedWorkspaceIdStorageKey); + } + }, [selectedWorkspaceId]); const [planModeEnabled, setPlanModeEnabled] = useState(false); const hasModelOptions = modelOptions.length > 0; const hasConfiguredModels = hasConfiguredModelsInCatalog(modelCatalog); @@ -377,7 +431,7 @@ export const AgentCreateForm: FC = ({ saveReasoningEffortForModel(selectedModel, value); }; - const isForbidden = !canCreateChat; + const isForbidden = !canCreateChat || noPermittedOrgs; // Filter workspaces by the selected organization. We use // client-side filtering of the full "owner:me" fetch rather @@ -388,8 +442,8 @@ export const AgentCreateForm: FC = ({ // enough to warrant pagination, this should switch to a // server-side organization: query filter. const filteredWorkspaces = - showOrganizations && selectedOrg - ? workspaceOptions.filter((ws) => ws.organization_id === selectedOrg.id) + showOrganizations && effectiveOrg + ? workspaceOptions.filter((ws) => ws.organization_id === effectiveOrg.id) : workspaceOptions; const effectiveWorkspaceId = @@ -398,6 +452,11 @@ export const AgentCreateForm: FC = ({ filteredWorkspaces.some((ws) => ws.id === selectedWorkspaceId)) ? selectedWorkspaceId : null; + // A stored workspace cannot be validated against the effective org + // until the list loads; sending then would silently drop the + // association, so Send stays disabled instead. + const workspaceValidationPending = + selectedWorkspaceId !== null && isWorkspacesLoading; const handleSend = async (message: string, fileIDs?: string[]) => { submitDraft(); @@ -420,6 +479,7 @@ export const AgentCreateForm: FC = ({ }; const { + organizationAdopted, attachments, textContents, uploadStates, @@ -427,10 +487,17 @@ export const AgentCreateForm: FC = ({ handleAttach, handleRemoveAttachment, resetAttachments, - } = useFileAttachments(organizationId || undefined, { - persist: true, - provider: getProviderForModelOption(modelOptions, selectedModel), - }); + } = useFileAttachments( + // Avoid restoring against effectiveOrg's fallback when no org is permitted; + // that would prune attachments persisted for other orgs. + orgSelectionSettled && !noPermittedOrgs + ? organizationId || undefined + : undefined, + { + persist: true, + provider: getProviderForModelOption(modelOptions, selectedModel), + }, + ); const handleSendWithAttachments = async (message: string) => { const fileIds: string[] = []; @@ -459,50 +526,6 @@ export const AgentCreateForm: FC = ({ } }; - const permittedOrgsQuery = useQuery({ - ...permittedOrganizations({ - object: { resource_type: "chat", owner_id: "me" }, - action: "create", - }), - enabled: showOrganizations, - }); - const permittedOrgs = permittedOrgsQuery.data ?? organizations; - - // Reconcile selectedOrg when permission filtering removes it. - // Only pure state setters run during render; side effects - // (localStorage, blob URL cleanup) run in the effect below. - const [prevPermittedOrgs, setPrevPermittedOrgs] = useState(permittedOrgs); - const [orgWasAdjusted, setOrgWasAdjusted] = useState(false); - if (permittedOrgs !== prevPermittedOrgs) { - setPrevPermittedOrgs(permittedOrgs); - if (selectedOrg && !permittedOrgs.some((o) => o.id === selectedOrg.id)) { - // Fall back through: first permitted org, then the - // dashboard default. Never null out selectedOrg. - // organizationId must always be a valid UUID for the - // create-chat request. - const nextOrg = permittedOrgs[0] ?? initialOrg ?? null; - setSelectedOrg(nextOrg); - if (nextOrg?.id !== selectedOrg.id) { - setOrgWasAdjusted(true); - } - } - } - - // Clean up workspace and attachment state after a programmatic - // org change from permission filtering. These calls have side - // effects (localStorage, blob URL revocation) that must not - // run during render. - const onOrgAdjusted = useEffectEvent(() => { - handleWorkspaceChange(null); - resetAttachments(); - }); - useEffect(() => { - if (orgWasAdjusted) { - setOrgWasAdjusted(false); - onOrgAdjusted(); - } - }, [orgWasAdjusted]); - return ( <>
@@ -542,23 +565,27 @@ export const AgentCreateForm: FC = ({ {permittedOrgsQuery.error != null && ( )} - {showOrganizations && permittedOrgs.length > 1 && ( - { - const orgChanged = newOrg.id !== selectedOrg?.id; - if (orgChanged && attachments.length > 0) { - setPendingOrgChange(newOrg); - return; - } - if (orgChanged) { - handleWorkspaceChange(null); - } - setSelectedOrg(newOrg); - }} - /> - )} + {/* The pre-settlement list is the unfiltered dashboard fallback; + selecting from it could destroy existing workspace state. */} + {showOrganizations && + orgSelectionSettled && + permittedOrgs.length > 1 && ( + { + const orgChanged = newOrg.id !== effectiveOrg?.id; + if (orgChanged && attachments.length > 0) { + setPendingOrgChange(newOrg); + return; + } + if (orgChanged) { + handleWorkspaceChange(null); + } + setSelectedOrg(newOrg); + }} + /> + )} = ({ isDisabled={ isCreating || isForbidden || + !orgSelectionSettled || + // Sending before adoption would omit persisted files not yet restored. + !organizationAdopted || + workspaceValidationPending || isPersonalModelOverridesLoading || !hasModelOptions || Boolean(aiGatewayDisabled) @@ -585,7 +616,9 @@ export const AgentCreateForm: FC = ({ planModeEnabled={planModeEnabled} onPlanModeToggle={setPlanModeEnabled} attachments={attachments} - onAttach={handleAttach} + // Files attached before org adoption cannot upload and would be discarded + // when restoration completes. + onAttach={organizationAdopted ? handleAttach : undefined} onRemoveAttachment={handleRemoveAttachment} uploadStates={uploadStates} previewUrls={previewUrls} @@ -599,7 +632,12 @@ export const AgentCreateForm: FC = ({ onMCPAuthComplete={onMCPAuthComplete} workspaceOptions={filteredWorkspaces} selectedWorkspaceId={effectiveWorkspaceId} - onWorkspaceChange={handleWorkspaceChange} + // Do not persist a workspace until its organization is authorized. + onWorkspaceChange={ + orgSelectionSettled && !noPermittedOrgs + ? handleWorkspaceChange + : undefined + } isWorkspaceLoading={isWorkspacesLoading} canConfigureAgentSetup={canConfigureAgentSetup} providerCount={providerCount} @@ -622,10 +660,18 @@ export const AgentCreateForm: FC = ({ hideCancel={false} confirmText="Continue" onConfirm={() => { + if (!pendingOrgChange) { + return; + } + setPendingOrgChange(null); + // Recheck authorization because a refetch may revoke the pending org + // after this render created the closure. + if (!permittedOrgs.some((org) => org.id === pendingOrgChange.id)) { + return; + } resetAttachments(); handleWorkspaceChange(null); setSelectedOrg(pendingOrgChange); - setPendingOrgChange(null); }} onClose={() => setPendingOrgChange(null)} /> diff --git a/site/src/pages/AgentsPage/hooks/useFileAttachments.test.tsx b/site/src/pages/AgentsPage/hooks/useFileAttachments.test.tsx new file mode 100644 index 00000000000..38d82886b33 --- /dev/null +++ b/site/src/pages/AgentsPage/hooks/useFileAttachments.test.tsx @@ -0,0 +1,216 @@ +import { act, render, renderHook, waitFor } from "@testing-library/react"; +import { type FC, Suspense } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { API } from "#/api/api"; +import { createDeferred } from "#/testHelpers/deferred"; +import { + persistedAttachmentsStorageKey, + useFileAttachments, +} from "./useFileAttachments"; + +const persistEntry = (fileId: string, fileName: string, orgId: string) => ({ + fileId, + fileName, + fileType: "text/plain", + lastModified: 1000, + organizationId: orgId, +}); + +const uploadedFileIds = ( + result: ReturnType, +): string[] => + [...result.uploadStates.values()].flatMap((state) => + state.status === "uploaded" && state.fileId ? [state.fileId] : [], + ); + +const renderAttachments = (initialProps: { orgId: string | undefined }) => + renderHook( + ({ orgId }: { orgId: string | undefined }) => + useFileAttachments(orgId, { persist: true }), + { initialProps }, + ); + +const mockDeferredUpload = (): ((value: { id: string }) => void) => { + const deferred = createDeferred<{ id: string }>(); + vi.spyOn(API.experimental, "uploadChatFile").mockReturnValue( + deferred.promise, + ); + return deferred.resolve; +}; + +describe("useFileAttachments org scoping", () => { + beforeEach(() => { + localStorage.clear(); + }); + + it("defers restoration until the org is known", () => { + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([persistEntry("file-b", "b.txt", "org-b")]), + ); + const { result, rerender } = renderAttachments({ orgId: undefined }); + + expect(result.current.attachments).toHaveLength(0); + expect(localStorage.getItem(persistedAttachmentsStorageKey)).toContain( + "file-b", + ); + + rerender({ orgId: "org-b" }); + expect(uploadedFileIds(result.current)).toStrictEqual(["file-b"]); + }); + + it("drops another org's attachments when the org changes", async () => { + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([persistEntry("file-a", "a.txt", "org-a")]), + ); + const { result, rerender } = renderAttachments({ orgId: "org-a" }); + expect(uploadedFileIds(result.current)).toStrictEqual(["file-a"]); + rerender({ orgId: "org-b" }); + await waitFor(() => { + expect(uploadedFileIds(result.current)).toStrictEqual([]); + }); + expect( + localStorage.getItem(persistedAttachmentsStorageKey) ?? "", + ).not.toContain("file-a"); + }); + + it("never exposes the previous org's file IDs in any render", async () => { + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([persistEntry("file-a", "a.txt", "org-a")]), + ); + // Log the hook output of every render, including the + // intermediate commit between the org changing and the + // adoption effect running; that window must expose nothing. + const renderLog: { orgId: string; fileIds: string[] }[] = []; + const Probe: FC<{ orgId: string }> = ({ orgId }) => { + const result = useFileAttachments(orgId, { persist: true }); + renderLog.push({ orgId, fileIds: uploadedFileIds(result) }); + return null; + }; + const { rerender } = render(); + expect(renderLog.at(-1)).toStrictEqual({ + orgId: "org-a", + fileIds: ["file-a"], + }); + + rerender(); + await waitFor(() => { + expect(renderLog.at(-1)?.fileIds).toStrictEqual([]); + }); + const leaked = renderLog.filter( + (entry) => entry.orgId === "org-b" && entry.fileIds.includes("file-a"), + ); + expect(leaked).toStrictEqual([]); + }); + + it("reports adoption only after the post-commit effect", () => { + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([persistEntry("file-a", "a.txt", "org-a")]), + ); + const log: { adopted: boolean; fileIds: string[] }[] = []; + const Probe: FC<{ orgId: string }> = ({ orgId }) => { + const result = useFileAttachments(orgId, { persist: true }); + log.push({ + adopted: result.organizationAdopted, + fileIds: uploadedFileIds(result), + }); + return null; + }; + render(); + expect(log[0]).toStrictEqual({ adopted: false, fileIds: [] }); + expect(log.at(-1)).toStrictEqual({ adopted: true, fileIds: ["file-a"] }); + }); + + it("discards an upload that completes after another org was adopted", async () => { + const resolveUpload = mockDeferredUpload(); + const { result, rerender } = renderAttachments({ orgId: "org-a" }); + act(() => { + result.current.startUpload(new File(["x"], "x.txt")); + }); + + rerender({ orgId: "org-b" }); + await waitFor(() => { + expect(result.current.organizationAdopted).toBe(true); + }); + await act(async () => { + resolveUpload({ id: "file-x" }); + }); + + expect(uploadedFileIds(result.current)).toStrictEqual([]); + expect( + localStorage.getItem(persistedAttachmentsStorageKey) ?? "", + ).not.toContain("file-x"); + }); + + it("discards an upload that spans an org round trip", async () => { + const resolveUpload = mockDeferredUpload(); + const { result, rerender } = renderAttachments({ orgId: "org-a" }); + act(() => { + result.current.startUpload(new File(["x"], "x.txt")); + }); + + // After A -> B -> A, only the adoption epoch distinguishes the stale upload. + rerender({ orgId: "org-b" }); + await waitFor(() => { + expect(result.current.organizationAdopted).toBe(true); + }); + rerender({ orgId: "org-a" }); + await waitFor(() => { + expect(result.current.organizationAdopted).toBe(true); + }); + await act(async () => { + resolveUpload({ id: "file-x" }); + }); + + expect(uploadedFileIds(result.current)).toStrictEqual([]); + expect( + localStorage.getItem(persistedAttachmentsStorageKey) ?? "", + ).not.toContain("file-x"); + }); + + it("hides attachments when authorization leaves no org", async () => { + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([persistEntry("file-a", "a.txt", "org-a")]), + ); + const { result, rerender } = renderAttachments({ orgId: "org-a" }); + await waitFor(() => { + expect(uploadedFileIds(result.current)).toStrictEqual(["file-a"]); + }); + + rerender({ orgId: undefined }); + expect(result.current.attachments).toStrictEqual([]); + expect(uploadedFileIds(result.current)).toStrictEqual([]); + expect(localStorage.getItem(persistedAttachmentsStorageKey)).toContain( + "file-a", + ); + + rerender({ orgId: "org-a" }); + await waitFor(() => { + expect(uploadedFileIds(result.current)).toStrictEqual(["file-a"]); + }); + }); + + it("does not prune storage during a render that never commits", () => { + localStorage.setItem( + persistedAttachmentsStorageKey, + JSON.stringify([persistEntry("file-a", "a.txt", "org-a")]), + ); + // Suspending after the hook runs simulates a render React abandons before commit. + const Suspender: FC<{ orgId: string }> = ({ orgId }) => { + useFileAttachments(orgId, { persist: true }); + throw new Promise(() => {}); + }; + render( + + + , + ); + expect(localStorage.getItem(persistedAttachmentsStorageKey)).toContain( + "file-a", + ); + }); +}); diff --git a/site/src/pages/AgentsPage/hooks/useFileAttachments.ts b/site/src/pages/AgentsPage/hooks/useFileAttachments.ts index eb6098c6312..b4b9193fb88 100644 --- a/site/src/pages/AgentsPage/hooks/useFileAttachments.ts +++ b/site/src/pages/AgentsPage/hooks/useFileAttachments.ts @@ -52,9 +52,7 @@ function restorePersistedAttachments(currentOrgId: string): { uploadStates: Map; previewUrls: Map; } { - // Skip when org ID isn't loaded yet so we don't prune valid - // entries. The initializer runs once, so callers must wait for - // the org ID before mounting. + // An unknown org must not prune entries persisted for the eventual org. if (!currentOrgId) { return { attachments: [], @@ -164,6 +162,11 @@ function clearPersistedAttachments() { } interface UseFileAttachmentsReturn { + /** + * True after the post-commit effect assigns in-memory attachment state to + * the supplied organization. Keep attach and send controls disabled until then. + */ + organizationAdopted: boolean; attachments: File[]; textContents: Map; uploadStates: Map; @@ -192,19 +195,14 @@ export function useFileAttachments( providerRef.current = provider; }, [provider]); - const [restored] = useState(() => - persist - ? restorePersistedAttachments(organizationId ?? "") - : { - attachments: [] as File[], - uploadStates: new Map(), - previewUrls: new Map(), - }, + const [attachments, setAttachments] = useState([]); + const [uploadStates, setUploadStates] = useState( + () => new Map(), ); - - const [attachments, setAttachments] = useState(restored.attachments); - const [uploadStates, setUploadStates] = useState(restored.uploadStates); - const [previewUrls, setPreviewUrls] = useState(restored.previewUrls); + const [previewUrls, setPreviewUrls] = useState(() => new Map()); + // Persisted state remains unowned until post-commit org adoption; restoring + // against a provisional org would prune entries for the eventual org. + const [stateOrgId, setStateOrgId] = useState(null); const [textContents, setTextContents] = useState( () => new Map(), ); @@ -220,6 +218,27 @@ export function useFileAttachments( return () => revokePreviewUrls(); }, []); + // Every adoption invalidates older upload completions, including A-to-B-to-A + // round trips where stateOrgId matches again. Otherwise a stale completion + // could persist and later restore an abandoned upload. + const adoptionEpochRef = useRef(0); + const commitUploadOutcome = useEffectEvent( + ( + file: File, + uploadOrgId: string, + uploadEpoch: number, + state: UploadState, + ) => { + if (persist && adoptionEpochRef.current !== uploadEpoch) { + return; + } + setUploadStates((prev) => new Map(prev).set(file, state)); + if (persist && state.status === "uploaded" && state.fileId) { + addPersistedAttachment(file, state.fileId, uploadOrgId); + } + }, + ); + const startUpload = (file: File) => { if (!organizationId) { setUploadStates((prev) => @@ -231,25 +250,18 @@ export function useFileAttachments( return; } - const shouldPersist = persist && Boolean(organizationId); + const uploadOrgId = organizationId; + const uploadEpoch = adoptionEpochRef.current; const isImage = file.type.startsWith("image/"); setUploadStates((prev) => new Map(prev).set(file, { status: "uploading" })); void (async () => { try { - const result = await API.experimental.uploadChatFile( - file, - organizationId, - ); - setUploadStates((prev) => - new Map(prev).set(file, { - status: "uploaded", - fileId: result.id, - }), - ); - if (shouldPersist) { - addPersistedAttachment(file, result.id, organizationId!); - } + const result = await API.experimental.uploadChatFile(file, uploadOrgId); + commitUploadOutcome(file, uploadOrgId, uploadEpoch, { + status: "uploaded", + fileId: result.id, + }); if (isImage) { // Pre-warm the HTTP cache so the timeline can // render the image instantly after send. Text @@ -257,13 +269,10 @@ export function useFileAttachments( void fetch(getChatFileURL(result.id)); } } catch (err: unknown) { - const errorMessage = formatAgentAttachmentUploadError(err); - setUploadStates((prev) => - new Map(prev).set(file, { - status: "error", - error: errorMessage, - }), - ); + commitUploadOutcome(file, uploadOrgId, uploadEpoch, { + status: "error", + error: formatAgentAttachmentUploadError(err), + }); } })(); }; @@ -273,6 +282,28 @@ export function useFileAttachments( // file can't be resurrected. WeakSet lets entries get GC'd. const abandonedResizesRef = useRef>(new WeakSet()); + // Permission refetches can change the org without user action. Replace state + // after commit so stale file IDs cannot cross orgs and an abandoned render + // cannot prune localStorage through restorePersistedAttachments. + const adoptOrganization = useEffectEvent((orgId: string) => { + adoptionEpochRef.current += 1; + for (const file of attachments) { + abandonedResizesRef.current.add(file); + } + revokePreviewUrls(); + setTextContents(new Map()); + setStateOrgId(orgId); + const restored = restorePersistedAttachments(orgId); + setAttachments(restored.attachments); + setUploadStates(restored.uploadStates); + setPreviewUrls(restored.previewUrls); + }); + useEffect(() => { + if (persist && organizationId && stateOrgId !== organizationId) { + adoptOrganization(organizationId); + } + }, [persist, stateOrgId, organizationId]); + type AttachItem = { file: File; needsResize: boolean }; const processResizes = async ( @@ -508,11 +539,18 @@ export function useFileAttachments( } }; + // Hide state that belongs to another organization. Exposing it could send + // stale file IDs or remove persisted attachments from the previous org. + const orgMismatch = + persist && stateOrgId !== null && stateOrgId !== organizationId; + return { - attachments, - textContents, - uploadStates, - previewUrls, + organizationAdopted: + !persist || (Boolean(organizationId) && stateOrgId === organizationId), + attachments: orgMismatch ? [] : attachments, + textContents: orgMismatch ? new Map() : textContents, + uploadStates: orgMismatch ? new Map() : uploadStates, + previewUrls: orgMismatch ? new Map() : previewUrls, handleAttach, handleRemoveAttachment, startUpload,