Simplify PortForwardButton

coder · BrunoQuaresma · May 29, 2025 · May 26, 2025 · May 28, 2025 · May 28, 2025
commit 54bf987d1acf6b52e7d8d993e637e134e8c4c8b6
diff --git a/site/src/modules/resources/AgentMetadata.tsx b/site/src/modules/resources/AgentMetadata.tsx
@@ -42,24 +42,24 @@ export const AgentMetadataView: FC<AgentMetadataViewProps> = ({ metadata }) => {
 
 interface AgentMetadataProps {
 	agent: WorkspaceAgent;
-	storybookMetadata?: WorkspaceAgentMetadata[];
+	initialMetadata?: WorkspaceAgentMetadata[];
 }
 
 const maxSocketErrorRetryCount = 3;
 
 export const AgentMetadata: FC<AgentMetadataProps> = ({
 	agent,
-	storybookMetadata,
+	initialMetadata,
 }) => {
-	const [activeMetadata, setActiveMetadata] = useState(storybookMetadata);
+	const [activeMetadata, setActiveMetadata] = useState(initialMetadata);
 	useEffect(() => {
 		// This is an unfortunate pitfall with this component's testing setup,
-		// but even though we use the value of storybookMetadata as the initial
+		// but even though we use the value of initialMetadata as the initial
 		// value of the activeMetadata, we cannot put activeMetadata itself into
 		// the dependency array. If we did, we would destroy and rebuild each
 		// connection every single time a new message comes in from the socket,
 		// because the socket has to be wired up to the state setter
-		if (storybookMetadata !== undefined) {
+		if (initialMetadata !== undefined) {
 			return;
 		}
 
@@ -118,7 +118,7 @@ export const AgentMetadata: FC<AgentMetadataProps> = ({
 			window.clearTimeout(timeoutId);
 			activeSocket?.close();
 		};
-	}, [agent.id, storybookMetadata]);
+	}, [agent.id, initialMetadata]);
 
 	if (activeMetadata === undefined) {
 		return (

diff --git a/site/src/modules/resources/AgentRow.stories.tsx b/site/src/modules/resources/AgentRow.stories.tsx
@@ -96,7 +96,7 @@ const meta: Meta<typeof AgentRow> = {
 		},
 		workspace: M.MockWorkspace,
 		showApps: true,
-		storybookAgentMetadata: defaultAgentMetadata,
+		initialMetadata: defaultAgentMetadata,
 	},
 	decorators: [withProxyProvider(), withDashboardProvider, withWebSocket],
 	parameters: {
@@ -162,7 +162,7 @@ export const BunchOfApps: Story = {
 export const Connecting: Story = {
 	args: {
 		agent: M.MockWorkspaceAgentConnecting,
-		storybookAgentMetadata: [],
+		initialMetadata: [],
 	},
 };
 
@@ -190,7 +190,7 @@ export const Started: Story = {
 export const StartedNoMetadata: Story = {
 	args: {
 		...Started.args,
-		storybookAgentMetadata: [],
+		initialMetadata: [],
 	},
 };
 

diff --git a/site/src/modules/resources/AgentRow.tsx b/site/src/modules/resources/AgentRow.tsx
@@ -43,14 +43,12 @@ import { TerminalLink } from "./TerminalLink/TerminalLink";
 import { VSCodeDesktopButton } from "./VSCodeDesktopButton/VSCodeDesktopButton";
 import { useAgentLogs } from "./useAgentLogs";
 import { useFeatureVisibility } from "modules/dashboard/useFeatureVisibility";
-import { useEmbeddedMetadata } from "hooks/useEmbeddedMetadata";
-import { buildInfo } from "api/queries/buildInfo";
 
 export interface AgentRowProps {
 	agent: WorkspaceAgent;
 	workspace: Workspace;
 	template: Template;
-	storybookAgentMetadata?: WorkspaceAgentMetadata[];
+	initialMetadata?: WorkspaceAgentMetadata[];
 	onUpdateAgent: () => void;
 }
 
@@ -59,7 +57,7 @@ export const AgentRow: FC<AgentRowProps> = ({
 	workspace,
 	template,
 	onUpdateAgent,
-	storybookAgentMetadata,
+	initialMetadata,
 }) => {
 	// Apps visibility
 	const { browser_only } = useFeatureVisibility();
@@ -196,10 +194,8 @@ export const AgentRow: FC<AgentRowProps> = ({
 						agent.display_apps.includes("port_forwarding_helper") && (
 							<PortForwardButton
 								host={proxy.preferredWildcardHostname}
-								workspaceName={workspace.name}
+								workspace={workspace}
 								agent={agent}
-								username={workspace.owner_name}
-								workspaceID={workspace.id}
 								template={template}
 							/>
 						)}
@@ -281,10 +277,7 @@ export const AgentRow: FC<AgentRowProps> = ({
 					</section>
 				)}
 
-				<AgentMetadata
-					storybookMetadata={storybookAgentMetadata}
-					agent={agent}
-				/>
+				<AgentMetadata initialMetadata={initialMetadata} agent={agent} />
 			</div>
 
 			{hasStartupFeatures && (

diff --git a/site/src/modules/resources/PortForwardButton.tsx b/site/src/modules/resources/PortForwardButton.tsx
@@ -18,8 +18,10 @@ import {
 import {
 	type Template,
 	type UpsertWorkspaceAgentPortShareRequest,
+	type Workspace,
 	type WorkspaceAgent,
 	type WorkspaceAgentListeningPort,
+	type WorkspaceAgentPortShare,
 	type WorkspaceAgentPortShareLevel,
 	type WorkspaceAgentPortShareProtocol,
 	WorkspaceAppSharingLevels,
@@ -64,121 +66,131 @@ import * as Yup from "yup";
 
 interface PortForwardButtonProps {
 	host: string;
-	username: string;
-	workspaceName: string;
-	workspaceID: string;
+	workspace: Workspace;
 	agent: WorkspaceAgent;
 	template: Template;
 }
 
-export const PortForwardButton: FC<PortForwardButtonProps> = (props) => {
-	const { agent } = props;
+export const PortForwardButton: FC<PortForwardButtonProps> = ({
+	host,
+	workspace,
+	template,
+	agent,
+}) => {
 	const { entitlements } = useDashboard();
 	const paper = useClassName(classNames.paper, []);
 
-	const portsQuery = useQuery({
+	const { data: listeningPorts } = useQuery({
 		queryKey: ["portForward", agent.id],
 		queryFn: () => API.getAgentListeningPorts(agent.id),
 		enabled: agent.status === "connected",
 		refetchInterval: 5_000,
+		select: (res) => res.ports,
+	});
+
+	const { data: sharedPorts, refetch: refetchSharedPorts } = useQuery({
+		...workspacePortShares(workspace.id),
+		enabled: agent.status === "connected",
+		select: (res) => res.shares,
+		initialData: { shares: [] },
 	});
 
 	return (
 		<Popover>
 			<PopoverTrigger>
-				<Button disabled={!portsQuery.data} size="sm" variant="subtle">
-					<Spinner loading={!portsQuery.data}>
-						<span css={styles.portCount}>{portsQuery.data?.ports.length}</span>
+				<Button disabled={!listeningPorts} size="sm" variant="subtle">
+					<Spinner loading={!listeningPorts}>
+						<span css={styles.portCount}>{listeningPorts?.length}</span>
 					</Spinner>
 					Open ports
 					<ChevronDownIcon className="size-4" />
 				</Button>
 			</PopoverTrigger>
 			<PopoverContent horizontal="right" classes={{ paper }}>
 				<PortForwardPopoverView
-					{...props}
-					listeningPorts={portsQuery.data?.ports}
+					host={host}
+					agent={agent}
+					workspace={workspace}
+					template={template}
+					sharedPorts={sharedPorts}
+					listeningPorts={listeningPorts ?? []}
 					portSharingControlsEnabled={
 						entitlements.features.control_shared_ports.enabled
 					}
+					refetchSharedPorts={refetchSharedPorts}
 				/>
 			</PopoverContent>
 		</Popover>
 	);
 };
 
-const getValidationSchema = (): Yup.AnyObjectSchema =>
+const openPortSchema = (): Yup.AnyObjectSchema =>
 	Yup.object({
 		port: Yup.number().required().min(9).max(65535),
 		share_level: Yup.string().required().oneOf(WorkspaceAppSharingLevels),
 	});
 
-interface PortForwardPopoverViewProps extends PortForwardButtonProps {
-	listeningPorts?: readonly WorkspaceAgentListeningPort[];
+interface PortForwardPopoverViewProps {
+	host: string;
+	workspace: Workspace;
+	agent: WorkspaceAgent;
+	template: Template;
+	sharedPorts: readonly WorkspaceAgentPortShare[];
+	listeningPorts: readonly WorkspaceAgentListeningPort[];
 	portSharingControlsEnabled: boolean;
+	refetchSharedPorts: () => void;
 }
 
-type Optional<T, K extends keyof T> = Pick<Partial<T>, K> & Omit<T, K>;
-
 export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 	host,
-	workspaceName,
-	workspaceID,
+	workspace,
 	agent,
 	template,
-	username,
+	sharedPorts,
 	listeningPorts,
 	portSharingControlsEnabled,
+	refetchSharedPorts,
 }) => {
 	const theme = useTheme();
 	const [listeningPortProtocol, setListeningPortProtocol] = useState(
-		getWorkspaceListeningPortsProtocol(workspaceID),
+		getWorkspaceListeningPortsProtocol(workspace.id),
 	);
 
-	const sharedPortsQuery = useQuery({
-		...workspacePortShares(workspaceID),
-		enabled: agent.status === "connected",
+	const upsertSharedPortMutation = useMutation({
+		...upsertWorkspacePortShare(workspace.id),
+		onSuccess: refetchSharedPorts,
 	});
-	const sharedPorts = sharedPortsQuery.data?.shares || [];
 
-	const upsertSharedPortMutation = useMutation(
-		upsertWorkspacePortShare(workspaceID),
-	);
-
-	const deleteSharedPortMutation = useMutation(
-		deleteWorkspacePortShare(workspaceID),
-	);
+	const deleteSharedPortMutation = useMutation({
+		...deleteWorkspacePortShare(workspace.id),
+		onSuccess: refetchSharedPorts,
+	});
 
-	// share port form
 	const {
 		mutateAsync: upsertWorkspacePortShareForm,
 		isPending: isSubmitting,
 		error: submitError,
-	} = useMutation(upsertWorkspacePortShare(workspaceID));
-	const validationSchema = getValidationSchema();
-	// TODO: do partial here
-	const form: FormikContextType<
-		Optional<UpsertWorkspaceAgentPortShareRequest, "port">
-	> = useFormik<Optional<UpsertWorkspaceAgentPortShareRequest, "port">>({
+	} = useMutation({
+		...upsertWorkspacePortShare(workspace.id),
+		onSuccess: refetchSharedPorts,
+	});
+
+	const form = useFormik({
 		initialValues: {
 			agent_name: agent.name,
-			port: undefined,
+			port: "",
 			protocol: "http",
 			share_level: "authenticated",
 		},
-		validationSchema,
-		onSubmit: async (values) => {
-			// we need port to be optional in the initialValues so it appears empty instead of 0.
-			// because of this we need to reset the form to clear the port field manually.
-			form.resetForm();
-			await form.setFieldValue("port", "");
-
-			const port = Number(values.port);
+		validationSchema: openPortSchema(),
+		onSubmit: async (values, { resetForm }) => {
+			resetForm();
 			await upsertWorkspacePortShareForm({
-				...values,
-				port,
+				agent_name: values.agent_name,
+				port: Number(values.port),
+				share_level: values.share_level as WorkspaceAgentPortShareLevel,
+				protocol: values.protocol as WorkspaceAgentPortShareProtocol,
 			});
-			await sharedPortsQuery.refetch();
 		},
 	});
 	const getFieldHelpers = getFormHelpers(form, submitError);
@@ -188,7 +200,7 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 		(port) => port.agent_name === agent.name,
 	);
 	// we don't want to show listening ports if it's a shared port
-	const filteredListeningPorts = (listeningPorts ?? []).filter((port) =>
+	const filteredListeningPorts = listeningPorts.filter((port) =>
 		filteredSharedPorts.every((sharedPort) => sharedPort.port !== port.port),
 	);
 	// only disable the form if shared port controls are entitled and the template doesn't allow sharing ports
@@ -257,7 +269,7 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 											| "https";
 										setListeningPortProtocol(selectedProtocol);
 										saveWorkspaceListeningPortsProtocol(
-											workspaceID,
+											workspace.id,
 											selectedProtocol,
 										);
 									}}
@@ -276,8 +288,8 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 										host,
 										port,
 										agent.name,
-										workspaceName,
-										username,
+										workspace.name,
+										workspace.owner_name,
 										listeningPortProtocol,
 									);
 									window.open(url, "_blank");
@@ -317,8 +329,8 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 							host,
 							port.port,
 							agent.name,
-							workspaceName,
-							username,
+							workspace.name,
+							workspace.owner_name,
 							listeningPortProtocol,
 						);
 						const label =
@@ -371,7 +383,6 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 																protocol: listeningPortProtocol,
 																share_level: "authenticated",
 															});
-															await sharedPortsQuery.refetch();
 														}}
 													>
 														<ShareIcon />
@@ -407,8 +418,8 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 								host,
 								share.port,
 								agent.name,
-								workspaceName,
-								username,
+								workspace.name,
+								workspace.owner_name,
 								share.protocol,
 							);
 							const label = share.port;
@@ -445,7 +456,6 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 														.value as WorkspaceAgentPortShareProtocol,
 													share_level: share.share_level,
 												});
-												await sharedPortsQuery.refetch();
 											}}
 										>
 											<MenuItem value="http">HTTP</MenuItem>
@@ -469,7 +479,6 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 														share_level: event.target
 															.value as WorkspaceAgentPortShareLevel,
 													});
-													await sharedPortsQuery.refetch();
 												}}
 											>
 												<MenuItem value="authenticated">Authenticated</MenuItem>
@@ -488,7 +497,6 @@ export const PortForwardPopoverView: FC<PortForwardPopoverViewProps> = ({
 													agent_name: agent.name,
 													port: share.port,
 												});
-												await sharedPortsQuery.refetch();
 											}}
 										>
 											<XIcon