chore(site): replace agent log service

coder · BrunoQuaresma · Sep 22, 2023 · Sep 21, 2023 · Sep 21, 2023 · Sep 22, 2023
commit a39b51e17d6f13f94db6af7ea2c3cb893ed246fc
diff --git a/site/src/api/api.ts b/site/src/api/api.ts
@@ -1392,7 +1392,7 @@ export const watchBuildLogsByTemplateVersionId = (
 type WatchWorkspaceAgentLogsOptions = {
   after: number;
   onMessage: (logs: TypesGen.WorkspaceAgentLog[]) => void;
-  onDone: () => void;
+  onDone?: () => void;
   onError: (error: Error) => void;
 };
 
@@ -1423,7 +1423,7 @@ export const watchWorkspaceAgentLogs = (
     onError(new Error("socket errored"));
   });
   socket.addEventListener("close", () => {
-    onDone();
+    onDone && onDone();
   });
 
   return socket;

diff --git a/site/src/components/Resources/AgentRow.tsx b/site/src/components/Resources/AgentRow.tsx
@@ -1,13 +1,17 @@
 import Popover from "@mui/material/Popover";
 import { makeStyles, useTheme } from "@mui/styles";
 import Skeleton from "@mui/material/Skeleton";
-import { useMachine } from "@xstate/react";
+import * as API from "api/api";
 import CodeOutlined from "@mui/icons-material/CodeOutlined";
 import {
   CloseDropdown,
   OpenDropdown,
 } from "components/DropdownArrows/DropdownArrows";
-import { LogLine, logLineHeight } from "components/WorkspaceBuildLogs/Logs";
+import {
+  Line,
+  LogLine,
+  logLineHeight,
+} from "components/WorkspaceBuildLogs/Logs";
 import { PortForwardButton } from "./PortForwardButton";
 import { VSCodeDesktopButton } from "components/Resources/VSCodeDesktopButton/VSCodeDesktopButton";
 import {
@@ -25,15 +29,11 @@ import AutoSizer from "react-virtualized-auto-sizer";
 import { FixedSizeList as List, ListOnScrollProps } from "react-window";
 import { colors } from "theme/colors";
 import { combineClasses } from "utils/combineClasses";
-import {
-  LineWithID,
-  workspaceAgentLogsMachine,
-} from "xServices/workspaceAgentLogs/workspaceAgentLogsXService";
 import {
   Workspace,
   WorkspaceAgent,
   WorkspaceAgentMetadata,
-} from "../../api/typesGenerated";
+} from "api/typesGenerated";
 import { AppLink } from "./AppLink/AppLink";
 import { SSHButton } from "./SSHButton/SSHButton";
 import { Stack } from "../Stack/Stack";
@@ -44,6 +44,14 @@ import { AgentVersion } from "./AgentVersion";
 import { AgentStatus } from "./AgentStatus";
 import Collapse from "@mui/material/Collapse";
 import { useProxy } from "contexts/ProxyContext";
+import { displayError } from "components/GlobalSnackbar/utils";
+
+// Logs are stored as the Line interface to make rendering
+// much more efficient. Instead of mapping objects each time, we're
+// able to just pass the array of logs to the component.
+export interface LineWithID extends Line {
+  id: number;
+}
 
 export interface AgentRowProps {
   agent: WorkspaceAgent;
@@ -68,24 +76,10 @@ export const AgentRow: FC<AgentRowProps> = ({
   hideVSCodeDesktopButton,
   serverVersion,
   onUpdateAgent,
-  storybookLogs,
   storybookAgentMetadata,
   sshPrefix,
 }) => {
   const styles = useStyles();
-  const [logsMachine, sendLogsEvent] = useMachine(workspaceAgentLogsMachine, {
-    context: { agentID: agent.id },
-    services: process.env.STORYBOOK
-      ? {
-          getLogs: async () => {
-            return storybookLogs || [];
-          },
-          streamLogs: () => async () => {
-            // noop
-          },
-        }
-      : undefined,
-  });
   const theme = useTheme();
   const startupScriptAnchorRef = useRef<HTMLButtonElement>(null);
   const [startupScriptOpen, setStartupScriptOpen] = useState(false);
@@ -94,36 +88,17 @@ export const AgentRow: FC<AgentRowProps> = ({
     showApps &&
     ((agent.status === "connected" && hasAppsToDisplay) ||
       agent.status === "connecting");
-  const hasStartupFeatures =
-    Boolean(agent.logs_length) || Boolean(logsMachine.context.logs?.length);
+  const hasStartupFeatures = Boolean(agent.logs_length);
   const { proxy } = useProxy();
-
   const [showLogs, setShowLogs] = useState(
     ["starting", "start_timeout"].includes(agent.lifecycle_state) &&
       hasStartupFeatures,
   );
-  useEffect(() => {
-    setShowLogs(agent.lifecycle_state !== "ready" && hasStartupFeatures);
-  }, [agent.lifecycle_state, hasStartupFeatures]);
-  // External applications can provide startup logs for an agent during it's spawn.
-  // These could be Kubernetes logs, or other logs that are useful to the user.
-  // For this reason, we want to fetch these logs when the agent is starting.
-  useEffect(() => {
-    if (agent.lifecycle_state === "starting") {
-      sendLogsEvent("FETCH_LOGS");
-    }
-  }, [sendLogsEvent, agent.lifecycle_state]);
-  useEffect(() => {
-    // We only want to fetch logs when they are actually shown,
-    // otherwise we can make a lot of requests that aren't necessary.
-    if (showLogs && logsMachine.can("FETCH_LOGS")) {
-      sendLogsEvent("FETCH_LOGS");
-    }
-  }, [logsMachine, sendLogsEvent, showLogs]);
+  const agentLogs = useAgentLogs(agent.id, { enabled: showLogs });
   const logListRef = useRef<List>(null);
   const logListDivRef = useRef<HTMLDivElement>(null);
   const startupLogs = useMemo(() => {
-    const allLogs = logsMachine.context.logs || [];
+    const allLogs = agentLogs || [];
 
     const logs = [...allLogs];
     if (agent.logs_overflowed) {
@@ -135,8 +110,13 @@ export const AgentRow: FC<AgentRowProps> = ({
       });
     }
     return logs;
-  }, [logsMachine.context.logs, agent.logs_overflowed]);
+  }, [agentLogs, agent.logs_overflowed]);
   const [bottomOfLogs, setBottomOfLogs] = useState(true);
+
+  useEffect(() => {
+    setShowLogs(agent.lifecycle_state !== "ready" && hasStartupFeatures);
+  }, [agent.lifecycle_state, hasStartupFeatures]);
+
   // This is a layout effect to remove flicker when we're scrolling to the bottom.
   useLayoutEffect(() => {
     // If we're currently watching the bottom, we always want to stay at the bottom.
@@ -396,6 +376,48 @@ export const AgentRow: FC<AgentRowProps> = ({
   );
 };
 
+const useAgentLogs = (agentId: string, { enabled }: { enabled: boolean }) => {
+  const [logs, setLogs] = useState<LineWithID[]>();
+  const socket = useRef<WebSocket | null>(null);
+
+  useEffect(() => {
+    if (!enabled) {
+      socket.current?.close();
+      return;
+    }
+
+    socket.current = API.watchWorkspaceAgentLogs(agentId, {
+      // Get all logs
+      after: 0,
+      onMessage: (logs) => {
+        setLogs((previousLogs) => {
+          const newLogs: LineWithID[] = logs.map((log) => ({
+            id: log.id,
+            level: log.level || "info",
+            output: log.output,
+            time: log.created_at,
+          }));
+
+          if (!previousLogs) {
+            return newLogs;
+          }
+
+          return [...previousLogs, ...newLogs];
+        });
+      },
+      onError: () => {
+        displayError("Error on gettings agent logs");
+      },
+    });
+
+    return () => {
+      socket.current?.close();
+    };
+  }, [agentId, enabled]);
+
+  return logs;
+};
+
 const useStyles = makeStyles((theme) => ({
   agentRow: {
     backgroundColor: theme.palette.background.paperLight,

diff --git a/site/src/components/WorkspaceBuildLogs/Logs.tsx b/site/src/components/WorkspaceBuildLogs/Logs.tsx
@@ -60,12 +60,25 @@ export const LogLine: FC<{
   const output = useMemo(() => {
     return convert.toHtml(line.output.split(/\r/g).pop() as string);
   }, [line.output]);
+  const isUsingLineNumber = number !== undefined;
 
   return (
-    <div className={combineClasses([styles.line, line.level])} style={style}>
+    <div
+      className={combineClasses([
+        styles.line,
+        line.level,
+        isUsingLineNumber && styles.lineNumber,
+      ])}
+      style={style}
+    >
       {!hideTimestamp && (
         <>
-          <span className={styles.time}>
+          <span
+            className={combineClasses([
+              styles.time,
+              isUsingLineNumber && styles.number,
+            ])}
+          >
             {number ? number : dayjs(line.time).format(`HH:mm:ss.SSS`)}
           </span>
           <span className={styles.space} />
@@ -119,6 +132,9 @@ const useStyles = makeStyles((theme) => ({
       backgroundColor: theme.palette.warning.dark,
     },
   },
+  lineNumber: {
+    paddingLeft: theme.spacing(2),
+  },
   space: {
     userSelect: "none",
     width: theme.spacing(3),
@@ -132,4 +148,8 @@ const useStyles = makeStyles((theme) => ({
     display: "inline-block",
     color: theme.palette.text.secondary,
   },
+  number: {
+    width: theme.spacing(4),
+    textAlign: "right",
+  },
 }));
diff --git a/site/src/pages/WorkspacePage/WorkspacePage.test.tsx b/site/src/pages/WorkspacePage/WorkspacePage.test.tsx
@@ -43,7 +43,7 @@ const renderWorkspacePage = async () => {
   jest
     .spyOn(api, "watchWorkspaceAgentLogs")
     .mockImplementation((_, options) => {
-      options.onDone();
+      options.onDone && options.onDone();
       return new WebSocket("");
     });
   renderWithAuth(<WorkspacePage />, {

diff --git a/site/src/xServices/workspaceAgentLogs/workspaceAgentLogsXService.ts b/site/src/xServices/workspaceAgentLogs/workspaceAgentLogsXService.ts