Skip to content

Commit 3f9a1bd

Browse files
BrunoQuaresmaBrunoQuaresma
committed
fix: duplicated agent logs
1 parent 64807e1 commit 3f9a1bd

12 files changed

+102
-174
lines changed

‎site/src/api/api.ts

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,11 @@
1919
*
2020
* For example, `utils/delay` must be imported using `../utils/delay` instead.
2121
*/
22-
import globalAxios, { type AxiosInstance, isAxiosError } from "axios";
22+
import globalAxios, {
23+
type AxiosInstance,
24+
type AxiosRequestConfig,
25+
isAxiosError,
26+
} from "axios";
2327
import type dayjs from "dayjs";
2428
import userAgentParser from "ua-parser-js";
2529
import { OneWayWebSocket } from "../utils/OneWayWebSocket";
@@ -221,11 +225,11 @@ export const watchBuildLogsByTemplateVersionId = (
221225

222226
export const watchWorkspaceAgentLogs = (
223227
agentId: string,
224-
{ after, onMessage, onDone, onError }: WatchWorkspaceAgentLogsOptions,
228+
params?: WatchWorkspaceAgentLogsParams,
225229
) => {
226230
const searchParams = new URLSearchParams({
227231
follow: "true",
228-
after: after.toString(),
232+
after: params?.after ? params?.after.toString() : "",
229233
});
230234

231235
/**
@@ -237,32 +241,14 @@ export const watchWorkspaceAgentLogs = (
237241
searchParams.set("no_compression", "");
238242
}
239243

240-
const socket = createWebSocket(
241-
`/api/v2/workspaceagents/${agentId}/logs`,
244+
return new OneWayWebSocket<TypesGen.WorkspaceAgentLog[]>({
245+
apiRoute: `/api/v2/workspaceagents/${agentId}/logs`,
242246
searchParams,
243-
);
244-
245-
socket.addEventListener("message", (event) => {
246-
const logs = JSON.parse(event.data) as TypesGen.WorkspaceAgentLog[];
247-
onMessage(logs);
248-
});
249-
250-
socket.addEventListener("error", () => {
251-
onError(new Error("socket errored"));
252247
});
253-
254-
socket.addEventListener("close", () => {
255-
onDone?.();
256-
});
257-
258-
return socket;
259248
};
260249

261-
type WatchWorkspaceAgentLogsOptions = {
262-
after: number;
263-
onMessage: (logs: TypesGen.WorkspaceAgentLog[]) => void;
264-
onDone?: () => void;
265-
onError: (error: Error) => void;
250+
type WatchWorkspaceAgentLogsParams = {
251+
after?: number;
266252
};
267253

268254
type WatchBuildLogsByBuildIdOptions = {

‎site/src/api/queries/workspaces.ts

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import type {
55
ProvisionerLogLevel,
66
UsageAppName,
77
Workspace,
8+
WorkspaceAgentLog,
89
WorkspaceBuild,
910
WorkspaceBuildParameter,
1011
WorkspacesRequest,
@@ -16,6 +17,7 @@ import type {
1617
QueryClient,
1718
QueryOptions,
1819
UseMutationOptions,
20+
UseQueryOptions,
1921
} from "react-query";
2022
import { disabledRefetchOptions } from "./util";
2123
import { workspaceBuildsKey } from "./workspaceBuilds";
@@ -337,20 +339,14 @@ export const buildLogs = (workspace: Workspace) => {
337339
};
338340
};
339341

340-
export const agentLogsKey = (workspaceId: string, agentId: string) => [
341-
"workspaces",
342-
workspaceId,
343-
"agents",
344-
agentId,
345-
"logs",
346-
];
342+
export const agentLogsKey = (agentId: string) => ["agents", agentId, "logs"];
347343

348-
export const agentLogs = (workspaceId: string, agentId: string) => {
344+
export const agentLogs = (agentId: string) => {
349345
return {
350-
queryKey: agentLogsKey(workspaceId, agentId),
346+
queryKey: agentLogsKey(agentId),
351347
queryFn: () => API.getWorkspaceAgentLogs(agentId),
352348
...disabledRefetchOptions,
353-
};
349+
} satisfies UseQueryOptions<WorkspaceAgentLog[]>;
354350
};
355351

356352
// workspace usage options

‎site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx

Lines changed: 12 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -2,44 +2,29 @@ import { act, renderHook, waitFor } from "@testing-library/react";
22
import { API } from "api/api";
33
import * as APIModule from "api/api";
44
import { agentLogsKey } from "api/queries/workspaces";
5-
import type { WorkspaceAgentLog } from "api/typesGenerated";
5+
import type { WorkspaceAgent, WorkspaceAgentLog } from "api/typesGenerated";
66
import WS from "jest-websocket-mock";
77
import { type QueryClient, QueryClientProvider } from "react-query";
8-
import { MockWorkspace, MockWorkspaceAgent } from "testHelpers/entities";
8+
import { MockWorkspaceAgent } from "testHelpers/entities";
99
import { createTestQueryClient } from "testHelpers/renderHelpers";
10-
import { type UseAgentLogsOptions, useAgentLogs } from "./useAgentLogs";
10+
import { useAgentLogs } from "./useAgentLogs";
1111

1212
afterEach(() => {
1313
WS.clean();
1414
});
1515

1616
describe("useAgentLogs", () => {
17-
it("should not fetch logs if disabled", async () => {
18-
const queryClient = createTestQueryClient();
19-
const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
20-
const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
21-
renderUseAgentLogs(queryClient, {
22-
workspaceId: MockWorkspace.id,
23-
agentId: MockWorkspaceAgent.id,
24-
agentLifeCycleState: "ready",
25-
enabled: false,
26-
});
27-
expect(fetchSpy).not.toHaveBeenCalled();
28-
expect(wsSpy).not.toHaveBeenCalled();
29-
});
30-
3117
it("should return existing logs without network calls if state is off", async () => {
3218
const queryClient = createTestQueryClient();
3319
queryClient.setQueryData(
34-
agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
20+
agentLogsKey(MockWorkspaceAgent.id),
3521
generateLogs(5),
3622
);
3723
const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
3824
const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
3925
const { result } = renderUseAgentLogs(queryClient, {
40-
workspaceId: MockWorkspace.id,
41-
agentId: MockWorkspaceAgent.id,
42-
agentLifeCycleState: "off",
26+
...MockWorkspaceAgent,
27+
lifecycle_state: "off",
4328
});
4429
await waitFor(() => {
4530
expect(result.current).toHaveLength(5);
@@ -55,9 +40,8 @@ describe("useAgentLogs", () => {
5540
.mockResolvedValueOnce(generateLogs(5));
5641
jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
5742
const { result } = renderUseAgentLogs(queryClient, {
58-
workspaceId: MockWorkspace.id,
59-
agentId: MockWorkspaceAgent.id,
60-
agentLifeCycleState: "ready",
43+
...MockWorkspaceAgent,
44+
lifecycle_state: "ready",
6145
});
6246
await waitFor(() => {
6347
expect(result.current).toHaveLength(5);
@@ -77,11 +61,7 @@ describe("useAgentLogs", () => {
7761
MockWorkspaceAgent.id
7862
}/logs?follow&after=${logs[logs.length - 1].id}`,
7963
);
80-
const { result } = renderUseAgentLogs(queryClient, {
81-
workspaceId: MockWorkspace.id,
82-
agentId: MockWorkspaceAgent.id,
83-
agentLifeCycleState: "starting",
84-
});
64+
const { result } = renderUseAgentLogs(queryClient, MockWorkspaceAgent);
8565
await waitFor(() => {
8666
expect(result.current).toHaveLength(5);
8767
});
@@ -102,11 +82,7 @@ describe("useAgentLogs", () => {
10282
MockWorkspaceAgent.id
10383
}/logs?follow&after=${logs[logs.length - 1].id}`,
10484
);
105-
const { result } = renderUseAgentLogs(queryClient, {
106-
workspaceId: MockWorkspace.id,
107-
agentId: MockWorkspaceAgent.id,
108-
agentLifeCycleState: "starting",
109-
});
85+
const { result } = renderUseAgentLogs(queryClient, MockWorkspaceAgent);
11086
await waitFor(() => {
11187
expect(result.current).toHaveLength(5);
11288
});
@@ -120,11 +96,8 @@ describe("useAgentLogs", () => {
12096
});
12197
});
12298

123-
function renderUseAgentLogs(
124-
queryClient: QueryClient,
125-
options: UseAgentLogsOptions,
126-
) {
127-
return renderHook(() => useAgentLogs(options), {
99+
function renderUseAgentLogs(queryClient: QueryClient, agent: WorkspaceAgent) {
100+
return renderHook(() => useAgentLogs(agent), {
128101
wrapper: ({ children }) => (
129102
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
130103
),

‎site/src/modules/resources/AgentLogs/useAgentLogs.ts

Lines changed: 43 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,95 +1,73 @@
11
import { watchWorkspaceAgentLogs } from "api/api";
22
import { agentLogs } from "api/queries/workspaces";
33
import type {
4+
WorkspaceAgent,
45
WorkspaceAgentLifecycle,
56
WorkspaceAgentLog,
67
} from "api/typesGenerated";
8+
import { displayError } from "components/GlobalSnackbar/utils";
79
import { useEffectEvent } from "hooks/hookPolyfills";
8-
import { useEffect, useRef } from "react";
10+
import { useEffect } from "react";
911
import { useQuery, useQueryClient } from "react-query";
1012

11-
export type UseAgentLogsOptions = Readonly<{
12-
workspaceId: string;
13-
agentId: string;
14-
agentLifeCycleState: WorkspaceAgentLifecycle;
15-
enabled?: boolean;
16-
}>;
13+
const ON_GOING_STATES: WorkspaceAgentLifecycle[] = ["starting", "created"];
1714

18-
/**
19-
* Defines a custom hook that gives you all workspace agent logs for a given
20-
* workspace.Depending on the status of the workspace, all logs may or may not
21-
* be available.
22-
*/
2315
export function useAgentLogs(
24-
options: UseAgentLogsOptions,
16+
agent: WorkspaceAgent,
2517
): readonly WorkspaceAgentLog[] | undefined {
26-
const { workspaceId, agentId, agentLifeCycleState, enabled = true } = options;
2718
const queryClient = useQueryClient();
28-
const queryOptions = agentLogs(workspaceId, agentId);
29-
const { data: logs, isFetched } = useQuery({ ...queryOptions, enabled });
30-
31-
// Track the ID of the last log received when the initial logs response comes
32-
// back. If the logs are not complete, the ID will mark the start point of the
33-
// Web sockets response so that the remaining logs can be received over time
34-
const lastQueriedLogId = useRef(0);
35-
useEffect(() => {
36-
const isAlreadyTracking = lastQueriedLogId.current !== 0;
37-
if (isAlreadyTracking) {
38-
return;
39-
}
19+
const agentLogsOptions = agentLogs(agent.id);
20+
const shouldUseSocket = ON_GOING_STATES.includes(agent.lifecycle_state);
21+
const { data: logs } = useQuery({
22+
...agentLogsOptions,
23+
enabled: !shouldUseSocket,
24+
});
4025

41-
const lastLog = logs?.at(-1);
42-
if (lastLog !== undefined) {
43-
lastQueriedLogId.current = lastLog.id;
44-
}
45-
}, [logs]);
26+
const appendAgentLogs = useEffectEvent(
27+
async (newLogs: WorkspaceAgentLog[]) => {
28+
await queryClient.cancelQueries(agentLogsOptions.queryKey);
29+
queryClient.setQueryData<WorkspaceAgentLog[]>(
30+
agentLogsOptions.queryKey,
31+
(oldLogs) => (oldLogs ? [...oldLogs, ...newLogs] : newLogs),
32+
);
33+
},
34+
);
4635

47-
const addLogs = useEffectEvent((newLogs: WorkspaceAgentLog[]) => {
48-
queryClient.setQueryData(
49-
queryOptions.queryKey,
50-
(oldData: WorkspaceAgentLog[] = []) => [...oldData, ...newLogs],
51-
);
36+
const refreshAgentLogs = useEffectEvent(() => {
37+
queryClient.invalidateQueries(agentLogsOptions.queryKey);
5238
});
5339

5440
useEffect(() => {
55-
// Stream data only for new logs. Old logs should be loaded beforehand
56-
// using a regular fetch to avoid overloading the websocket with all
57-
// logs at once.
58-
if (!isFetched) {
41+
if (!shouldUseSocket) {
5942
return;
6043
}
6144

62-
// If the agent is off, we don't need to stream logs. This is the only state
63-
// where the Coder API can't receive logs for the agent from third-party
64-
// apps like envbuilder.
65-
if (agentLifeCycleState === "off") {
66-
return;
67-
}
45+
const socket = watchWorkspaceAgentLogs(agent.id, { after: 0 });
46+
socket.addEventListener("message", (e) => {
47+
if (e.parseError) {
48+
console.warn("Error parsing agent log: ", e.parseError);
49+
return;
50+
}
51+
appendAgentLogs(e.parsedMessage);
52+
});
6853

69-
const socket = watchWorkspaceAgentLogs(agentId, {
70-
after: lastQueriedLogId.current,
71-
onMessage: (newLogs) => {
72-
// Prevent new logs getting added when a connection is not open
73-
if (socket.readyState !== WebSocket.OPEN) {
74-
return;
75-
}
76-
addLogs(newLogs);
77-
},
78-
onError: (error) => {
79-
// For some reason Firefox and Safari throw an error when a websocket
80-
// connection is close in the middle of a message and because of that we
81-
// can't safely show to the users an error message since most of the
82-
// time they are just internal stuff. This does not happen to Chrome at
83-
// all and I tried to find better way to "soft close" a WS connection on
84-
// those browsers without success.
85-
console.error(error);
86-
},
54+
socket.addEventListener("error", (e) => {
55+
console.error("Error in agent log socket: ", e);
56+
displayError(
57+
"Unable to watch the agent logs",
58+
"Please try refreshing the browser",
59+
);
60+
socket.close();
8761
});
8862

8963
return () => {
9064
socket.close();
65+
// For some reason, after closing the socket, a few logs still getting
66+
// generated in the BE. This is a workaround to avoid we don't display
67+
// them in the UI.
68+
refreshAgentLogs();
9169
};
92-
}, [addLogs, agentId, agentLifeCycleState, isFetched]);
70+
}, [agent.id, shouldUseSocket, appendAgentLogs, refreshAgentLogs]);
9371

9472
return logs;
9573
}

‎site/src/modules/resources/AgentRow.tsx

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,7 @@ export const AgentRow: FC<AgentRowProps> = ({
8888
["starting", "start_timeout"].includes(agent.lifecycle_state) &&
8989
hasStartupFeatures,
9090
);
91-
const agentLogs = useAgentLogs({
92-
workspaceId: workspace.id,
93-
agentId: agent.id,
94-
agentLifeCycleState: agent.lifecycle_state,
95-
enabled: showLogs,
96-
});
91+
const agentLogs = useAgentLogs(agent);
9792
const logListRef = useRef<List>(null);
9893
const logListDivRef = useRef<HTMLDivElement>(null);
9994
const startupLogs = useMemo(() => {
@@ -160,7 +155,13 @@ export const AgentRow: FC<AgentRowProps> = ({
160155
select: (res) => res.containers.filter((c) => c.status === "running"),
161156
// TODO: Implement a websocket connection to get updates on containers
162157
// without having to poll.
163-
refetchInterval: 10_000,
158+
refetchInterval: (data) => {
159+
if (!data) {
160+
return false;
161+
}
162+
163+
return data.length > 0 ? 10_000 : false;
164+
},
164165
});
165166

166167
return (

‎site/src/modules/resources/DownloadAgentLogsButton.stories.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ const meta: Meta<typeof DownloadAgentLogsButton> = {
1515
parameters: {
1616
queries: [
1717
{
18-
key: agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
18+
key: agentLogsKey(MockWorkspaceAgent.id),
1919
data: generateLogs(5),
2020
},
2121
],

‎site/src/modules/resources/DownloadAgentLogsButton.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ export const DownloadAgentLogsButton: FC<DownloadAgentLogsButtonProps> = ({
2323
const [isDownloading, setIsDownloading] = useState(false);
2424

2525
const fetchLogs = async () => {
26-
const queryOpts = agentLogs(workspaceId, agent.id);
26+
const queryOpts = agentLogs(agent.id);
2727
let logs = queryClient.getQueryData<WorkspaceAgentLog[]>(
2828
queryOpts.queryKey,
2929
);

0 commit comments

Comments
 (0)