Skip to content

fix: fix duplicated agent logs #17806

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
May 15, 2025
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 6 additions & 24 deletions site/src/api/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -221,11 +221,11 @@ export const watchBuildLogsByTemplateVersionId = (

export const watchWorkspaceAgentLogs = (
agentId: string,
{ after, onMessage, onDone, onError }: WatchWorkspaceAgentLogsOptions,
params?: WatchWorkspaceAgentLogsParams,
) => {
const searchParams = new URLSearchParams({
follow: "true",
after: after.toString(),
after: params?.after ? params?.after.toString() : "",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we make this params?.after.toString() ?? ""?

});

/**
Expand All @@ -237,32 +237,14 @@ export const watchWorkspaceAgentLogs = (
searchParams.set("no_compression", "");
}

const socket = createWebSocket(
`/api/v2/workspaceagents/${agentId}/logs`,
return new OneWayWebSocket<TypesGen.WorkspaceAgentLog[]>({
apiRoute: `/api/v2/workspaceagents/${agentId}/logs`,
searchParams,
);

socket.addEventListener("message", (event) => {
const logs = JSON.parse(event.data) as TypesGen.WorkspaceAgentLog[];
onMessage(logs);
});

socket.addEventListener("error", () => {
onError(new Error("socket errored"));
});

socket.addEventListener("close", () => {
onDone?.();
});

return socket;
};

type WatchWorkspaceAgentLogsOptions = {
after: number;
onMessage: (logs: TypesGen.WorkspaceAgentLog[]) => void;
onDone?: () => void;
onError: (error: Error) => void;
type WatchWorkspaceAgentLogsParams = {
after?: number;
};

type WatchBuildLogsByBuildIdOptions = {
Expand Down
16 changes: 6 additions & 10 deletions site/src/api/queries/workspaces.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import type {
ProvisionerLogLevel,
UsageAppName,
Workspace,
WorkspaceAgentLog,
WorkspaceBuild,
WorkspaceBuildParameter,
WorkspacesRequest,
Expand All @@ -16,6 +17,7 @@ import type {
QueryClient,
QueryOptions,
UseMutationOptions,
UseQueryOptions,
} from "react-query";
import { disabledRefetchOptions } from "./util";
import { workspaceBuildsKey } from "./workspaceBuilds";
Expand Down Expand Up @@ -337,20 +339,14 @@ export const buildLogs = (workspace: Workspace) => {
};
};

export const agentLogsKey = (workspaceId: string, agentId: string) => [
"workspaces",
workspaceId,
"agents",
agentId,
"logs",
];
export const agentLogsKey = (agentId: string) => ["agents", agentId, "logs"];

export const agentLogs = (workspaceId: string, agentId: string) => {
export const agentLogs = (agentId: string) => {
return {
queryKey: agentLogsKey(workspaceId, agentId),
queryKey: agentLogsKey(agentId),
queryFn: () => API.getWorkspaceAgentLogs(agentId),
...disabledRefetchOptions,
};
} satisfies UseQueryOptions<WorkspaceAgentLog[]>;
};

// workspace usage options
Expand Down
51 changes: 12 additions & 39 deletions site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -2,44 +2,29 @@ import { act, renderHook, waitFor } from "@testing-library/react";
import { API } from "api/api";
import * as APIModule from "api/api";
import { agentLogsKey } from "api/queries/workspaces";
import type { WorkspaceAgentLog } from "api/typesGenerated";
import type { WorkspaceAgent, WorkspaceAgentLog } from "api/typesGenerated";
import WS from "jest-websocket-mock";
import { type QueryClient, QueryClientProvider } from "react-query";
import { MockWorkspace, MockWorkspaceAgent } from "testHelpers/entities";
import { MockWorkspaceAgent } from "testHelpers/entities";
import { createTestQueryClient } from "testHelpers/renderHelpers";
import { type UseAgentLogsOptions, useAgentLogs } from "./useAgentLogs";
import { useAgentLogs } from "./useAgentLogs";

afterEach(() => {
WS.clean();
});

describe("useAgentLogs", () => {
it("should not fetch logs if disabled", async () => {
const queryClient = createTestQueryClient();
const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
renderUseAgentLogs(queryClient, {
workspaceId: MockWorkspace.id,
agentId: MockWorkspaceAgent.id,
agentLifeCycleState: "ready",
enabled: false,
});
expect(fetchSpy).not.toHaveBeenCalled();
expect(wsSpy).not.toHaveBeenCalled();
});

it("should return existing logs without network calls if state is off", async () => {
const queryClient = createTestQueryClient();
queryClient.setQueryData(
agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
agentLogsKey(MockWorkspaceAgent.id),
generateLogs(5),
);
const fetchSpy = jest.spyOn(API, "getWorkspaceAgentLogs");
const wsSpy = jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
const { result } = renderUseAgentLogs(queryClient, {
workspaceId: MockWorkspace.id,
agentId: MockWorkspaceAgent.id,
agentLifeCycleState: "off",
...MockWorkspaceAgent,
lifecycle_state: "off",
});
await waitFor(() => {
expect(result.current).toHaveLength(5);
Expand All @@ -55,9 +40,8 @@ describe("useAgentLogs", () => {
.mockResolvedValueOnce(generateLogs(5));
jest.spyOn(APIModule, "watchWorkspaceAgentLogs");
const { result } = renderUseAgentLogs(queryClient, {
workspaceId: MockWorkspace.id,
agentId: MockWorkspaceAgent.id,
agentLifeCycleState: "ready",
...MockWorkspaceAgent,
lifecycle_state: "ready",
});
await waitFor(() => {
expect(result.current).toHaveLength(5);
Expand All @@ -77,11 +61,7 @@ describe("useAgentLogs", () => {
MockWorkspaceAgent.id
}/logs?follow&after=${logs[logs.length - 1].id}`,
);
const { result } = renderUseAgentLogs(queryClient, {
workspaceId: MockWorkspace.id,
agentId: MockWorkspaceAgent.id,
agentLifeCycleState: "starting",
});
const { result } = renderUseAgentLogs(queryClient, MockWorkspaceAgent);
await waitFor(() => {
expect(result.current).toHaveLength(5);
});
Expand All @@ -102,11 +82,7 @@ describe("useAgentLogs", () => {
MockWorkspaceAgent.id
}/logs?follow&after=${logs[logs.length - 1].id}`,
);
const { result } = renderUseAgentLogs(queryClient, {
workspaceId: MockWorkspace.id,
agentId: MockWorkspaceAgent.id,
agentLifeCycleState: "starting",
});
const { result } = renderUseAgentLogs(queryClient, MockWorkspaceAgent);
await waitFor(() => {
expect(result.current).toHaveLength(5);
});
Expand All @@ -120,11 +96,8 @@ describe("useAgentLogs", () => {
});
});

function renderUseAgentLogs(
queryClient: QueryClient,
options: UseAgentLogsOptions,
) {
return renderHook(() => useAgentLogs(options), {
function renderUseAgentLogs(queryClient: QueryClient, agent: WorkspaceAgent) {
return renderHook(() => useAgentLogs(agent), {
wrapper: ({ children }) => (
<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
),
Expand Down
108 changes: 43 additions & 65 deletions site/src/modules/resources/AgentLogs/useAgentLogs.ts
Original file line number Diff line number Diff line change
@@ -1,95 +1,73 @@
import { watchWorkspaceAgentLogs } from "api/api";
import { agentLogs } from "api/queries/workspaces";
import type {
WorkspaceAgent,
WorkspaceAgentLifecycle,
WorkspaceAgentLog,
} from "api/typesGenerated";
import { displayError } from "components/GlobalSnackbar/utils";
import { useEffectEvent } from "hooks/hookPolyfills";
import { useEffect, useRef } from "react";
import { useEffect } from "react";
import { useQuery, useQueryClient } from "react-query";

export type UseAgentLogsOptions = Readonly<{
workspaceId: string;
agentId: string;
agentLifeCycleState: WorkspaceAgentLifecycle;
enabled?: boolean;
}>;
const ON_GOING_STATES: WorkspaceAgentLifecycle[] = ["starting", "created"];

/**
* 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,
agent: WorkspaceAgent,
): 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 });

// Track the ID of the last log received when the initial logs response comes
// back. If the logs are not complete, the ID will mark the start point of the
// Web sockets response so that the remaining logs can be received over time
const lastQueriedLogId = useRef(0);
useEffect(() => {
const isAlreadyTracking = lastQueriedLogId.current !== 0;
if (isAlreadyTracking) {
return;
}
const agentLogsOptions = agentLogs(agent.id);
const shouldUseSocket = ON_GOING_STATES.includes(agent.lifecycle_state);
const { data: logs } = useQuery({
...agentLogsOptions,
enabled: !shouldUseSocket,
});

const lastLog = logs?.at(-1);
if (lastLog !== undefined) {
lastQueriedLogId.current = lastLog.id;
}
}, [logs]);
const appendAgentLogs = useEffectEvent(
async (newLogs: WorkspaceAgentLog[]) => {
await queryClient.cancelQueries(agentLogsOptions.queryKey);
queryClient.setQueryData<WorkspaceAgentLog[]>(
agentLogsOptions.queryKey,
(oldLogs) => (oldLogs ? [...oldLogs, ...newLogs] : newLogs),
);
},
);

const addLogs = useEffectEvent((newLogs: WorkspaceAgentLog[]) => {
queryClient.setQueryData(
queryOptions.queryKey,
(oldData: WorkspaceAgentLog[] = []) => [...oldData, ...newLogs],
);
const refreshAgentLogs = useEffectEvent(() => {
queryClient.invalidateQueries(agentLogsOptions.queryKey);
});

useEffect(() => {
// 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) {
if (!shouldUseSocket) {
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;
}
const socket = watchWorkspaceAgentLogs(agent.id, { after: 0 });
socket.addEventListener("message", (e) => {
if (e.parseError) {
console.warn("Error parsing agent log: ", e.parseError);
return;
}
appendAgentLogs(e.parsedMessage);
});

const socket = watchWorkspaceAgentLogs(agentId, {
after: lastQueriedLogId.current,
onMessage: (newLogs) => {
// Prevent new logs getting added when a connection is not open
if (socket.readyState !== WebSocket.OPEN) {
return;
}
addLogs(newLogs);
},
onError: (error) => {
// For some reason Firefox and Safari throw an error when a websocket
// connection is close in the middle of a message and because of that we
// can't safely show to the users an error message since most of the
// time they are just internal stuff. This does not happen to Chrome at
// all and I tried to find better way to "soft close" a WS connection on
// those browsers without success.
console.error(error);
},
socket.addEventListener("error", (e) => {
console.error("Error in agent log socket: ", e);
displayError(
"Unable to watch the agent logs",
"Please try refreshing the browser",
);
socket.close();
});

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

return logs;
}
15 changes: 8 additions & 7 deletions site/src/modules/resources/AgentRow.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -88,12 +88,7 @@ export const AgentRow: FC<AgentRowProps> = ({
["starting", "start_timeout"].includes(agent.lifecycle_state) &&
hasStartupFeatures,
);
const agentLogs = useAgentLogs({
workspaceId: workspace.id,
agentId: agent.id,
agentLifeCycleState: agent.lifecycle_state,
enabled: showLogs,
});
const agentLogs = useAgentLogs(agent);
const logListRef = useRef<List>(null);
const logListDivRef = useRef<HTMLDivElement>(null);
const startupLogs = useMemo(() => {
Expand Down Expand Up @@ -160,7 +155,13 @@ export const AgentRow: FC<AgentRowProps> = ({
select: (res) => res.containers.filter((c) => c.status === "running"),
// TODO: Implement a websocket connection to get updates on containers
// without having to poll.
refetchInterval: 10_000,
refetchInterval: (data) => {
if (!data) {
return false;
}

return data.length > 0 ? 10_000 : false;
},
});

return (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ const meta: Meta<typeof DownloadAgentLogsButton> = {
parameters: {
queries: [
{
key: agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
key: agentLogsKey(MockWorkspaceAgent.id),
data: generateLogs(5),
},
],
Expand Down
2 changes: 1 addition & 1 deletion site/src/modules/resources/DownloadAgentLogsButton.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ export const DownloadAgentLogsButton: FC<DownloadAgentLogsButtonProps> = ({
const [isDownloading, setIsDownloading] = useState(false);

const fetchLogs = async () => {
const queryOpts = agentLogs(workspaceId, agent.id);
const queryOpts = agentLogs(agent.id);
let logs = queryClient.getQueryData<WorkspaceAgentLog[]>(
queryOpts.queryKey,
);
Expand Down
Loading
Loading