feat: connect to dynamic parameters websocket

jaaydenh · jaaydenh · commit 9673441d6c21 · 2025-04-14T19:27:48.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
@@ -7,6 +7,7 @@ import {
 } from "api/queries/templates";
 import { autoCreateWorkspace, createWorkspace } from "api/queries/workspaces";
 import type {
+	DynamicParametersRequest,
 	DynamicParametersResponse,
 	Template,
 	Workspace,
@@ -31,17 +32,13 @@ 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 { useWebSocket } from "hooks/useWebsocket";
+import { API } from "api/api";
 import {
 	type CreateWorkspacePermissions,
 	createWorkspaceChecks,
 } from "./permissions";
 export type ExternalAuthPollingState = "idle" | "polling" | "abandoned";
 
-const serverAddress = "localhost:8100";
-const urlTestdata = "demo";
-const wsUrl = `ws://${serverAddress}/ws/${encodeURIComponent(urlTestdata)}`;
-
 const CreateWorkspacePageExperimental: FC = () => {
 	const { organization: organizationName = "default", template: templateName } =
 		useParams() as { organization?: string; template: string };
@@ -51,20 +48,8 @@ const CreateWorkspacePageExperimental: FC = () => {
 
 	const [currentResponse, setCurrentResponse] =
 		useState<DynamicParametersResponse | null>(null);
-	const [wsResponseId, setWSResponseId] = useState<number>(0);
-	const { message: webSocketResponse, sendMessage } =
-		useWebSocket<DynamicParametersResponse>(wsUrl, urlTestdata, "", "");
-
-	useEffect(() => {
-		if (webSocketResponse && webSocketResponse.id >= wsResponseId) {
-			setCurrentResponse((prev) => {
-				if (prev?.id === webSocketResponse.id) {
-					return prev;
-				}
-				return webSocketResponse;
-			});
-		}
-	}, [webSocketResponse, wsResponseId]);
+	const [wsResponseId, setWSResponseId] = useState<number>(-1);
+	const webSocket = useRef<WebSocket | null>(null);
 
 	const customVersionId = searchParams.get("version") ?? undefined;
 	const defaultName = searchParams.get("name");
@@ -96,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 {
@@ -225,7 +263,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
@@ -1,6 +1,5 @@
 import type * as TypesGen from "api/typesGenerated";
 import type {
-	DynamicParametersRequest,
 	PreviewDiagnostics,
 	PreviewParameter,
 } from "api/typesGenerated";
@@ -66,8 +65,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;
 }
 
@@ -94,7 +92,6 @@ export const CreateWorkspacePageViewExperimental: FC<
 	onCancel,
 	resetMutation,
 	sendMessage,
-	setWSResponseId,
 	startPollingExternalAuth,
 }) => {
 	const [owner, setOwner] = useState(defaultOwner);
@@ -218,35 +215,24 @@ export const CreateWorkspacePageViewExperimental: FC<
 		parameterField: string,
 		parameter: PreviewParameter,
 	) => {
-		// Update form value immediately for all types
 		await form.setFieldValue(parameterField, {
-			name: parameter.form_type,
+			name: parameter.name,
 			value,
 		});
 
-		// Create the request object
 		const createRequest = () => {
 			// Convert the rich_parameter_values array to a key-value object
-			const newInputs = (form.values.rich_parameter_values ?? []).reduce(
+			const formValues = (form.values.rich_parameter_values ?? []).reduce(
 				(acc, param) => {
 					acc[param.name] = param.value;
 					return acc;
 				},
 				{} as Record<string, string>,
 			);
 
-			// Update the input for the changed parameter
-			newInputs[parameter.name] = value;
-
-			setWSResponseId((prevId) => {
-				const newId = prevId + 1;
-				const request: DynamicParametersRequest = {
-					id: newId,
-					inputs: newInputs,
-				};
-				sendMessage(request);
-				return newId;
-			});
+			formValues[parameter.name] = value;
+
+			sendMessage(formValues);
 		};
 
 		// Clear any existing timer
