mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(site): prefer permitted organization for chat creation (#28078)
Fixes first-send 403s for multi-org users who lack `chat:create` in the deployment default organization (CODAGT-892). Two frontend defects combined to send chat creation to the wrong organization: 1. The Agents create form initialized its organization selection permission-blind to the default org and only corrected it when the `permittedOrganizations` authcheck result *transitioned*, so a fast first send raced the check, and on warm-query-cache remounts the correction never ran at all: the wrong org stayed selected permanently while the org picker was hidden. 2. The permitted-organizations authcheck itself sent no `owner_id`, so roles that grant `chat:create` at member (owner) scope, such as `agents-access` (Coder Agents User, the exact role in the customer report), were denied in every organization and the form always fell back to the default org. This part was split out and already landed on main via #28076; after rebasing, this PR relies on that fix and keeps its stricter regression stories around it. ## Changes - Derive the effective organization at render time: keep the user's explicit pick only while it is still permitted, otherwise fall back to the permitted default org, then the first permitted org, then the dashboard default. Replaces the transition-based reconciliation, which could not fire when the query cache was already warm on mount. - Keep user-driven org cleanup (workspace selection, attachments) in the picker/dialog event handlers; permission-driven changes are handled by render-time state adjustments and the attachment hook's post-commit adoption effect. - Disable Send until the permitted-organizations check settles and the attachment hook has adopted the effective org; hide the org picker and disable the workspace picker until the check settles (their pre-settlement options come from the unfiltered dashboard fallback, so a pick could persist a foreign-org workspace). - Scope persisted attachments to their organization in `useFileAttachments`: restoration defers until a permitted org is known, permission-driven org changes replace attachment state post-commit, in-flight uploads are invalidated by an adoption epoch (including A-to-B-to-A round trips), and no render exposes another org's file IDs, including when authorization resolves to no org at all. - Revalidate org-scoped state on permission refetches: a revoked explicit selection clears instead of lying latent, a settled effective-org change drops the stored workspace, and the org-change confirmation dialog closes (and re-checks on confirm) when its pending org is revoked. - Rebase reconciliation with #28076: the stories use its `permittedOrganizationsKey` helper and retain its `MemberScopedPermissionsShowOrgPicker` regression story alongside this PR's stricter member-scope stories. The backend RBAC rejection was correct; this is frontend-only. ## Testing - Red-green: every guard above was verified by reverting it and confirming exactly its guarding story or unit test fails (whole-file runs). - `pnpm -C site check`, `pnpm -C site lint`, `pnpm -C site lint:types`, full `AgentCreateForm.stories.tsx` (42 pass), `useFileAttachments.test.tsx` (8 pass), re-run after the rebase onto main. - Dogfood UAT on a licensed multi-org dev deployment at this branch: a restricted user with `agents-access` only in a non-default org sends first and warm-remount messages successfully (201, payload carries the permitted org, no 403); admin picker, workspace filtering, and attachment org-change dialog verified. Round 2 of UAT caught the missing `owner_id` (now landed via #28076); round 3 re-verified end to end. > Mux created this PR on Mike's behalf.
This commit is contained in:
@@ -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<typeof AgentCreateForm>;
|
||||
|
||||
const defaultArgs = meta.args;
|
||||
|
||||
const mockPermittedOrganizations = (permissions: Record<string, boolean>) => {
|
||||
const mockPermittedOrganizations = (
|
||||
permissions: Record<string, boolean>,
|
||||
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<string, boolean> = {};
|
||||
let revocableQueryClient: QueryClient | undefined;
|
||||
|
||||
const withRevocableQueryClient: Decorator = (Story) => {
|
||||
const [queryClient] = useState(
|
||||
() =>
|
||||
new QueryClient({
|
||||
defaultOptions: {
|
||||
queries: {
|
||||
staleTime: Number.POSITIVE_INFINITY,
|
||||
retry: false,
|
||||
},
|
||||
},
|
||||
}),
|
||||
);
|
||||
revocableQueryClient = queryClient;
|
||||
return (
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<Story />
|
||||
</QueryClientProvider>
|
||||
);
|
||||
};
|
||||
|
||||
const mockRevocablePermissions = (permissions: Record<string, boolean>) => {
|
||||
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<typeof fn>).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");
|
||||
},
|
||||
};
|
||||
|
||||
|
||||
@@ -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<AgentCreateFormProps> = ({
|
||||
: 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<string | null>(
|
||||
() => {
|
||||
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<TypesGen.Organization | null>(
|
||||
initialOrg ?? null,
|
||||
null,
|
||||
);
|
||||
const [pendingOrgChange, setPendingOrgChange] =
|
||||
useState<TypesGen.Organization | null>(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<string | null>(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<AgentCreateFormProps> = ({
|
||||
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<AgentCreateFormProps> = ({
|
||||
// enough to warrant pagination, this should switch to a
|
||||
// server-side organization:<name> 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<AgentCreateFormProps> = ({
|
||||
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<AgentCreateFormProps> = ({
|
||||
};
|
||||
|
||||
const {
|
||||
organizationAdopted,
|
||||
attachments,
|
||||
textContents,
|
||||
uploadStates,
|
||||
@@ -427,10 +487,17 @@ export const AgentCreateForm: FC<AgentCreateFormProps> = ({
|
||||
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<AgentCreateFormProps> = ({
|
||||
}
|
||||
};
|
||||
|
||||
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 (
|
||||
<>
|
||||
<div className="order-last flex min-h-0 flex-none items-end justify-center overflow-auto px-4 pb-4 sm:order-none sm:h-full sm:flex-1 sm:items-center">
|
||||
@@ -542,23 +565,27 @@ export const AgentCreateForm: FC<AgentCreateFormProps> = ({
|
||||
{permittedOrgsQuery.error != null && (
|
||||
<ErrorAlert error={permittedOrgsQuery.error} />
|
||||
)}
|
||||
{showOrganizations && permittedOrgs.length > 1 && (
|
||||
<CompactOrgSelector
|
||||
value={selectedOrg}
|
||||
options={permittedOrgs}
|
||||
onChange={(newOrg) => {
|
||||
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 && (
|
||||
<CompactOrgSelector
|
||||
value={effectiveOrg}
|
||||
options={permittedOrgs}
|
||||
onChange={(newOrg) => {
|
||||
const orgChanged = newOrg.id !== effectiveOrg?.id;
|
||||
if (orgChanged && attachments.length > 0) {
|
||||
setPendingOrgChange(newOrg);
|
||||
return;
|
||||
}
|
||||
if (orgChanged) {
|
||||
handleWorkspaceChange(null);
|
||||
}
|
||||
setSelectedOrg(newOrg);
|
||||
}}
|
||||
/>
|
||||
)}
|
||||
<AgentChatInput
|
||||
onSend={handleSendWithAttachments}
|
||||
sendShortcut={sendShortcut}
|
||||
@@ -566,6 +593,10 @@ export const AgentCreateForm: FC<AgentCreateFormProps> = ({
|
||||
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<AgentCreateFormProps> = ({
|
||||
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<AgentCreateFormProps> = ({
|
||||
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<AgentCreateFormProps> = ({
|
||||
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)}
|
||||
/>
|
||||
|
||||
@@ -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<typeof useFileAttachments>,
|
||||
): 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(<Probe orgId="org-a" />);
|
||||
expect(renderLog.at(-1)).toStrictEqual({
|
||||
orgId: "org-a",
|
||||
fileIds: ["file-a"],
|
||||
});
|
||||
|
||||
rerender(<Probe orgId="org-b" />);
|
||||
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(<Probe orgId="org-a" />);
|
||||
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(
|
||||
<Suspense fallback={null}>
|
||||
<Suspender orgId="org-b" />
|
||||
</Suspense>,
|
||||
);
|
||||
expect(localStorage.getItem(persistedAttachmentsStorageKey)).toContain(
|
||||
"file-a",
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -52,9 +52,7 @@ function restorePersistedAttachments(currentOrgId: string): {
|
||||
uploadStates: Map<File, UploadState>;
|
||||
previewUrls: Map<File, string>;
|
||||
} {
|
||||
// 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<File, string>;
|
||||
uploadStates: Map<File, UploadState>;
|
||||
@@ -192,19 +195,14 @@ export function useFileAttachments(
|
||||
providerRef.current = provider;
|
||||
}, [provider]);
|
||||
|
||||
const [restored] = useState(() =>
|
||||
persist
|
||||
? restorePersistedAttachments(organizationId ?? "")
|
||||
: {
|
||||
attachments: [] as File[],
|
||||
uploadStates: new Map<File, UploadState>(),
|
||||
previewUrls: new Map<File, string>(),
|
||||
},
|
||||
const [attachments, setAttachments] = useState<File[]>([]);
|
||||
const [uploadStates, setUploadStates] = useState(
|
||||
() => new Map<File, UploadState>(),
|
||||
);
|
||||
|
||||
const [attachments, setAttachments] = useState<File[]>(restored.attachments);
|
||||
const [uploadStates, setUploadStates] = useState(restored.uploadStates);
|
||||
const [previewUrls, setPreviewUrls] = useState(restored.previewUrls);
|
||||
const [previewUrls, setPreviewUrls] = useState(() => new Map<File, string>());
|
||||
// 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<string | null>(null);
|
||||
const [textContents, setTextContents] = useState(
|
||||
() => new Map<File, string>(),
|
||||
);
|
||||
@@ -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<WeakSet<File>>(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<File, string>() : textContents,
|
||||
uploadStates: orgMismatch ? new Map<File, UploadState>() : uploadStates,
|
||||
previewUrls: orgMismatch ? new Map<File, string>() : previewUrls,
|
||||
handleAttach,
|
||||
handleRemoveAttachment,
|
||||
startUpload,
|
||||
|
||||
Reference in New Issue
Block a user