feat: connect to dynamic parameters websocket

jaaydenh · jaaydenh · commit 2da7d9948cda · 2025-04-16T09:13:03.000Z
diff --git a/site/src/api/api.ts b/site/src/api/api.ts
@@ -1009,6 +1009,13 @@ class ApiMethods {
 		return response.data;
 	};
 
+	templateVersionDynamicParameters = (versionId: string): WebSocket => {
+		const socket = createWebSocket(
+			`/api/v2/templateversions/${versionId}/dynamic-parameters`,
+		);
+		return socket;
+	};
+
 	/**
 	 * @param organization Can be the organization's ID or name
 	 */
diff --git a/site/src/pages/CreateWorkspacePage/CreateWorkspacePageExperimental.tsx b/site/src/pages/CreateWorkspacePage/CreateWorkspacePageExperimental.tsx
@@ -32,6 +32,7 @@ import type { AutofillBuildParameter } from "utils/richParameters";
 import { CreateWorkspacePageViewExperimental } from "./CreateWorkspacePageViewExperimental";
 export const createWorkspaceModes = ["form", "auto", "duplicate"] as const;
 export type CreateWorkspaceMode = (typeof createWorkspaceModes)[number];
+import { API } from "api/api";
 import {
 	type CreateWorkspacePermissions,
 	createWorkspaceChecks,
@@ -47,8 +48,8 @@ const CreateWorkspacePageExperimental: FC = () => {
 
 	const [currentResponse, setCurrentResponse] =
 		useState<DynamicParametersResponse | null>(null);
-	const [wsResponseId, setWSResponseId] = useState<number>(0);
-	const sendMessage = (message: DynamicParametersRequest) => {};
+	const [wsResponseId, setWSResponseId] = useState<number>(-1);
+	const webSocket = useRef<WebSocket | null>(null);
 
 	const customVersionId = searchParams.get("version") ?? undefined;
 	const defaultName = searchParams.get("name");
@@ -80,6 +81,59 @@ const CreateWorkspacePageExperimental: FC = () => {
 	const realizedVersionId =
 		customVersionId ?? templateQuery.data?.active_version_id;
 
+	// Initialize the WebSocket connection when there is a valid template version ID
+	useEffect(() => {
+		if (!realizedVersionId) {
+			return;
+		}
+
+		if (webSocket.current) {
+			webSocket.current.close();
+		}
+
+		const socket = API.templateVersionDynamicParameters(realizedVersionId);
+
+		socket.addEventListener("message", (event) => {
+			try {
+				const response = JSON.parse(event.data) as DynamicParametersResponse;
+
+				if (response && response.id >= wsResponseId) {
+					setCurrentResponse((prev) => {
+						if (prev?.id === response.id) {
+							return prev;
+						}
+						return response;
+					});
+				}
+			} catch (error) {
+				console.error("Failed to parse WebSocket message:", error);
+			}
+		});
+
+		webSocket.current = socket;
+
+		return () => {
+			if (webSocket.current) {
+				webSocket.current.close();
+			}
+		};
+	}, [realizedVersionId]);
+
+	const sendMessage =
+		(formValues: Record<string, string>) => {
+			setWSResponseId(prevId => {
+				const request: DynamicParametersRequest = {
+					id: prevId + 1,
+					inputs: formValues,
+				};
+				if (webSocket.current && webSocket.current.readyState === WebSocket.OPEN) {
+					webSocket.current.send(JSON.stringify(request));
+					return prevId + 1;
+				}
+				return prevId;
+			})
+		};
+
 	const organizationId = templateQuery.data?.organization_id;
 
 	const {
@@ -210,7 +264,6 @@ const CreateWorkspacePageExperimental: FC = () => {
 					parameters={sortedParams}
 					presets={templateVersionPresetsQuery.data ?? []}
 					creatingWorkspace={createWorkspaceMutation.isLoading}
-					setWSResponseId={setWSResponseId}
 					sendMessage={sendMessage}
 					onCancel={() => {
 						navigate(-1);
diff --git a/site/src/pages/CreateWorkspacePage/CreateWorkspacePageViewExperimental.tsx b/site/src/pages/CreateWorkspacePage/CreateWorkspacePageViewExperimental.tsx
@@ -67,8 +67,7 @@ export interface CreateWorkspacePageViewExperimentalProps {
 		owner: TypesGen.User,
 	) => void;
 	resetMutation: () => void;
-	sendMessage: (message: DynamicParametersRequest) => void;
-	setWSResponseId: (value: React.SetStateAction<number>) => void;
+	sendMessage: (message: Record<string, string>) => void;
 	startPollingExternalAuth: () => void;
 }
 
@@ -95,7 +94,6 @@ export const CreateWorkspacePageViewExperimental: FC<
 	onCancel,
 	resetMutation,
 	sendMessage,
-	setWSResponseId,
 	startPollingExternalAuth,
 }) => {
 	const [owner, setOwner] = useState(defaultOwner);
