From 7632e736080ce568e4bcf075cf7e2515a4c50aab Mon Sep 17 00:00:00 2001
From: Ethan <39577870+ethanndickson@users.noreply.github.com>
Date: Fri, 12 Jun 2026 11:51:20 +1000
Subject: [PATCH] refactor(site/src): extract shared workspace-app frame module
 (#26088)

Extracts the workspace app iframe, wildcard warning, and workspace-app
helper functions out of TaskPage into shared `site/src/modules/apps`
modules. Existing agent and app lookups in the task chat helpers,
download-logs dialog, and workspaces table now route through the shared
`workspaceApps` helpers instead of duplicating resource-flattening
logic. The extracted frame preserves the existing preview-only toolbar
behavior, and its open-in-new-tab link gains `rel="noreferrer"` to
harden against tabnabbing.

Relates to CODAGT-346
---
 .../apps/WorkspaceAppFrame.stories.tsx        |  45 ++++++++
 .../apps/WorkspaceAppFrame.tsx}               |  43 +++++---
 .../WorkspaceWildcardWarning.stories.tsx}     |  11 +-
 .../apps/WorkspaceWildcardWarning.tsx}        |   2 +-
 site/src/modules/apps/apps.test.ts            |  48 +++++++++
 site/src/modules/apps/apps.ts                 |  19 ++++
 site/src/modules/apps/workspaceApps.test.ts   | 100 ++++++++++++++++++
 site/src/modules/apps/workspaceApps.ts        |  37 +++++++
 .../src/modules/resources/AppLink/AppLink.tsx |   8 +-
 site/src/modules/tasks/apps.ts                |  22 ----
 .../DownloadLogsDialog.tsx                    |  10 +-
 .../ChatConversation/chatHelpers.ts           |   5 +-
 site/src/pages/TaskPage/TaskApps.tsx          |  11 +-
 site/src/pages/TaskPage/TaskPage.tsx          |  13 +--
 .../pages/WorkspacesPage/WorkspacesTable.tsx  |  16 +--
 site/src/utils/workspace.test.ts              |  76 +++++++++++++
 site/src/utils/workspace.tsx                  |  31 ++++--
 17 files changed, 411 insertions(+), 86 deletions(-)
 create mode 100644 site/src/modules/apps/WorkspaceAppFrame.stories.tsx
 rename site/src/{pages/TaskPage/TaskAppIframe.tsx => modules/apps/WorkspaceAppFrame.tsx} (77%)
 rename site/src/{pages/TaskPage/TaskWildcardWarning.stories.tsx => modules/apps/WorkspaceWildcardWarning.stories.tsx} (63%)
 rename site/src/{pages/TaskPage/TaskWildcardWarning.tsx => modules/apps/WorkspaceWildcardWarning.tsx} (96%)
 create mode 100644 site/src/modules/apps/workspaceApps.test.ts
 create mode 100644 site/src/modules/apps/workspaceApps.ts
 delete mode 100644 site/src/modules/tasks/apps.ts

diff --git a/site/src/modules/apps/WorkspaceAppFrame.stories.tsx b/site/src/modules/apps/WorkspaceAppFrame.stories.tsx
new file mode 100644
index 0000000000..f7e56fad84
--- /dev/null
+++ b/site/src/modules/apps/WorkspaceAppFrame.stories.tsx
@@ -0,0 +1,45 @@
+import type { Meta, StoryObj } from "@storybook/react-vite";
+import {
+	MockUserOwner,
+	MockWorkspace,
+	MockWorkspaceAgent,
+	MockWorkspaceApp,
+} from "#/testHelpers/entities";
+import { withAuthProvider, withProxyProvider } from "#/testHelpers/storybook";
+import { WorkspaceAppFrame } from "./WorkspaceAppFrame";
+import type { WorkspaceAppWithAgent } from "./workspaceApps";
+
+const meta: Meta<typeof WorkspaceAppFrame> = {
+	title: "modules/apps/WorkspaceAppFrame",
+	component: WorkspaceAppFrame,
+	decorators: [withAuthProvider, withProxyProvider()],
+	parameters: {
+		layout: "fullscreen",
+		user: MockUserOwner,
+	},
+	args: {
+		workspace: MockWorkspace,
+		app: buildWorkspaceApp(),
+		active: true,
+	},
+};
+
+export default meta;
+type Story = StoryObj<typeof meta>;
+
+export const Unhealthy: Story = {
+	args: {
+		app: buildWorkspaceApp({ health: "unhealthy" }),
+	},
+};
+
+function buildWorkspaceApp(
+	overrides: Partial<WorkspaceAppWithAgent> = {},
+): WorkspaceAppWithAgent {
+	return {
+		...MockWorkspaceApp,
+		agent: MockWorkspaceAgent,
+		health: "healthy",
+		...overrides,
+	};
+}
diff --git a/site/src/pages/TaskPage/TaskAppIframe.tsx b/site/src/modules/apps/WorkspaceAppFrame.tsx
similarity index 77%
rename from site/src/pages/TaskPage/TaskAppIframe.tsx
rename to site/src/modules/apps/WorkspaceAppFrame.tsx
index cfbca72746..4ff610f88a 100644
--- a/site/src/pages/TaskPage/TaskAppIframe.tsx
+++ b/site/src/modules/apps/WorkspaceAppFrame.tsx
@@ -3,7 +3,7 @@ import {
 	ExternalLinkIcon,
 	HouseIcon,
 } from "lucide-react";
-import { type FC, type HTMLProps, useRef } from "react";
+import { type ComponentProps, type FC, useRef } from "react";
 import { Link as RouterLink } from "react-router";
 import type { Workspace } from "#/api/typesGenerated";
 import { Button } from "#/components/Button/Button";
@@ -15,18 +15,20 @@ import {
 } from "#/components/DropdownMenu/DropdownMenu";
 import { Spinner } from "#/components/Spinner/Spinner";
 import { useProxy } from "#/contexts/ProxyContext";
-import { useAppLink } from "#/modules/apps/useAppLink";
-import type { WorkspaceAppWithAgent } from "#/modules/tasks/apps";
 import { cn } from "#/utils/cn";
-import { TaskWildcardWarning } from "./TaskWildcardWarning";
+import { isAppBlockedByMissingWildcard } from "./apps";
+import { useAppLink } from "./useAppLink";
+import { WorkspaceWildcardWarning } from "./WorkspaceWildcardWarning";
+import type { WorkspaceAppWithAgent } from "./workspaceApps";
 
-type TaskAppIFrameProps = {
+type WorkspaceAppFrameProps = {
 	workspace: Workspace;
 	app: WorkspaceAppWithAgent;
+	// Keep the iframe mounted while hidden so callers can preserve app state.
 	active: boolean;
 };
 
-export const TaskAppIFrame: FC<TaskAppIFrameProps> = ({
+export const WorkspaceAppFrame: FC<WorkspaceAppFrameProps> = ({
 	workspace,
 	app,
 	active,
@@ -37,20 +39,24 @@ export const TaskAppIFrame: FC<TaskAppIFrameProps> = ({
 	});
 	const proxy = useProxy();
 	const frameRef = useRef<HTMLIFrameElement>(null);
-	const shouldDisplayWildcardWarning =
-		app.subdomain && !proxy.proxy?.preferredWildcardHostname;
+	const shouldDisplayWildcardWarning = isAppBlockedByMissingWildcard(
+		app,
+		proxy.proxy?.preferredWildcardHostname,
+	);
+	// The "preview" app renders a navigation toolbar above its iframe.
+	const showToolbar = app.slug === "preview";
 
 	if (shouldDisplayWildcardWarning) {
 		return (
 			<div className="h-full flex items-center justify-center pb-4">
-				<TaskWildcardWarning />
+				<WorkspaceWildcardWarning />
 			</div>
 		);
 	}
 
 	return (
 		<div className={cn([active ? "flex" : "hidden", "w-full h-full flex-col"])}>
-			{app.slug === "preview" && (
+			{showToolbar && (
 				<div className="bg-surface-tertiary flex items-center p-2 py-1 gap-1">
 					<Button
 						size="icon"
@@ -66,9 +72,7 @@ export const TaskAppIFrame: FC<TaskAppIFrameProps> = ({
 						<span className="sr-only">Home</span>
 					</Button>
 
-					{/* Possibly we will put a URL bar here, but for now we cannot due to
-					 * cross-origin restrictions in iframes. */}
-					<div className="w-full"></div>
+					<div className="w-full" />
 
 					<DropdownMenu>
 						<DropdownMenuTrigger asChild>
@@ -79,7 +83,7 @@ export const TaskAppIFrame: FC<TaskAppIFrameProps> = ({
 						</DropdownMenuTrigger>
 						<DropdownMenuContent align="end">
 							<DropdownMenuItem asChild>
-								<RouterLink to={link.href} target="_blank">
+								<RouterLink to={link.href} target="_blank" rel="noreferrer">
 									<ExternalLinkIcon />
 									Open app in new tab
 								</RouterLink>
@@ -90,7 +94,7 @@ export const TaskAppIFrame: FC<TaskAppIFrameProps> = ({
 			)}
 
 			{app.health === "healthy" || app.health === "disabled" ? (
-				<TaskIframe ref={frameRef} src={link.href} title={link.label} />
+				<WorkspaceIframe ref={frameRef} src={link.href} title={link.label} />
 			) : app.health === "unhealthy" ? (
 				<div className="w-full h-full flex flex-col items-center justify-center p-4">
 					<h3 className="m-0 font-medium text-content-primary text-base text-center">
@@ -143,11 +147,16 @@ export const TaskAppIFrame: FC<TaskAppIFrameProps> = ({
 	);
 };
 
-type TaskIframeProps = HTMLProps<HTMLIFrameElement>;
+type WorkspaceIframeProps = ComponentProps<"iframe">;
 
-export const TaskIframe: FC<TaskIframeProps> = ({ className, ...props }) => {
+export const WorkspaceIframe: FC<WorkspaceIframeProps> = ({
+	className,
+	ref,
+	...props
+}) => {
 	return (
 		<iframe
+			ref={ref}
 			loading="eager"
 			className={cn("w-full h-full border-0", className)}
 			allow="clipboard-read; clipboard-write"
diff --git a/site/src/pages/TaskPage/TaskWildcardWarning.stories.tsx b/site/src/modules/apps/WorkspaceWildcardWarning.stories.tsx
similarity index 63%
rename from site/src/pages/TaskPage/TaskWildcardWarning.stories.tsx
rename to site/src/modules/apps/WorkspaceWildcardWarning.stories.tsx
index c98b2ccee3..c1c755f5db 100644
--- a/site/src/pages/TaskPage/TaskWildcardWarning.stories.tsx
+++ b/site/src/modules/apps/WorkspaceWildcardWarning.stories.tsx
@@ -1,12 +1,11 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import { MockUserOwner } from "#/testHelpers/entities";
 import { withAuthProvider } from "#/testHelpers/storybook";
-import type { TaskApps } from "./TaskApps";
-import { TaskWildcardWarning } from "./TaskWildcardWarning";
+import { WorkspaceWildcardWarning } from "./WorkspaceWildcardWarning";
 
-const meta: Meta<typeof TaskWildcardWarning> = {
-	title: "pages/TaskPage/TaskWildcardWarning",
-	component: TaskWildcardWarning,
+const meta: Meta<typeof WorkspaceWildcardWarning> = {
+	title: "modules/apps/WorkspaceWildcardWarning",
+	component: WorkspaceWildcardWarning,
 	decorators: [withAuthProvider],
 	parameters: {
 		layout: "fullscreen",
@@ -15,7 +14,7 @@ const meta: Meta<typeof TaskWildcardWarning> = {
 };
 
 export default meta;
-type Story = StoryObj<typeof TaskApps>;
+type Story = StoryObj<typeof meta>;
 
 export const WithoutEditPermission: Story = {};
 
diff --git a/site/src/pages/TaskPage/TaskWildcardWarning.tsx b/site/src/modules/apps/WorkspaceWildcardWarning.tsx
similarity index 96%
rename from site/src/pages/TaskPage/TaskWildcardWarning.tsx
rename to site/src/modules/apps/WorkspaceWildcardWarning.tsx
index 04095e8a9d..b5ded62350 100644
--- a/site/src/pages/TaskPage/TaskWildcardWarning.tsx
+++ b/site/src/modules/apps/WorkspaceWildcardWarning.tsx
@@ -4,7 +4,7 @@ import { Button } from "#/components/Button/Button";
 import { useAuthenticated } from "#/hooks/useAuthenticated";
 import { docs } from "#/utils/docs";
 
-export const TaskWildcardWarning = () => {
+export const WorkspaceWildcardWarning = () => {
 	const { permissions } = useAuthenticated();
 
 	return (
diff --git a/site/src/modules/apps/apps.test.ts b/site/src/modules/apps/apps.test.ts
index 6f8bc73f67..d9ac821cc8 100644
--- a/site/src/modules/apps/apps.test.ts
+++ b/site/src/modules/apps/apps.test.ts
@@ -1,3 +1,4 @@
+import type { WorkspaceApp } from "#/api/typesGenerated";
 import {
 	MockWorkspace,
 	MockWorkspaceAgent,
@@ -6,10 +7,20 @@ import {
 import {
 	getAppHref,
 	getVSCodeHref,
+	isAppBlockedByMissingWildcard,
+	isWorkspaceAppEmbeddable,
 	openAppInNewWindow,
 	SESSION_TOKEN_PLACEHOLDER,
 } from "./apps";
 
+function buildApp(overrides: Partial<WorkspaceApp> = {}): WorkspaceApp {
+	return {
+		...MockWorkspaceApp,
+		health: "healthy",
+		...overrides,
+	};
+}
+
 describe("getVSCodeHref", () => {
 	it("includes the chat ID when provided", () => {
 		const folder = "/workspace/test";
@@ -219,3 +230,40 @@ describe("openAppInNewWindow", () => {
 		expect(popup.location.href).toBe("https://app.example.com");
 	});
 });
+
+describe("isWorkspaceAppEmbeddable", () => {
+	it("returns true for visible path-based apps", () => {
+		expect(isWorkspaceAppEmbeddable(buildApp())).toBe(true);
+	});
+
+	it("returns false for command apps, hidden apps, and external apps", () => {
+		expect(isWorkspaceAppEmbeddable(buildApp({ command: "run-preview" }))).toBe(
+			false,
+		);
+		expect(isWorkspaceAppEmbeddable(buildApp({ hidden: true }))).toBe(false);
+		expect(
+			isWorkspaceAppEmbeddable(
+				buildApp({ external: true, url: "https://example.com" }),
+			),
+		).toBe(false);
+	});
+});
+
+describe("isAppBlockedByMissingWildcard", () => {
+	it("blocks subdomain apps when no wildcard host is configured", () => {
+		const subdomainApp = buildApp({ subdomain: true });
+
+		expect(isAppBlockedByMissingWildcard(subdomainApp, "")).toBe(true);
+		expect(isAppBlockedByMissingWildcard(subdomainApp, undefined)).toBe(true);
+		expect(
+			isAppBlockedByMissingWildcard(subdomainApp, "*.apps.example.com"),
+		).toBe(false);
+	});
+
+	it("never blocks path-based apps", () => {
+		const pathApp = buildApp({ subdomain: false });
+
+		expect(isAppBlockedByMissingWildcard(pathApp, "")).toBe(false);
+		expect(isAppBlockedByMissingWildcard(pathApp, undefined)).toBe(false);
+	});
+});
diff --git a/site/src/modules/apps/apps.ts b/site/src/modules/apps/apps.ts
index d3f5b7c8e2..f987a992cd 100644
--- a/site/src/modules/apps/apps.ts
+++ b/site/src/modules/apps/apps.ts
@@ -170,3 +170,22 @@ export const needsSessionToken = (app: ExternalWorkspaceApp) => {
 	const requiresSessionToken = app.url.includes(SESSION_TOKEN_PLACEHOLDER);
 	return requiresSessionToken && !isHttp;
 };
+
+/**
+ * True for apps that can be rendered inside a dashboard iframe. Command apps
+ * open in terminal tabs instead.
+ */
+export const isWorkspaceAppEmbeddable = (app: WorkspaceApp): boolean => {
+	return !app.hidden && !isExternalApp(app) && !app.command;
+};
+
+/**
+ * True when an app requires subdomain access but the deployment has no wildcard
+ * access URL configured, so the app cannot be launched or embedded.
+ */
+export const isAppBlockedByMissingWildcard = (
+	app: WorkspaceApp,
+	wildcardHostname: string | undefined,
+): boolean => {
+	return app.subdomain && !wildcardHostname;
+};
diff --git a/site/src/modules/apps/workspaceApps.test.ts b/site/src/modules/apps/workspaceApps.test.ts
new file mode 100644
index 0000000000..c28f4cd9c8
--- /dev/null
+++ b/site/src/modules/apps/workspaceApps.test.ts
@@ -0,0 +1,100 @@
+import type {
+	Workspace,
+	WorkspaceAgent,
+	WorkspaceApp,
+} from "#/api/typesGenerated";
+import {
+	MockWorkspace,
+	MockWorkspaceAgent,
+	MockWorkspaceApp,
+} from "#/testHelpers/entities";
+import {
+	findWorkspaceAppWithAgent,
+	getAllAppsWithAgent,
+} from "./workspaceApps";
+
+describe("workspaceApps", () => {
+	describe("getAllAppsWithAgent", () => {
+		it("flattens workspace apps with their owning agent", () => {
+			const workspace = buildWorkspace([
+				[buildAgent("agent-1", [buildApp("app-1")])],
+				[buildAgent("agent-2", [buildApp("app-2")])],
+			]);
+
+			expect(
+				getAllAppsWithAgent(workspace).map((app) => ({
+					appId: app.id,
+					agentId: app.agent.id,
+				})),
+			).toEqual([
+				{ appId: "app-1", agentId: "agent-1" },
+				{ appId: "app-2", agentId: "agent-2" },
+			]);
+		});
+
+		it("returns an empty list when the workspace has no agents", () => {
+			const workspace = buildWorkspace([]);
+
+			expect(getAllAppsWithAgent(workspace)).toEqual([]);
+		});
+	});
+
+	describe("findWorkspaceAppWithAgent", () => {
+		it("returns the matching app with its owning agent", () => {
+			const workspace = buildWorkspace([
+				[buildAgent("agent-1", [buildApp("app-1")])],
+				[buildAgent("agent-2", [buildApp("app-2")])],
+			]);
+
+			expect(
+				findWorkspaceAppWithAgent(workspace, "agent-2", "app-2"),
+			).toMatchObject({
+				id: "app-2",
+				agent: { id: "agent-2" },
+			});
+			expect(
+				findWorkspaceAppWithAgent(workspace, "agent-1", "app-2"),
+			).toBeUndefined();
+		});
+	});
+});
+
+function buildWorkspace(
+	resourceAgents: readonly WorkspaceAgent[][],
+): Workspace {
+	const resourceTemplate = MockWorkspace.latest_build.resources[0];
+	return {
+		...MockWorkspace,
+		latest_build: {
+			...MockWorkspace.latest_build,
+			resources: resourceAgents.map((agents) => ({
+				...resourceTemplate,
+				agents,
+			})),
+		},
+	};
+}
+
+function buildAgent(id: string, apps: WorkspaceApp[]): WorkspaceAgent {
+	return {
+		...MockWorkspaceAgent,
+		id,
+		name: id,
+		apps,
+	};
+}
+
+function buildApp(
+	id: string,
+	overrides: Partial<WorkspaceApp> = {},
+): WorkspaceApp {
+	return {
+		...MockWorkspaceApp,
+		id,
+		slug: id,
+		display_name: id,
+		health: "healthy",
+		statuses: [],
+		...overrides,
+	};
+}
diff --git a/site/src/modules/apps/workspaceApps.ts b/site/src/modules/apps/workspaceApps.ts
new file mode 100644
index 0000000000..0869206562
--- /dev/null
+++ b/site/src/modules/apps/workspaceApps.ts
@@ -0,0 +1,37 @@
+import type {
+	Workspace,
+	WorkspaceAgent,
+	WorkspaceApp,
+} from "#/api/typesGenerated";
+import { findWorkspaceAgent, getWorkspaceAgents } from "#/utils/workspace";
+
+export type WorkspaceAppWithAgent = WorkspaceApp & {
+	agent: WorkspaceAgent;
+};
+
+export function getAllAppsWithAgent(
+	workspace: Workspace,
+): WorkspaceAppWithAgent[] {
+	return getWorkspaceAgents(workspace).flatMap((agent) =>
+		agent.apps.map((app) => ({
+			...app,
+			agent,
+		})),
+	);
+}
+
+export function findWorkspaceAppWithAgent(
+	workspace: Workspace,
+	agentId: string,
+	appId: string,
+): WorkspaceAppWithAgent | undefined {
+	const agent = findWorkspaceAgent(workspace, agentId);
+	if (!agent) {
+		return undefined;
+	}
+	const app = agent.apps.find((workspaceApp) => workspaceApp.id === appId);
+	if (!app) {
+		return undefined;
+	}
+	return { ...app, agent };
+}
diff --git a/site/src/modules/resources/AppLink/AppLink.tsx b/site/src/modules/resources/AppLink/AppLink.tsx
index 5b8d7aa56f..93319ab008 100644
--- a/site/src/modules/resources/AppLink/AppLink.tsx
+++ b/site/src/modules/resources/AppLink/AppLink.tsx
@@ -18,7 +18,11 @@ import {
 	TooltipTrigger,
 } from "#/components/Tooltip/Tooltip";
 import { useProxy } from "#/contexts/ProxyContext";
-import { isExternalApp, needsSessionToken } from "#/modules/apps/apps";
+import {
+	isAppBlockedByMissingWildcard,
+	isExternalApp,
+	needsSessionToken,
+} from "#/modules/apps/apps";
 import { useAppLink } from "#/modules/apps/useAppLink";
 import { docs } from "#/utils/docs";
 import { AgentButton } from "../AgentButton";
@@ -76,7 +80,7 @@ export const AppLink: FC<AppLinkProps> = ({
 		primaryTooltip = "Unhealthy";
 	}
 
-	if (!host && app.subdomain) {
+	if (isAppBlockedByMissingWildcard(app, host)) {
 		canClick = false;
 		icon = (
 			<CircleAlertIcon
diff --git a/site/src/modules/tasks/apps.ts b/site/src/modules/tasks/apps.ts
deleted file mode 100644
index 52200819b4..0000000000
--- a/site/src/modules/tasks/apps.ts
+++ /dev/null
@@ -1,22 +0,0 @@
-import type {
-	Workspace,
-	WorkspaceAgent,
-	WorkspaceApp,
-} from "#/api/typesGenerated";
-
-export type WorkspaceAppWithAgent = WorkspaceApp & {
-	agent: WorkspaceAgent;
-};
-
-export function getAllAppsWithAgent(
-	workspace: Workspace,
-): WorkspaceAppWithAgent[] {
-	return workspace.latest_build.resources
-		.flatMap((r) => r.agents ?? [])
-		.flatMap((agent) =>
-			agent.apps.map((app) => ({
-				...app,
-				agent,
-			})),
-		);
-}
diff --git a/site/src/modules/workspaces/WorkspaceMoreActions/DownloadLogsDialog.tsx b/site/src/modules/workspaces/WorkspaceMoreActions/DownloadLogsDialog.tsx
index 044134e321..004913c8ec 100644
--- a/site/src/modules/workspaces/WorkspaceMoreActions/DownloadLogsDialog.tsx
+++ b/site/src/modules/workspaces/WorkspaceMoreActions/DownloadLogsDialog.tsx
@@ -13,6 +13,7 @@ import {
 } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog";
 import { Skeleton } from "#/components/Skeleton/Skeleton";
 import { cn } from "#/utils/cn";
+import { getWorkspaceAgents } from "#/utils/workspace";
 
 type DownloadLogsDialogProps = Pick<
 	ConfirmDialogProps,
@@ -39,15 +40,12 @@ export const DownloadLogsDialog: FC<DownloadLogsDialogProps> = ({
 	});
 
 	const allUniqueAgents = useMemo<readonly WorkspaceAgent[]>(() => {
-		const allAgents = workspace.latest_build.resources.flatMap(
-			(resource) => resource.agents ?? [],
-		);
+		const allAgents = getWorkspaceAgents(workspace);
 
 		// Can't use the "new Set()" trick because we're not dealing with primitives
 		const uniqueAgents = new Map(allAgents.map((agent) => [agent.id, agent]));
-		const iterable = [...uniqueAgents.values()];
-		return iterable;
-	}, [workspace.latest_build.resources]);
+		return [...uniqueAgents.values()];
+	}, [workspace]);
 
 	const agentLogQueries = useQueries({
 		queries: allUniqueAgents.map((agent) => ({
diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts
index 18f21df8ea..e0877b0525 100644
--- a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts
+++ b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts
@@ -1,4 +1,5 @@
 import type * as TypesGen from "#/api/typesGenerated";
+import { getWorkspaceAgents } from "#/utils/workspace";
 import type { AgentContextUsage } from "../AgentChatInput";
 import type { ModelSelectorOption } from "../ChatElements";
 import { asString } from "../ChatElements/runtimeTypeUtils";
@@ -114,9 +115,7 @@ export const getWorkspaceAgent = (
 	if (!workspace) {
 		return undefined;
 	}
-	const agents = workspace.latest_build.resources.flatMap(
-		(resource) => resource.agents ?? [],
-	);
+	const agents = getWorkspaceAgents(workspace);
 	if (agents.length === 0) {
 		return undefined;
 	}
diff --git a/site/src/pages/TaskPage/TaskApps.tsx b/site/src/pages/TaskPage/TaskApps.tsx
index b195c5c395..d415514be2 100644
--- a/site/src/pages/TaskPage/TaskApps.tsx
+++ b/site/src/pages/TaskPage/TaskApps.tsx
@@ -15,13 +15,16 @@ import { Link } from "#/components/Link/Link";
 import { ScrollArea, ScrollBar } from "#/components/ScrollArea/ScrollArea";
 import { getTerminalHref } from "#/modules/apps/apps";
 import { useAppLink } from "#/modules/apps/useAppLink";
+import {
+	WorkspaceAppFrame,
+	WorkspaceIframe,
+} from "#/modules/apps/WorkspaceAppFrame";
 import {
 	getAllAppsWithAgent,
 	type WorkspaceAppWithAgent,
-} from "#/modules/tasks/apps";
+} from "#/modules/apps/workspaceApps";
 import { cn } from "#/utils/cn";
 import { docs } from "#/utils/docs";
-import { TaskAppIFrame, TaskIframe } from "./TaskAppIframe";
 
 type TaskAppsProps = {
 	task: Task;
@@ -94,7 +97,7 @@ export const TaskApps: FC<TaskAppsProps> = ({ task, workspace }) => {
 			{embeddedApps.length > 0 ? (
 				<div className="flex-1">
 					{embeddedApps.map((app) => (
-						<TaskAppIFrame
+						<WorkspaceAppFrame
 							key={app.id}
 							active={activeAppId === app.id}
 							app={app}
@@ -102,7 +105,7 @@ export const TaskApps: FC<TaskAppsProps> = ({ task, workspace }) => {
 						/>
 					))}
 
-					<TaskIframe
+					<WorkspaceIframe
 						src={terminalHref}
 						title="Terminal"
 						className={cn({
diff --git a/site/src/pages/TaskPage/TaskPage.tsx b/site/src/pages/TaskPage/TaskPage.tsx
index 0efc87cca5..d6bc91a4c6 100644
--- a/site/src/pages/TaskPage/TaskPage.tsx
+++ b/site/src/pages/TaskPage/TaskPage.tsx
@@ -43,9 +43,10 @@ import { Margins } from "#/components/Margins/Margins";
 import { ScrollArea } from "#/components/ScrollArea/ScrollArea";
 import { Spinner } from "#/components/Spinner/Spinner";
 import { useWorkspaceBuildLogs } from "#/hooks/useWorkspaceBuildLogs";
+import { WorkspaceAppFrame } from "#/modules/apps/WorkspaceAppFrame";
+import { getAllAppsWithAgent } from "#/modules/apps/workspaceApps";
 import { AgentLogs } from "#/modules/resources/AgentLogs/AgentLogs";
 import { useAgentLogs } from "#/modules/resources/useAgentLogs";
-import { getAllAppsWithAgent } from "#/modules/tasks/apps";
 import { TasksSidebar } from "#/modules/tasks/TasksSidebar/TasksSidebar";
 import { isPauseDisabled } from "#/modules/tasks/taskActions";
 import { WorkspaceErrorDialog } from "#/modules/workspaces/ErrorDialog/WorkspaceErrorDialog";
@@ -54,13 +55,13 @@ import { WorkspaceOutdatedTooltip } from "#/modules/workspaces/WorkspaceOutdated
 import { cn } from "#/utils/cn";
 import { pageTitle } from "#/utils/page";
 import { relativeTime } from "#/utils/time";
+import { getWorkspaceAgents } from "#/utils/workspace";
 import {
 	getActiveTransitionStats,
 	WorkspaceBuildProgress,
 } from "../WorkspacePage/WorkspaceBuildProgress";
 import { FollowUpDialog } from "./FollowUpDialog";
 import { ModifyPromptDialog } from "./ModifyPromptDialog";
-import { TaskAppIFrame } from "./TaskAppIframe";
 import { TaskApps } from "./TaskApps";
 import { TaskTopbar } from "./TaskTopbar";
 
@@ -334,7 +335,7 @@ const TaskPage = () => {
 			<PanelGroup autoSaveId="task" direction="horizontal">
 				<Panel defaultSize={25} minSize={20}>
 					{chatApp ? (
-						<TaskAppIFrame active workspace={workspace} app={chatApp} />
+						<WorkspaceAppFrame active workspace={workspace} app={chatApp} />
 					) : (
 						<div className="h-full flex items-center justify-center p-6 text-center">
 							<div className="flex flex-col items-center">
@@ -898,9 +899,5 @@ const TaskStartingAgent: FC<TaskStartingAgentProps> = ({ task, agent }) => {
 };
 
 function selectAgent(workspace: Workspace) {
-	const agents = workspace.latest_build.resources
-		.flatMap((r) => r.agents)
-		.filter(Boolean);
-
-	return agents.at(0);
+	return getWorkspaceAgents(workspace).at(0);
 }
diff --git a/site/src/pages/WorkspacesPage/WorkspacesTable.tsx b/site/src/pages/WorkspacesPage/WorkspacesTable.tsx
index bcf0571d90..ae33dec515 100644
--- a/site/src/pages/WorkspacesPage/WorkspacesTable.tsx
+++ b/site/src/pages/WorkspacesPage/WorkspacesTable.tsx
@@ -70,6 +70,7 @@ import {
 	openAppInNewWindow,
 } from "#/modules/apps/apps";
 import { useAppLink } from "#/modules/apps/useAppLink";
+import { findWorkspaceAppWithAgent } from "#/modules/apps/workspaceApps";
 import { useDashboard } from "#/modules/dashboard/useDashboard";
 import { abilitiesByWorkspaceStatus } from "#/modules/workspaces/actions";
 import { WorkspaceBuildCancelDialog } from "#/modules/workspaces/WorkspaceBuildCancelDialog/WorkspaceBuildCancelDialog";
@@ -759,15 +760,18 @@ const WorkspaceAppStatusLinks: FC<WorkspaceAppStatusLinksProps> = ({
 	workspace,
 }) => {
 	const status = workspace.latest_app_status;
-	const agent = workspace.latest_build.resources
-		.flatMap((r) => r.agents)
-		.find((a) => a?.id === status?.agent_id);
-	const app = agent?.apps.find((a) => a.id === status?.app_id);
+	const appWithAgent = status
+		? findWorkspaceAppWithAgent(workspace, status.agent_id, status.app_id)
+		: undefined;
 
 	return (
 		<>
-			{agent && app && (
-				<IconAppLink app={app} workspace={workspace} agent={agent} />
+			{appWithAgent && (
+				<IconAppLink
+					app={appWithAgent}
+					workspace={workspace}
+					agent={appWithAgent.agent}
+				/>
 			)}
 
 			{status?.uri && status?.uri !== "n/a" && (
diff --git a/site/src/utils/workspace.test.ts b/site/src/utils/workspace.test.ts
index cb7480ed9e..a011bfe652 100644
--- a/site/src/utils/workspace.test.ts
+++ b/site/src/utils/workspace.test.ts
@@ -4,12 +4,39 @@ import * as Mocks from "#/testHelpers/entities";
 import {
 	agentVersionStatus,
 	defaultWorkspaceExtension,
+	findWorkspaceAgent,
 	getDisplayVersionStatus,
 	getDisplayWorkspaceBuildInitiatedBy,
 	getDisplayWorkspaceTemplateName,
+	getMatchingAgentOrFirst,
+	getWorkspaceAgents,
 	isWorkspaceOn,
 } from "./workspace";
 
+function buildWorkspace(
+	resourceAgents: readonly TypesGen.WorkspaceAgent[][],
+): TypesGen.Workspace {
+	const resourceTemplate = Mocks.MockWorkspace.latest_build.resources[0];
+	return {
+		...Mocks.MockWorkspace,
+		latest_build: {
+			...Mocks.MockWorkspace.latest_build,
+			resources: resourceAgents.map((agents) => ({
+				...resourceTemplate,
+				agents,
+			})),
+		},
+	};
+}
+
+function buildAgent(id: string): TypesGen.WorkspaceAgent {
+	return {
+		...Mocks.MockWorkspaceAgent,
+		id,
+		name: id,
+	};
+}
+
 describe("util > workspace", () => {
 	describe("isWorkspaceOn", () => {
 		it.each<
@@ -148,4 +175,53 @@ describe("util > workspace", () => {
 			expect(displayed).toEqual(workspace.template_display_name);
 		});
 	});
+
+	describe("getWorkspaceAgents", () => {
+		it("flattens agents across workspace resources", () => {
+			const workspace = buildWorkspace([
+				[buildAgent("agent-1")],
+				[buildAgent("agent-2")],
+			]);
+
+			expect(getWorkspaceAgents(workspace).map((agent) => agent.id)).toEqual([
+				"agent-1",
+				"agent-2",
+			]);
+		});
+	});
+
+	describe("findWorkspaceAgent", () => {
+		it("returns the matching agent by ID", () => {
+			const workspace = buildWorkspace([[buildAgent("agent-1")]]);
+
+			expect(findWorkspaceAgent(workspace, "agent-1")?.name).toBe("agent-1");
+			expect(findWorkspaceAgent(workspace, "missing")).toBeUndefined();
+		});
+	});
+
+	describe("getMatchingAgentOrFirst", () => {
+		it("returns the agent matching by name across resources", () => {
+			const workspace = buildWorkspace([
+				[buildAgent("agent-1")],
+				[buildAgent("agent-2")],
+			]);
+
+			expect(getMatchingAgentOrFirst(workspace, "agent-2")?.id).toBe("agent-2");
+		});
+
+		it("returns the first agent when no name is given", () => {
+			const workspace = buildWorkspace([
+				[buildAgent("agent-1")],
+				[buildAgent("agent-2")],
+			]);
+
+			expect(getMatchingAgentOrFirst(workspace, undefined)?.id).toBe("agent-1");
+		});
+
+		it("returns undefined when no agent matches the name", () => {
+			const workspace = buildWorkspace([[buildAgent("agent-1")]]);
+
+			expect(getMatchingAgentOrFirst(workspace, "missing")).toBeUndefined();
+		});
+	});
 });
diff --git a/site/src/utils/workspace.tsx b/site/src/utils/workspace.tsx
index 9bda4f60a0..4ec7ba6c86 100644
--- a/site/src/utils/workspace.tsx
+++ b/site/src/utils/workspace.tsx
@@ -239,21 +239,30 @@ export const getDisplayWorkspaceStatus = (
 	}
 };
 
+export const getWorkspaceAgents = (
+	workspace: TypesGen.Workspace,
+): TypesGen.WorkspaceAgent[] => {
+	return workspace.latest_build.resources.flatMap(
+		(resource) => resource.agents ?? [],
+	);
+};
+
+export const findWorkspaceAgent = (
+	workspace: TypesGen.Workspace,
+	agentId: string,
+): TypesGen.WorkspaceAgent | undefined => {
+	return getWorkspaceAgents(workspace).find((agent) => agent.id === agentId);
+};
+
 export const getMatchingAgentOrFirst = (
 	workspace: TypesGen.Workspace,
 	agentName: string | undefined,
 ): TypesGen.WorkspaceAgent | undefined => {
-	return workspace.latest_build.resources
-		.map((resource) => {
-			if (!resource.agents || resource.agents.length === 0) {
-				return;
-			}
-			if (!agentName) {
-				return resource.agents[0];
-			}
-			return resource.agents.find((agent) => agent.name === agentName);
-		})
-		.filter((a) => a)[0];
+	const agents = getWorkspaceAgents(workspace);
+	if (!agentName) {
+		return agents[0];
+	}
+	return agents.find((agent) => agent.name === agentName);
 };
 
 export const mustUpdateWorkspace = (