diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 21b7853879..9a20437dfe 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -2,9 +2,11 @@ import { act, renderHook } from "@testing-library/react"; import { createRef } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ChatQueuedMessage } from "#/api/typesGenerated"; +import { MockUserOwner, MockWorkspace } from "#/testHelpers/entities"; import { draftInputStorageKeyPrefix, getPersistedDraftInputValue, + getWorkspaceOptionsWithLinkedWorkspace, restoreOptimisticRequestSnapshot, runPromoteQueuedMessage, submitEditAndScroll, @@ -93,6 +95,42 @@ const createDeferred = (): Deferred => { return { promise, resolve, reject }; }; +describe("getWorkspaceOptionsWithLinkedWorkspace", () => { + it("includes a missing linked workspace only when the current user owns it", () => { + const existingWorkspace = { + ...MockWorkspace, + id: "existing-workspace", + }; + const ownerWorkspaceOptions = [existingWorkspace]; + const linkedWorkspace = { + ...MockWorkspace, + id: "linked-workspace", + owner_id: MockUserOwner.id, + }; + + expect( + getWorkspaceOptionsWithLinkedWorkspace( + ownerWorkspaceOptions, + linkedWorkspace, + MockUserOwner.id, + ), + ).toEqual([linkedWorkspace, existingWorkspace]); + + const sharedWorkspace = { + ...linkedWorkspace, + owner_id: "another-user", + }; + + expect( + getWorkspaceOptionsWithLinkedWorkspace( + ownerWorkspaceOptions, + sharedWorkspace, + MockUserOwner.id, + ), + ).toBe(ownerWorkspaceOptions); + }); +}); + describe("waitForPendingChatSettingsSyncs", () => { it("waits for plan-mode and workspace updates before resolving", async () => { const planModeUpdate = createDeferred(); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index a6d968d8af..b4196b7280 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -275,6 +275,32 @@ export const filterWorkspaceOptionsByOrganization = ( ); }; +/** @internal Exported for testing. */ +export const getWorkspaceOptionsWithLinkedWorkspace = ( + workspaceOptions: readonly TypesGen.Workspace[], + workspace: TypesGen.Workspace | undefined, + ownerID: string, +): readonly TypesGen.Workspace[] => { + if (!workspace || workspace.owner_id !== ownerID) { + return workspaceOptions; + } + + const existingIndex = workspaceOptions.findIndex( + (candidate) => candidate.id === workspace.id, + ); + if (existingIndex === -1) { + return [workspace, ...workspaceOptions]; + } + + if (workspaceOptions[existingIndex] === workspace) { + return workspaceOptions; + } + + const nextWorkspaceOptions = [...workspaceOptions]; + nextWorkspaceOptions[existingIndex] = workspace; + return nextWorkspaceOptions; +}; + const buildAttachmentMediaTypes = ( attachments?: readonly PendingAttachment[], ): ReadonlyMap | undefined => { @@ -723,6 +749,7 @@ const AgentChatPage: FC = () => { ...workspaceById(workspaceId ?? ""), enabled: Boolean(workspaceId), }); + const workspace = workspaceQuery.data; const chatModelsQuery = useQuery(chatModels()); const chatModelConfigsQuery = useQuery(chatModelConfigs()); @@ -736,7 +763,11 @@ const AgentChatPage: FC = () => { const userDebugLoggingQuery = useQuery(userChatDebugLogging()); const mcpServersQuery = useQuery(mcpServerConfigs()); const workspacesQuery = useQuery(workspaces({ q: "owner:me", limit: 0 })); - const workspaceOptions = workspacesQuery.data?.workspaces ?? []; + const workspaceOptions = getWorkspaceOptionsWithLinkedWorkspace( + workspacesQuery.data?.workspaces ?? [], + workspace, + currentUser.id, + ); const desktopEnabled = desktopEnabledQuery.data?.enable_desktop ?? false; const debugLoggingEnabled = userDebugLoggingQuery.data?.debug_logging_enabled ?? false; @@ -835,7 +866,6 @@ const AgentChatPage: FC = () => { }); }, [workspaceId, queryClient]); const sshConfigQuery = useQuery(deploymentSSHConfig()); - const workspace = workspaceQuery.data; const workspaceAgent = getWorkspaceAgent(workspace, undefined); const { proxy } = useProxy();