diff --git a/src/background.ts b/src/background.ts
index 196b9b0..46e2ab7 100644
--- a/src/background.ts
+++ b/src/background.ts
@@ -7,7 +7,7 @@ const TOGGL_KEY = "toggl_api_token";
 const HABITIFY_KEY = "habitify_api_token";
 
 const projectMap = new Map<string, Project>();
-async function getProjectByName(name: string) {
+const getProjectByName = async (name: string) => {
   if (projectMap.size === 0) {
     const toggl = await useToggl(await getTogglApiToken());
     const workspaces = await toggl.getWorkspaces();
@@ -19,45 +19,45 @@ async function getProjectByName(name: string) {
     }
   }
   return projectMap.get(name);
-}
+};
 
-async function getTogglApiToken() {
+const getTogglApiToken = async () => {
   const token = (await browser.storage.local.get(TOGGL_KEY))[TOGGL_KEY];
   if (typeof token !== "string") throw "Toggl API Token is empty";
   return token;
-}
+};
 
-async function getHabitifyApiToken() {
+const getHabitifyApiToken = async () => {
   const token = (await browser.storage.local.get(HABITIFY_KEY))[HABITIFY_KEY];
   if (typeof token !== "string") throw "Habitify API Token is empty";
   return token;
-}
+};
 
-async function processVerifyTogglMessage(_: VerifyTogglMessage) {
+const processVerifyTogglMessage = async (_: VerifyTogglMessage) => {
   try {
     return await verifyTogglApiToken(await getTogglApiToken());
   } catch {
     return false;
   }
-}
+};
 
-async function processVerifyHabitifyMessage(_: VerifyHabitifyMessage) {
+const processVerifyHabitifyMessage = async (_: VerifyHabitifyMessage) => {
   try {
     return await verifyHabitifyApiToken(await getHabitifyApiToken());
   } catch {
     return false;
   }
-}
+};
 
-async function processTokenTogglMessage(message: TokenTogglMessage) {
+const processTokenTogglMessage = async (message: TokenTogglMessage) => {
   return await browser.storage.local.set({ [TOGGL_KEY]: message.token });
-}
+};
 
-async function processTokenHabitifyMessage(message: TokenHabitifyMessage) {
+const processTokenHabitifyMessage = async (message: TokenHabitifyMessage) => {
   return await browser.storage.local.set({ [HABITIFY_KEY]: message.token });
-}
+};
 
-async function processTimerMessage(message: TimerMessage) {
+const processTimerMessage = async (message: TimerMessage) => {
   const habitify = await useHabitify(await getHabitifyApiToken());
   const habit = (await habitify.getHabit(message.habit)).data;
   const project = await getProjectByName(habit?.area?.name ?? "");
@@ -68,7 +68,7 @@ async function processTimerMessage(message: TimerMessage) {
   );
   const toggl = await useToggl(await getTogglApiToken());
   return await toggl.startTimer(message.description, project?.id);
-}
+};
 
 browser.runtime.onMessage.addListener(async (message: Message) => {
   if (message.type === "verify_toggl") {
diff --git a/src/content.ts b/src/content.ts
index 384c930..38ec5a2 100644
--- a/src/content.ts
+++ b/src/content.ts
@@ -1,9 +1,9 @@
 import { debounce } from "ts-debounce";
 import { browser } from "webextension-polyfill-ts";
-import { wait } from "./utils";
+import { notNull, retry } from "./utils";
 import "./style.css";
 
-async function verifyTogglApiToken() {
+const verifyTogglApiToken = async () => {
   while (true) {
     const isVerified = await browser.runtime.sendMessage({
       type: "verify_toggl",
@@ -16,9 +16,9 @@ async function verifyTogglApiToken() {
       token,
     } as TokenTogglMessage);
   }
-}
+};
 
-async function verifyHabitifyApiToken() {
+const verifyHabitifyApiToken = async () => {
   while (true) {
     const isVerified = await browser.runtime.sendMessage({
       type: "verify_habitify",
@@ -31,42 +31,50 @@ async function verifyHabitifyApiToken() {
       token,
     } as TokenHabitifyMessage);
   }
-}
+};
 
-async function startTimer(description: string, habit: string) {
+const startTimer = async (description: string, habit: string) => {
   await browser.runtime.sendMessage({
     type: "timer",
     description,
     habit,
   } as TimerMessage);
-}
+};
 
-async function getRoot() {
-  return wait(() => document.querySelector("#root"));
-}
+const getRoot = async () => {
+  return retry(() => document.querySelector<HTMLElement>("#root"));
+};
 
-async function getHabitContainers() {
-  const todo = await wait(() =>
+const getMainContainer = async () => {
+  return retry(() =>
+    document.querySelector<HTMLElement>(
+      "#root > div > div.css-76h34y > div.css-y3isu0 > div.css-1mwek1r"
+    )
+  );
+};
+
+const getHabitLists = async () => {
+  const todo = await retry(() =>
     document.querySelector<HTMLElement>(
       "#root > div > div.css-76h34y > div.css-y3isu0 > div.css-1mwek1r > div.css-0"
     )
   );
-  const done = await wait(() =>
+  const done = await retry(() =>
     document.querySelector<HTMLElement>(
       "#root > div > div.css-76h34y > div.css-y3isu0 > div.css-1mwek1r > div:nth-child(4) > div > div.chakra-collapse > div"
     )
   );
-  return [todo, done];
-}
+  return [todo, done].filter(notNull);
+};
 
-function createTogglButton(onclick: (e: MouseEvent) => void) {
+const createTogglButton = (onclick: (e: MouseEvent) => void) => {
   const button = document.createElement("div");
   button.className = "toggl-button";
   button.onclick = onclick;
   return button;
-}
+};
 
-function appendTogglButton($item: HTMLElement) {
+const appendTogglButton = ($item: HTMLElement) => {
   if ($item.querySelector(".toggl-button")) return;
   const $info = $item.querySelector(".item-habit-info");
   $info?.append(
@@ -84,35 +92,46 @@ function appendTogglButton($item: HTMLElement) {
       startTimer(description, habit);
     })
   );
-}
+};
 
-const appendTogglButtons = debounce(async ($containers: HTMLElement[]) => {
-  for (const $container of $containers) {
-    const $items = $container.querySelectorAll<HTMLElement>(":scope > div");
+const appendTogglButtons = debounce(async ($lists: HTMLElement[]) => {
+  for (const $list of $lists) {
+    const $items = $list.querySelectorAll<HTMLElement>(":scope > div");
     for (const $item of $items) {
       appendTogglButton($item);
     }
   }
 });
 
-const setupObserverOnContainers = debounce(async () => {
-  const $containers = await getHabitContainers();
-  const observer = new MutationObserver(() => appendTogglButtons($containers));
-  for (const $container of $containers) {
-    observer.observe($container, { childList: true });
+const setupHabitListObserver = debounce(async () => {
+  const $lists = await getHabitLists();
+  const observer = new MutationObserver(() => appendTogglButtons($lists));
+  for (const $list of $lists) {
+    observer.observe($list, { childList: true });
   }
-  appendTogglButtons($containers);
+  appendTogglButtons($lists);
 });
 
-async function initialize() {
+const setupMainContainerObserver = debounce(async () => {
+  const $container = await getMainContainer();
+  if (!$container) return;
+  const observer = new MutationObserver(setupHabitListObserver);
+  observer.observe($container, { childList: true });
+  setupHabitListObserver();
+});
+
+const setupRootObserver = debounce(async () => {
   const $root = await getRoot();
-  const observer = new MutationObserver(setupObserverOnContainers);
+  if (!$root) return;
+  const observer = new MutationObserver(setupHabitListObserver);
   observer.observe($root, { childList: true });
-  setupObserverOnContainers();
-}
+  setupMainContainerObserver();
+});
 
-(async () => {
+const initialize = async () => {
   await verifyTogglApiToken();
   await verifyHabitifyApiToken();
-  await initialize();
-})();
+  await setupRootObserver();
+};
+
+void initialize();
diff --git a/src/habitify.ts b/src/habitify.ts
index 16f4545..4ee16a0 100644
--- a/src/habitify.ts
+++ b/src/habitify.ts
@@ -2,16 +2,16 @@ import { Response, Area, Habit } from "habitify";
 
 const HABITIFY_API_URL = "https://api.habitify.me";
 
-export async function verifyHabitifyApiToken(token: string) {
+export const verifyHabitifyApiToken = async (token: string) => {
   const result = await fetch(`${HABITIFY_API_URL}/areas`, {
     headers: {
       Authorization: token,
     },
   });
   return result.status === 200;
-}
+};
 
-export function useHabitify(token: string) {
+export const useHabitify = (token: string) => {
   const headers = {
     Authorization: token,
   };
@@ -29,4 +29,4 @@ export function useHabitify(token: string) {
     getAreas,
     getHabit,
   };
-}
+};
diff --git a/src/toggl.ts b/src/toggl.ts
index 31457cb..e3f8c46 100644
--- a/src/toggl.ts
+++ b/src/toggl.ts
@@ -2,16 +2,16 @@ import { Response, Project, TimeEntry, Workspace } from "toggl";
 
 const TOGGL_API_URL = "https://api.track.toggl.com/api/v8";
 
-export async function verifyTogglApiToken(token: string) {
+export const verifyTogglApiToken = async (token: string) => {
   const result = await fetch(`${TOGGL_API_URL}/me`, {
     headers: {
       Authorization: `Basic ${btoa(`${token}:api_token`)}`,
     },
   });
   return result.status === 200;
-}
+};
 
-export function useToggl(token: string) {
+export const useToggl = (token: string) => {
   const headers = {
     Authorization: `Basic ${btoa(`${token}:api_token`)}`,
   };
@@ -60,4 +60,4 @@ export function useToggl(token: string) {
     getWorkspaceProjects,
     startTimer,
   };
-}
+};
diff --git a/src/utils.ts b/src/utils.ts
index 0d1ffb8..24c84bd 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -1,13 +1,15 @@
 export const sleep = (ms: number) =>
   new Promise((resolve) => setTimeout(resolve, ms));
 
-export const wait = async <T>(
-  callback: () => T | undefined | null,
-  ms: number = 100
-): Promise<T> => {
-  while (true) {
-    const ret = callback();
-    if (ret) return ret;
-    await sleep(ms);
-  }
+export const retry = async <T>(
+  fn: () => T | undefined | null,
+  ms: number = 100,
+  count: number = 10
+): Promise<T | null> => {
+  if (count === 0) return null;
+  await sleep(ms);
+  return (await fn()) ?? retry(fn, ms, count - 1);
 };
+
+export const notNull = <T>(value: T | null | undefined): value is T =>
+  value !== null && value !== undefined;