diff --git a/ui/src/components/InlineEntitySelector.tsx b/ui/src/components/InlineEntitySelector.tsx
index cae1897712e..0ac378b5fee 100644
--- a/ui/src/components/InlineEntitySelector.tsx
+++ b/ui/src/components/InlineEntitySelector.tsx
@@ -1,4 +1,4 @@
-import { forwardRef, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
+import { forwardRef, useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
 import { Check } from "lucide-react";
 import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
 import { orderItemsBySelectedAndRecent } from "../lib/recent-selections";
@@ -75,17 +75,17 @@ export const InlineEntitySelector = forwardRef<HTMLButtonElement, InlineEntitySe
 
     const currentOption = options.find((option) => option.id === value) ?? null;
 
-    function setHighlightedIndexValue(next: number | ((current: number) => number)) {
+    const setHighlightedIndexValue = useCallback((next: number | ((current: number) => number)) => {
       const resolved = typeof next === "function" ? next(highlightedIndexRef.current) : next;
       highlightedIndexRef.current = resolved;
       setHighlightedIndex(resolved);
-    }
+    }, []);
 
     useEffect(() => {
       if (!open) return;
       const selectedIndex = filteredOptions.findIndex((option) => option.id === value);
       setHighlightedIndexValue(selectedIndex >= 0 ? selectedIndex : 0);
-    }, [filteredOptions, open, value]);
+    }, [filteredOptions, open, setHighlightedIndexValue, value]);
 
     const commitSelection = (index: number, moveNext: boolean) => {
       const option = filteredOptions[index] ?? filteredOptions[0];
diff --git a/ui/src/context/CompanyContext.test.tsx b/ui/src/context/CompanyContext.test.tsx
index 0e26a5b82cf..d31b6e33d2e 100644
--- a/ui/src/context/CompanyContext.test.tsx
+++ b/ui/src/context/CompanyContext.test.tsx
@@ -5,6 +5,7 @@ import { createRoot, type Root } from "react-dom/client";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { Company } from "@paperclipai/shared";
+import { queryKeys } from "../lib/queryKeys";
 import {
   CompanyProvider,
   resolveBootstrapCompanySelection,
@@ -141,8 +142,8 @@ describe("CompanyProvider", () => {
     });
   });
 
-  afterEach(() => {
-    act(() => {
+  afterEach(async () => {
+    await act(async () => {
       root.unmount();
     });
     queryClient.clear();
@@ -170,11 +171,11 @@ describe("CompanyProvider", () => {
 
   it("replaces a stale stored company id with the first loaded company", async () => {
     localStorage.setItem("paperclip.selectedCompanyId", "stale-company");
-    let resolveCompanies!: (companies: Company[]) => void;
-    const companiesPromise = new Promise<Company[]>((resolve) => {
-      resolveCompanies = resolve;
+    queryClient.setQueryData(queryKeys.companies.all, {
+      companies: [makeCompany("company-1")],
+      unauthorized: false,
     });
-    mockCompaniesApi.list.mockReturnValue(companiesPromise);
+    mockCompaniesApi.list.mockImplementation(() => new Promise(() => {}));
     const seen: Array<string | null> = [];
 
     await act(async () => {
@@ -187,14 +188,7 @@ describe("CompanyProvider", () => {
       );
     });
 
-    await act(async () => {
-      resolveCompanies([makeCompany("company-1")]);
-      await companiesPromise;
-    });
-
-    await vi.waitFor(() => {
-      expect(seen).toEqual([null, "company-1"]);
-      expect(localStorage.getItem("paperclip.selectedCompanyId")).toBe("company-1");
-    });
+    expect(seen).toEqual([null, "company-1"]);
+    expect(localStorage.getItem("paperclip.selectedCompanyId")).toBe("company-1");
   });
 });