coder · BrunoQuaresma · May 15, 2025 · May 13, 2025 · May 13, 2025 · May 14, 2025
diff --git a/site/src/api/api.ts b/site/src/api/api.ts
@@ -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?.toString() ?? "",
 	});
 
 	/**
@@ -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 = {

diff --git a/site/src/api/queries/workspaces.ts b/site/src/api/queries/workspaces.ts
@@ -5,6 +5,7 @@ import type {
 	ProvisionerLogLevel,
 	UsageAppName,
 	Workspace,
+	WorkspaceAgentLog,
 	WorkspaceBuild,
 	WorkspaceBuildParameter,
 	WorkspacesRequest,
@@ -20,6 +21,7 @@ import type {
 	QueryClient,
 	QueryOptions,
 	UseMutationOptions,
+	UseQueryOptions,
 } from "react-query";
 import { checkAuthorization } from "./authCheck";
 import { disabledRefetchOptions } from "./util";
@@ -342,20 +344,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

diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx b/site/src/modules/resources/AgentLogs/useAgentLogs.test.tsx
diff --git a/site/src/modules/resources/AgentLogs/useAgentLogs.ts b/site/src/modules/resources/AgentLogs/useAgentLogs.ts
diff --git a/site/src/modules/resources/AgentRow.tsx b/site/src/modules/resources/AgentRow.tsx
@@ -30,7 +30,6 @@ import { AgentDevcontainerCard } from "./AgentDevcontainerCard";
 import { AgentLatency } from "./AgentLatency";
 import { AGENT_LOG_LINE_HEIGHT } from "./AgentLogs/AgentLogLine";
 import { AgentLogs } from "./AgentLogs/AgentLogs";
-import { useAgentLogs } from "./AgentLogs/useAgentLogs";
 import { AgentMetadata } from "./AgentMetadata";
 import { AgentStatus } from "./AgentStatus";
 import { AgentVersion } from "./AgentVersion";
@@ -40,6 +39,7 @@ import { PortForwardButton } from "./PortForwardButton";
 import { AgentSSHButton } from "./SSHButton/SSHButton";
 import { TerminalLink } from "./TerminalLink/TerminalLink";
 import { VSCodeDesktopButton } from "./VSCodeDesktopButton/VSCodeDesktopButton";
+import { useAgentLogs } from "./useAgentLogs";
 
 export interface AgentRowProps {
 	agent: WorkspaceAgent;
@@ -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, showLogs);
 	const logListRef = useRef<List>(null);
 	const logListDivRef = useRef<HTMLDivElement>(null);
 	const startupLogs = useMemo(() => {

diff --git a/site/src/modules/resources/DownloadAgentLogsButton.stories.tsx b/site/src/modules/resources/DownloadAgentLogsButton.stories.tsx
@@ -15,7 +15,7 @@ const meta: Meta<typeof DownloadAgentLogsButton> = {
 	parameters: {
 		queries: [
 			{
-				key: agentLogsKey(MockWorkspace.id, MockWorkspaceAgent.id),
+				key: agentLogsKey(MockWorkspaceAgent.id),
 				data: generateLogs(5),
 			},
 		],

diff --git a/site/src/modules/resources/DownloadAgentLogsButton.tsx b/site/src/modules/resources/DownloadAgentLogsButton.tsx
@@ -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,
 		);