Skip to content

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

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 2 commits into from
Sep 3, 2024
Merged
Show file tree
Hide file tree
Changes from all 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
11 changes: 5 additions & 6 deletions site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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);
Expand All @@ -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,
Expand All @@ -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
Expand Down
19 changes: 13 additions & 6 deletions site/src/modules/resources/AgentLogs/useAgentLogs.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 });
Expand Down Expand Up @@ -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;
}

Expand Down
Loading