From 97dde1f8246348c4b32172801f623268885dff33 Mon Sep 17 00:00:00 2001
From: Ethan <39577870+ethanndickson@users.noreply.github.com>
Date: Tue, 2 Jun 2026 14:37:12 +1000
Subject: [PATCH] fix: refresh attach workspace picker dynamically (#25834)

After the chat agent creates a workspace via the `create_workspace`
tool, opening the composer `+` menu and clicking "Attach workspace"
could show "No workspaces found" until a full page refresh, even though
the workspace pill already rendered the linked workspace correctly.

The picker was sourced only from the `owner:me` workspace list query,
whose cache could be stale right after `create_workspace` completed. The
fix derives the picker options at render time from both the owner
workspace list and the linked workspace already fetched by ID for the
pill, prepending or replacing the linked workspace only when the current
user owns it. This keeps the picker consistent with the pill without
broadening visibility beyond `owner:me` or invalidating workspace lists
on chat link updates.

Relates to CODAGT-510
---
 .../pages/AgentsPage/AgentChatPage.test.ts    | 38 +++++++++++++++++++
 site/src/pages/AgentsPage/AgentChatPage.tsx   | 34 ++++++++++++++++-
 2 files changed, 70 insertions(+), 2 deletions(-)

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 = <T>(): Deferred<T> => {
 	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<void>();
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<string, string> | 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();