diff --git a/ui/src/hooks/useWebsocket.tsx b/ui/src/hooks/useWebsocket.tsx
index 1a95236..5db879c 100644
--- a/ui/src/hooks/useWebsocket.tsx
+++ b/ui/src/hooks/useWebsocket.tsx
@@ -4,14 +4,14 @@ import { getLastMessage, setLastMessage } from "@/services/store";
 import { initializeConnection } from "@/services/websocket";
 import { useEffect, useState } from "react";
 
-export function useWebsocket(url: string) {
+export function useWebsocket(agent: string) {
     const [status, setStatus] = useState<WebsocketStatus>();
     const [message, setMessage] = useState<StatusMessage | null>(
         getLastMessage(),
     );
 
     useEffect(() => {
-        initializeConnection(url, {
+        initializeConnection(agent, {
             onMessage: (data) => {
                 const parsed = JSON.parse(data);
 
@@ -23,7 +23,7 @@ export function useWebsocket(url: string) {
             onClose: () => setStatus("disconnected"),
             onError: () => setStatus("error"),
         });
-    }, [url]);
+    }, [agent]);
 
     return { status, message };
 }
diff --git a/ui/src/pages/agent.tsx b/ui/src/pages/agent.tsx
index 3d44479..c2b111c 100644
--- a/ui/src/pages/agent.tsx
+++ b/ui/src/pages/agent.tsx
@@ -20,18 +20,16 @@ import { TimePeriodSelector } from "@/components/TimePeriodSelector";
 
 export function AgentPage() {
     const { agent } = useParams();
-    const { status, message } = useWebsocket(`ws://localhost:3000/ws/${agent}`);
+    const { status, message } = useWebsocket(agent!);
 
     const [period, setPeriod] = useState<TimePeriod>("realtime");
     const [history, setHistory] = useState<StatusMessage[] | null>(null);
 
     useEffect(() => {
-        if (period === "realtime") return;
+        if (!agent || period === "realtime") return;
 
         const fetchData = async () => {
-            const data = await getHistoricalData(
-                `http://localhost:3000/history/${agent}/${period}`,
-            );
+            const data = await getHistoricalData(agent, period);
             setHistory(data);
         };
 
diff --git a/ui/src/services/store.ts b/ui/src/services/store.ts
index 37116c8..0821bc3 100644
--- a/ui/src/services/store.ts
+++ b/ui/src/services/store.ts
@@ -20,8 +20,15 @@ export function getRealtimeData(): StatusMessage[] {
     return realtimeData ?? [];
 }
 
-export async function getHistoricalData(url: string): Promise<StatusMessage[]> {
-    return await fetch(url)
+export async function getHistoricalData(
+    agent: string,
+    period: string,
+): Promise<StatusMessage[]> {
+    return await fetch(
+        import.meta.env.DEV
+            ? `http://localhost:3000/agent/${agent}/${period}`
+            : `/agent/${agent}`,
+    )
         .then((res) => res.json())
         .then((data) => data as StatusMessage[]);
 }
diff --git a/ui/src/services/websocket.ts b/ui/src/services/websocket.ts
index 7cd5fa9..5c192ae 100644
--- a/ui/src/services/websocket.ts
+++ b/ui/src/services/websocket.ts
@@ -5,8 +5,12 @@ type Callbacks = {
     onClose?: () => void;
 };
 
-export function initializeConnection(url: string, callbacks: Callbacks) {
-    const ws = new WebSocket(url);
+export function initializeConnection(agent: string, callbacks: Callbacks) {
+    const ws = new WebSocket(
+        import.meta.env.DEV
+            ? `ws://localhost:3000/ws/${agent}`
+            : `/ws/${agent}`,
+    );
 
     ws.onopen = () => callbacks.onOpen?.();
     ws.onmessage = (event) => callbacks.onMessage(event.data);