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

This commit is contained in:
Bruno Quaresma
2024-09-03 15:01:02 -03:00
committed by GitHub
parent fcb0ce1f1b
commit 242b1ea4ca
2 changed files with 18 additions and 12 deletions
@@ -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
@@ -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;
}