coder · Parkreiner · Jul 3, 2024 · Jul 1, 2024 · Jul 1, 2024 · Jul 1, 2024
diff --git a/site/src/api/errors.ts b/site/src/api/errors.ts
@@ -110,15 +110,15 @@ export const getValidationErrorMessage = (error: unknown): string => {
   return validationErrors.map((error) => error.detail).join("\n");
 };
 
-export const getErrorDetail = (error: unknown): string | undefined | null => {
+export const getErrorDetail = (error: unknown): string | null => {
   if (error instanceof Error) {
     return "Please check the developer console for more details.";
   }
   if (isApiError(error)) {
-    return error.response.data.detail;
+    return error.response.data.detail ?? null;
   }
   if (isApiErrorResponse(error)) {
-    return error.detail;
+    return error.detail ?? null;
   }
   return null;
 };
diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.ts b/site/src/modules/resources/AgentLogs/useAgentLogs.ts
@@ -15,22 +15,38 @@ export type UseAgentLogsOptions = Readonly<{
   enabled?: boolean;
 }>;
 
+/**
+ * Defines a custom hook that gives you all workspace agent logs for a given
+ * workspace.
+ *
+ * Depending on the status of the workspace, all logs may or may not be
+ * available.
+ */
 export function useAgentLogs(
   options: UseAgentLogsOptions,
 ): readonly WorkspaceAgentLog[] | undefined {
   const { workspaceId, agentId, agentLifeCycleState, enabled = true } = options;
+
   const queryClient = useQueryClient();
   const queryOptions = agentLogs(workspaceId, agentId);
-  const query = useQuery({
-    ...queryOptions,
-    enabled,
-  });
-  const logs = query.data;
+  const query = useQuery({ ...queryOptions, enabled });
+
+  // One pitfall with the current approach: the enabled property does NOT
+  // prevent the useQuery call above from eventually having data. All it does
+  // is prevent it from getting data on its own. Let's say a different useQuery
+  // call elsewhere in the app has the same query key and is enabled. When it
+  // gets data back from the server, the useQuery call here will re-render with
+  // that same new data, even though this state is "disabled". This can EASILY
+  // cause bugs.
+  const logs = enabled ? query.data : undefined;
 
   const lastQueriedLogId = useRef(0);
   useEffect(() => {
-    if (logs && lastQueriedLogId.current === 0) {
-      lastQueriedLogId.current = logs[logs.length - 1].id;
+    const lastLog = logs?.at(-1);
+    const canSetLogId = lastLog !== undefined && lastQueriedLogId.current === 0;
+
+    if (canSetLogId) {
+      lastQueriedLogId.current = lastLog.id;
     }
   }, [logs]);
 

diff --git a/site/src/pages/WorkspacePage/WorkspaceActions/DownloadLogsDialog.tsx b/site/src/pages/WorkspacePage/WorkspaceActions/DownloadLogsDialog.tsx
@@ -1,17 +1,25 @@
 import { useTheme, type Interpolation, type Theme } from "@emotion/react";
 import Skeleton from "@mui/material/Skeleton";
+import ErrorIcon from "@mui/icons-material/ErrorOutline";
 import { saveAs } from "file-saver";
 import JSZip from "jszip";
-import { useMemo, useState, type FC } from "react";
-import { useQueries, useQuery } from "react-query";
+import { type FC, useMemo, useState, useRef, useEffect } from "react";
+import { UseQueryOptions, useQueries, useQuery } from "react-query";
 import { agentLogs, buildLogs } from "api/queries/workspaces";
-import type { Workspace, WorkspaceAgent } from "api/typesGenerated";
+import type {
+  Workspace,
+  WorkspaceAgent,
+  WorkspaceAgentLog,
+} from "api/typesGenerated";
 import {
   ConfirmDialog,
   type ConfirmDialogProps,
 } from "components/Dialogs/ConfirmDialog/ConfirmDialog";
 import { displayError } from "components/GlobalSnackbar/utils";
 import { Stack } from "components/Stack/Stack";
+import { ErrorAlert } from "components/Alert/ErrorAlert";
+
+const BLOB_SIZE_UNITS = ["B", "KB", "MB", "GB", "TB"] as const;
 
 type DownloadLogsDialogProps = Pick<
   ConfirmDialogProps,
@@ -28,70 +36,124 @@
 
 export const DownloadLogsDialog: FC<DownloadLogsDialogProps> = ({
   workspace,
+  open,
+  onClose,
   download = saveAs,
-  ...dialogProps
 }) => {
   const theme = useTheme();
-  const agents = selectAgents(workspace);
-  const agentLogResults = useQueries({
-    queries: agents.map((a) => ({
-      ...agentLogs(workspace.id, a.id),
-      enabled: dialogProps.open,
-    })),
-  });
+
   const buildLogsQuery = useQuery({
     ...buildLogs(workspace),
-    enabled: dialogProps.open,
+    enabled: open,
   });
-  const downloadableFiles: DownloadableFile[] = useMemo(() => {
-    const files: DownloadableFile[] = [
-      {
-        name: `${workspace.name}-build-logs.txt`,
-        blob: buildLogsQuery.data
-          ? new Blob([buildLogsQuery.data.map((l) => l.output).join("\n")], {
-              type: "text/plain",
-            })
-          : undefined,
-      },
+  const buildLogsFile = useMemo<DownloadableFile>(() => {
+    return {
+      name: `${workspace.name}-build-logs.txt`,
+      blob: buildLogsQuery.data
+        ? new Blob([buildLogsQuery.data.map((l) => l.output).join("\n")], {
+            type: "text/plain",
+          })
+        : undefined,
+    };
+  }, [workspace.name, buildLogsQuery.data]);
+
+  // This is clunky, but we have to memoize in two steps to make sure that we
+  // don't accidentally break the memo cache every render. We can't tuck
+  // everything into a single memo call, because we need to set up React Query
+  // state between processing the agents, but we can't violate rules of hooks by
+  // putting hooks inside of hooks
+  type AgentInfo = Readonly<{
+    agents: readonly WorkspaceAgent[];
+    logOptionsArray: readonly UseQueryOptions<readonly WorkspaceAgentLog[]>[];
+  }>;
+
+  const { agents, logOptionsArray } = useMemo<AgentInfo>(() => {
+    const allAgents = workspace.latest_build.resources.flatMap(
+      (resource) => resource.agents ?? [],
+    );
+
+    // Can't use the "new Set()" trick because we're not dealing with primitives
+    const uniqueAgents = [
+      ...new Map(allAgents.map((agent) => [agent.id, agent])).values(),
     ];
 
+    return {
+      agents: uniqueAgents,
+      logOptionsArray: uniqueAgents.map((agent) => {
+        return {
+          ...agentLogs(workspace.id, agent.id),
+          enabled: open,
+        };
+      }),
+    };
+  }, [workspace, open]);
+
+  const agentLogQueries = useQueries({ queries: logOptionsArray });
+  const allFiles = useMemo<readonly DownloadableFile[]>(() => {
+    const files: DownloadableFile[] = [buildLogsFile];
+
     agents.forEach((a, i) => {
       const name = `${a.name}-logs.txt`;
-      const logs = agentLogResults[i].data;
-      const txt = logs?.map((l) => l.output).join("\n");
+      const txt = agentLogQueries[i]?.data?.map((l) => l.output).join("\n");
+
       let blob: Blob | undefined;
       if (txt) {
         blob = new Blob([txt], { type: "text/plain" });
       }
+
       files.push({ name, blob });
     });
 
     return files;
-  }, [agentLogResults, agents, buildLogsQuery.data, workspace.name]);
-  const isLoadingFiles = downloadableFiles.some((f) => f.blob === undefined);
+  }, [agentLogQueries, agents, buildLogsFile]);
+
   const [isDownloading, setIsDownloading] = useState(false);
+  const isWorkspaceHealthy = workspace.health.healthy;
+  const isLoadingFiles = allFiles.some((f) => f.blob === undefined);
+
+  const resetDownloadStateIdRef = useRef<number | undefined>(undefined);
+  useEffect(() => {
+    const clearTimeoutOnUnmount = () => {
+      window.clearTimeout(resetDownloadStateIdRef.current);
+    };
+
+    return clearTimeoutOnUnmount;
+  }, []);
 
   return (
     <ConfirmDialog
-      {...dialogProps}
+      open={open}
+      onClose={onClose}
       hideCancel={false}
       title="Download logs"
-      confirmText="Download"
-      disabled={isLoadingFiles}
       confirmLoading={isDownloading}
+      confirmText={
+        <>
+          Download
+          {!isWorkspaceHealthy && <> {isLoadingFiles ? "partial" : "all"}</>}
+        </>
+      }
+      disabled={
+        isDownloading ||
+        // If a workspace isn't healthy, let the user download as many logs as
+        // they can. Otherwise, wait for everything to come in
+        (isWorkspaceHealthy && isLoadingFiles)
+      }
       onConfirm={async () => {
+        setIsDownloading(true);
+        const zip = new JSZip();
+        allFiles.forEach((f) => {
+          if (f.blob) {
+            zip.file(f.name, f.blob);
+          }
+        });
+
         try {
-          setIsDownloading(true);
-          const zip = new JSZip();
-          downloadableFiles.forEach((f) => {
-            if (f.blob) {
-              zip.file(f.name, f.blob);
-            }
-          });
           const content = await zip.generateAsync({ type: "blob" });
           download(content, `${workspace.name}-logs.zip`);
-          dialogProps.onClose();
-          setTimeout(() => {
+          onClose();
+
+          resetDownloadStateIdRef.current = window.setTimeout(() => {
             setIsDownloading(false);
           }, theme.transitions.duration.leavingScreen);
         } catch (error) {
@@ -106,18 +168,18 @@
             Downloading logs will create a zip file containing all logs from all
             jobs in this workspace. This may take a while.
           </p>
+
+          {!isWorkspaceHealthy && isLoadingFiles && (
+            <ErrorAlert error="Your workspace is not healthy. Some logs may be unavailable." />
+          )}
+
           <ul css={styles.list}>
-            {downloadableFiles.map((f) => (
-              <li key={f.name} css={styles.listItem}>
-                <span css={styles.listItemPrimary}>{f.name}</span>
-                <span css={styles.listItemSecondary}>
-                  {f.blob ? (
-                    humanBlobSize(f.blob.size)
-                  ) : (
-                    <Skeleton variant="text" width={48} height={12} />
-                  )}
-                </span>
-              </li>
+            {allFiles.map((f) => (
+              <DownloadingItem
+                key={f.name}
+                file={f}
+                giveUpTimeMs={isWorkspaceHealthy ? undefined : 5_000}
+              />
             ))}
           </ul>
         </Stack>
@@ -126,20 +188,70 @@
   );
 };
 
+type DownloadingItemProps = Readonly<{
+  // A value of undefined indicates that the component will wait forever
+  giveUpTimeMs?: number;
+  file: DownloadableFile;
+}>;
+
+const DownloadingItem: FC<DownloadingItemProps> = ({ file, giveUpTimeMs }) => {
+  const theme = useTheme();
+  const [isWaiting, setIsWaiting] = useState(true);
+  useEffect(() => {
+    if (giveUpTimeMs === undefined || file.blob !== undefined) {
+      setIsWaiting(true);
+      return;
+    }
+
+    const timeoutId = window.setTimeout(
+      () => setIsWaiting(false),
+      giveUpTimeMs,
+    );
+
+    return () => window.clearTimeout(timeoutId);
+  }, [giveUpTimeMs, file]);
+
+  return (
+    <li css={styles.listItem}>
+      <span
+        css={[
+          styles.listItemPrimary,
+          !isWaiting && { color: theme.palette.text.disabled },
+        ]}
+      >
+        {file.name}
+      </span>
+
+      <span css={styles.listItemSecondary}>
+        {file.blob ? (
+          humanBlobSize(file.blob.size)
+        ) : isWaiting ? (
+          <Skeleton variant="text" width={48} height={12} />
+        ) : (
+          <div css={styles.notAvailableText}>
+            <span aria-hidden>
+              <ErrorIcon fontSize={"inherit"} />
+            </span>
+
+            <p>N/A</p>
+          </div>
+        )}
+      </span>
+    </li>
+  );
+};
+
 function humanBlobSize(size: number) {
-  const units = ["B", "KB", "MB", "GB", "TB"];
   let i = 0;
-  while (size > 1024 && i < units.length) {
+  while (size > 1024 && i < BLOB_SIZE_UNITS.length) {
     size /= 1024;
     i++;
   }
-  return `${size.toFixed(2)} ${units[i]}`;
-}
 
-function selectAgents(workspace: Workspace): WorkspaceAgent[] {
-  return workspace.latest_build.resources
-    .flatMap((r) => r.agents)
-    .filter((a) => a !== undefined) as WorkspaceAgent[];
+  // The while condition can break if we accidentally exceed the bounds of the
+  // array. Have to be extra sure we have a unit at the very end.
+  const finalUnit = BLOB_SIZE_UNITS[i] ?? BLOB_SIZE_UNITS.at(-1) ?? "TB";
+  return `${size.toFixed(2)} ${finalUnit}`;
 }
 
 const styles = {
@@ -163,4 +275,22 @@
   listItemSecondary: {
     fontSize: 14,
   },
+
+  notAvailableText: (theme) => ({
+    display: "flex",
+    flexFlow: "row nowrap",
+    alignItems: "center",
+    columnGap: "4px",
+
+    "& > span": {
+      maxHeight: "fit-content",
+      display: "flex",
+      alignItems: "center",
+      color: theme.palette.error.light,
+    },
+
+    "& > p": {
+      opacity: "80%",
+    },
+  }),
 } satisfies Record<string, Interpolation<Theme>>;