fix(site): fix agent logs streaming for third party apps (#14541)

BrunoQuaresma · aslilac · commit d6d568b7aea8 · 2024-09-04T20:29:31.000Z
(cherry picked from commit 242b1ea)
diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx b/site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx
@@ -28,7 +28,7 @@ describe("useAgentLogs", () => {
     expect(wsSpy).not.toHaveBeenCalled();
   });
 
-  it("should return existing logs without network calls", async () => {
+  it("should return existing logs without network calls if state is off", async () => {
     const queryClient = createTestQueryClient();
     queryClient.setQueryData(
       agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
@@ -39,7 +39,7 @@ describe("useAgentLogs", () => {
     const { result } = renderUseAgentLogs(queryClient, {
       workspaceId: MockWorkspace.id,
       agentId: MockWorkspaceAgent.id,
-      agentLifeCycleState: "ready",
+      agentLifeCycleState: "off",
     });
     await waitFor(() => {
       expect(result.current).toHaveLength(5);
@@ -48,12 +48,12 @@ describe("useAgentLogs", () => {
     expect(wsSpy).not.toHaveBeenCalled();
   });
 
-  it("should fetch logs when empty and should not connect to WebSocket when not starting", async () => {
+  it("should fetch logs when empty", async () => {
     const queryClient = createTestQueryClient();
     const fetchSpy = jest
       .spyOn(API, "getWorkspaceAgentLogs")
       .mockResolvedValueOnce(generateLogs(5));
-    const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
+    jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
     const { result } = renderUseAgentLogs(queryClient, {
       workspaceId: MockWorkspace.id,
       agentId: MockWorkspaceAgent.id,
@@ -63,10 +63,9 @@ describe("useAgentLogs", () => {
       expect(result.current).toHaveLength(5);
     });
     expect(fetchSpy).toHaveBeenCalledWith(MockWorkspaceAgent.id);
-    expect(wsSpy).not.toHaveBeenCalled();
   });
 
-  it("should fetch logs and connect to websocket when agent is starting", async () => {
+  it("should fetch logs and connect to websocket", async () => {
     const queryClient = createTestQueryClient();
     const logs = generateLogs(5);
     const fetchSpy = jest
diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.ts b/site/src/modules/resources/AgentLogs/useAgentLogs.ts
@@ -17,16 +17,13 @@ export type UseAgentLogsOptions = Readonly<{
 
 /**
  * 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.
+ * 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 { data: logs, isFetched } = useQuery({ ...queryOptions, enabled });
@@ -55,7 +52,17 @@ export function useAgentLogs(
   });
 
   useEffect(() => {
-    if (agentLifeCycleState !== "starting" || !isFetched) {
+    // Stream data only for new logs. Old logs should be loaded beforehand
+    // using a regular fetch to avoid overloading the websocket with all
+    // logs at once.
+    if (!isFetched) {
+      return;
+    }
+
+    // If the agent is off, we don't need to stream logs. This is the only state
+    // where the Coder API can't receive logs for the agent from third-party
+    // apps like envbuilder.
+    if (agentLifeCycleState === "off") {
       return;
     }
 
