Support Supabase branches (#1394)

drizzle/0015_complete_old_lace.sql

+ALTER TABLE `apps` ADD `supabase_parent_project_id` text;
\ No newline at end of file

drizzle/meta/_journal.json

@@ -106,6 +106,13 @@
       "when": 1760034009367,
       "tag": "0014_needy_vertigo",
       "breakpoints": true
+    },
+    {
+      "idx": 15,
+      "version": "6",
+      "when": 1760474402750,
+      "tag": "0015_complete_old_lace",
+      "breakpoints": true
     }
   ]
 }
\ No newline at end of file

e2e-tests/snapshots/supabase_branch.spec.ts_supabase-branch-selection-works-1.aria.yml

+- text: "/Tokens: \\d+,\\d+ \\d+% of [\\d,.]+[bkmBKM]+ Optimize your tokens with Dyad Pro's Smart Context/"
\ No newline at end of file

e2e-tests/supabase_branch.spec.ts

+import { testSkipIfWindows } from "./helpers/test_helper";
+import { expect } from "@playwright/test";
+
+testSkipIfWindows("supabase branch selection works", async ({ po }) => {
+  await po.setUp({ autoApprove: true });
+  await po.importApp("minimal");
+  await po.sendPrompt("tc=add-supabase");
+
+  // Connect to Supabase
+  await po.page.getByText("Set up supabase").click();
+  await po.clickConnectSupabaseButton();
+  await po.clickBackButton();
+  await po.page.getByTestId("token-bar-toggle").click();
+  // The default branch has a small context.
+  await expect(po.page.getByTestId("token-bar")).toContainText("6% of 128K");
+
+  await po.getTitleBarAppNameButton().click();
+  await po.page.getByTestId("supabase-branch-select").click();
+  await po.page.getByRole("option", { name: "Test Branch" }).click();
+
+  await po.clickBackButton();
+  // The test branch has a large context (200k tokens) so it'll hit the 100% limit.
+  // This is to make sure we're connecting to the right supabase project for the branch.
+  await expect(po.page.getByTestId("token-bar")).toContainText("100% of 128K");
+});

src/atoms/supabaseAtoms.ts

 import { atom } from "jotai";
+import { SupabaseBranch } from "@/ipc/ipc_types";
 
 // Define atom for storing the list of Supabase projects
 export const supabaseProjectsAtom = atom<any[]>([]);
+export const supabaseBranchesAtom = atom<SupabaseBranch[]>([]);
 
 // Define atom for tracking loading state
 export const supabaseLoadingAtom = atom<boolean>(false);

src/components/SupabaseConnector.tsx

@@ -56,6 +56,8 @@ export function SupabaseConnector({ appId }: { appId: number }) {
     loading,
     error,
     loadProjects,
+    branches,
+    loadBranches,
     setAppProject,
     unsetAppProject,
   } = useSupabase();
@@ -70,7 +72,7 @@ export function SupabaseConnector({ appId }: { appId: number }) {
 
   const handleProjectSelect = async (projectId: string) => {
     try {
-      await setAppProject(projectId, appId);
+      await setAppProject({ projectId, appId });
       toast.success("Project connected to app successfully");
       await refreshApp();
     } catch (error) {
@@ -78,6 +80,14 @@ export function SupabaseConnector({ appId }: { appId: number }) {
     }
   };
 
+  const projectIdForBranches =
+    app?.supabaseParentProjectId || app?.supabaseProjectId;
+  useEffect(() => {
+    if (projectIdForBranches) {
+      loadBranches(projectIdForBranches);
+    }
+  }, [projectIdForBranches, loadBranches]);
+
   const handleUnsetProject = async () => {
     try {
       await unsetAppProject(appId);
@@ -122,9 +132,56 @@ export function SupabaseConnector({ appId }: { appId: number }) {
             </CardDescription>
           </CardHeader>
           <CardContent>
-            <Button variant="destructive" onClick={handleUnsetProject}>
-              Disconnect Project
-            </Button>
+            <div className="space-y-4">
+              <div className="space-y-2">
+                <Label htmlFor="supabase-branch-select">Database Branch</Label>
+                <Select
+                  value={app.supabaseProjectId || ""}
+                  onValueChange={async (supabaseBranchProjectId) => {
+                    try {
+                      const branch = branches.find(
+                        (b) => b.projectRef === supabaseBranchProjectId,
+                      );
+                      if (!branch) {
+                        throw new Error("Branch not found");
+                      }
+                      await setAppProject({
+                        projectId: branch.projectRef,
+                        parentProjectId: branch.parentProjectRef,
+                        appId,
+                      });
+                      toast.success("Branch selected");
+                      await refreshApp();
+                    } catch (error) {
+                      toast.error("Failed to set branch: " + error);
+                    }
+                  }}
+                  disabled={loading}
+                >
+                  <SelectTrigger
+                    id="supabase-branch-select"
+                    data-testid="supabase-branch-select"
+                  >
+                    <SelectValue placeholder="Select a branch" />
+                  </SelectTrigger>
+                  <SelectContent>
+                    {branches.map((branch) => (
+                      <SelectItem
+                        key={branch.projectRef}
+                        value={branch.projectRef}
+                      >
+                        {branch.name}
+                        {branch.isDefault && " (Default)"}
+                      </SelectItem>
+                    ))}
+                  </SelectContent>
+                </Select>
+              </div>
+
+              <Button variant="destructive" onClick={handleUnsetProject}>
+                Disconnect Project
+              </Button>
+            </div>
           </CardContent>
         </Card>
       );

src/components/chat/ChatInput.tsx

@@ -351,6 +351,7 @@ export function ChatInput({ chatId }: { chatId?: number }) {
                       showTokenBar ? "text-purple-500 bg-purple-100" : ""
                     }`}
                     size="sm"
+                    data-testid="token-bar-toggle"
                   >
                     <ChartColumnIncreasing size={14} />
                   </Button>

src/components/chat/TokenBar.tsx

@@ -67,7 +67,7 @@ export function TokenBar({ chatId }: TokenBarProps) {
   const inputPercent = (inputTokens / contextWindow) * 100;
 
   return (
-    <div className="px-4 pb-2 text-xs">
+    <div className="px-4 pb-2 text-xs" data-testid="token-bar">
       <TooltipProvider>
         <Tooltip>
           <TooltipTrigger asChild>

src/db/schema.ts

@@ -29,6 +29,12 @@ export const apps = sqliteTable("apps", {
   githubRepo: text("github_repo"),
   githubBranch: text("github_branch"),
   supabaseProjectId: text("supabase_project_id"),
+  // If supabaseProjectId is a branch, then the parent project id set.
+  // This is because there's no way to retrieve ALL the branches for ALL projects
+  // in a single API call
+  // This is only used for display purposes but is NOT used for any actual
+  // supabase management logic.
+  supabaseParentProjectId: text("supabase_parent_project_id"),
   neonProjectId: text("neon_project_id"),
   neonDevelopmentBranchId: text("neon_development_branch_id"),
   neonPreviewBranchId: text("neon_preview_branch_id"),

src/hooks/useSupabase.ts

@@ -2,14 +2,17 @@ import { useCallback } from "react";
 import { useAtom } from "jotai";
 import {
   supabaseProjectsAtom,
+  supabaseBranchesAtom,
   supabaseLoadingAtom,
   supabaseErrorAtom,
   selectedSupabaseProjectAtom,
 } from "@/atoms/supabaseAtoms";
 import { IpcClient } from "@/ipc/ipc_client";
+import { SetSupabaseAppProjectParams } from "@/ipc/ipc_types";
 
 export function useSupabase() {
   const [projects, setProjects] = useAtom(supabaseProjectsAtom);
+  const [branches, setBranches] = useAtom(supabaseBranchesAtom);
   const [loading, setLoading] = useAtom(supabaseLoadingAtom);
   const [error, setError] = useAtom(supabaseErrorAtom);
   const [selectedProject, setSelectedProject] = useAtom(
@@ -35,14 +38,34 @@ export function useSupabase() {
     }
   }, [ipcClient, setProjects, setError, setLoading]);
 
+  /**
+   * Load branches for a Supabase project
+   */
+  const loadBranches = useCallback(
+    async (projectId: string) => {
+      setLoading(true);
+      try {
+        const list = await ipcClient.listSupabaseBranches({ projectId });
+        setBranches(Array.isArray(list) ? list : []);
+        setError(null);
+      } catch (error) {
+        console.error("Error loading Supabase branches:", error);
+        setError(error instanceof Error ? error : new Error(String(error)));
+      } finally {
+        setLoading(false);
+      }
+    },
+    [ipcClient, setBranches, setError, setLoading],
+  );
+
   /**
    * Associate a Supabase project with an app
    */
   const setAppProject = useCallback(
-    async (projectId: string, appId: number) => {
+    async (params: SetSupabaseAppProjectParams) => {
       setLoading(true);
       try {
-        await ipcClient.setSupabaseAppProject(projectId, appId);
+        await ipcClient.setSupabaseAppProject(params);
         setError(null);
       } catch (error) {
         console.error("Error setting Supabase project for app:", error);
@@ -87,10 +110,12 @@ export function useSupabase() {
 
   return {
     projects,
+    branches,
     loading,
     error,
     selectedProject,
     loadProjects,
+    loadBranches,
     setAppProject,
     unsetAppProject,
     selectProject,

src/ipc/handlers/app_handlers.ts

@@ -724,7 +724,9 @@ export function registerAppHandlers() {
     let supabaseProjectName: string | null = null;
     const settings = readSettings();
     if (app.supabaseProjectId && settings.supabase?.accessToken?.value) {
-      supabaseProjectName = await getSupabaseProjectName(app.supabaseProjectId);
+      supabaseProjectName = await getSupabaseProjectName(
+        app.supabaseParentProjectId || app.supabaseProjectId,
+      );
     }
 
     let vercelTeamSlug: string | null = null;

src/ipc/handlers/supabase_handlers.ts

@@ -2,7 +2,10 @@ import log from "electron-log";
 import { db } from "../../db";
 import { eq } from "drizzle-orm";
 import { apps } from "../../db/schema";
-import { getSupabaseClient } from "../../supabase_admin/supabase_management_client";
+import {
+  getSupabaseClient,
+  listSupabaseBranches,
+} from "../../supabase_admin/supabase_management_client";
 import {
   createLoggedHandler,
   createTestOnlyLoggedHandler,
@@ -10,6 +13,8 @@ import {
 import { handleSupabaseOAuthReturn } from "../../supabase_admin/supabase_return_handler";
 import { safeSend } from "../utils/safe_sender";
 
+import { SetSupabaseAppProjectParams, SupabaseBranch } from "../ipc_types";
+
 const logger = log.scope("supabase_handlers");
 const handle = createLoggedHandler(logger);
 const testOnlyHandle = createTestOnlyLoggedHandler(logger);
@@ -20,16 +25,44 @@ export function registerSupabaseHandlers() {
     return supabase.getProjects();
   });
 
+  // List branches for a Supabase project (database branches)
+  handle(
+    "supabase:list-branches",
+    async (
+      _,
+      { projectId }: { projectId: string },
+    ): Promise<Array<SupabaseBranch>> => {
+      const branches = await listSupabaseBranches({
+        supabaseProjectId: projectId,
+      });
+      return branches.map((branch) => ({
+        id: branch.id,
+        name: branch.name,
+        isDefault: branch.is_default,
+        projectRef: branch.project_ref,
+        parentProjectRef: branch.parent_project_ref,
+      }));
+    },
+  );
+
   // Set app project - links a Dyad app to a Supabase project
   handle(
     "supabase:set-app-project",
-    async (_, { project, app }: { project: string; app: number }) => {
+    async (
+      _,
+      { projectId, appId, parentProjectId }: SetSupabaseAppProjectParams,
+    ) => {
       await db
         .update(apps)
-        .set({ supabaseProjectId: project })
-        .where(eq(apps.id, app));
+        .set({
+          supabaseProjectId: projectId,
+          supabaseParentProjectId: parentProjectId,
+        })
+        .where(eq(apps.id, appId));
 
-      logger.info(`Associated app ${app} with Supabase project ${project}`);
+      logger.info(
+        `Associated app ${appId} with Supabase project ${projectId} ${parentProjectId ? `and parent project ${parentProjectId}` : ""}`,
+      );
     },
   );
 
@@ -37,7 +70,7 @@ export function registerSupabaseHandlers() {
   handle("supabase:unset-app-project", async (_, { app }: { app: number }) => {
     await db
       .update(apps)
-      .set({ supabaseProjectId: null })
+      .set({ supabaseProjectId: null, supabaseParentProjectId: null })
       .where(eq(apps.id, app));
 
     logger.info(`Removed Supabase project association for app ${app}`);

src/ipc/ipc_client.ts

@@ -66,6 +66,8 @@ import type {
   McpServerUpdate,
   CreateMcpServer,
   CloneRepoParams,
+  SupabaseBranch,
+  SetSupabaseAppProjectParams,
 } from "./ipc_types";
 import type { Template } from "../shared/templates";
 import type {
@@ -961,14 +963,16 @@ export class IpcClient {
     return this.ipcRenderer.invoke("supabase:list-projects");
   }
 
+  public async listSupabaseBranches(params: {
+    projectId: string;
+  }): Promise<SupabaseBranch[]> {
+    return this.ipcRenderer.invoke("supabase:list-branches", params);
+  }
+
   public async setSupabaseAppProject(
-    project: string,
-    app: number,
+    params: SetSupabaseAppProjectParams,
   ): Promise<void> {
-    await this.ipcRenderer.invoke("supabase:set-app-project", {
-      project,
-      app,
-    });
+    await this.ipcRenderer.invoke("supabase:set-app-project", params);
   }
 
   public async unsetSupabaseAppProject(app: number): Promise<void> {

src/ipc/ipc_types.ts

@@ -90,6 +90,7 @@ export interface App {
   githubRepo: string | null;
   githubBranch: string | null;
   supabaseProjectId: string | null;
+  supabaseParentProjectId: string | null;
   supabaseProjectName: string | null;
   neonProjectId: string | null;
   neonDevelopmentBranchId: string | null;
@@ -508,3 +509,17 @@ export type CloneRepoReturnType =
   | {
       error: string;
     };
+
+export interface SupabaseBranch {
+  id: string;
+  name: string;
+  isDefault: boolean;
+  projectRef: string;
+  parentProjectRef: string;
+}
+
+export interface SetSupabaseAppProjectParams {
+  projectId: string;
+  parentProjectId?: string;
+  appId: number;
+}

src/preload.ts

@@ -76,6 +76,7 @@ const validInvokeChannels = [
   "reject-proposal",
   "get-system-debug-info",
   "supabase:list-projects",
+  "supabase:list-branches",
   "supabase:set-app-project",
   "supabase:unset-app-project",
   "local-models:list-ollama",

src/supabase_admin/supabase_context.ts

@@ -56,6 +56,9 @@ export async function getSupabaseContext({
   supabaseProjectId: string;
 }) {
   if (IS_TEST_BUILD) {
+    if (supabaseProjectId === "test-branch-project-id") {
+      return "1234".repeat(200_000);
+    }
     return "[[TEST_BUILD_SUPABASE_CONTEXT]]";
   }
 

src/supabase_admin/supabase_management_client.ts

@@ -168,6 +168,61 @@ export async function deleteSupabaseFunction({
   );
 }
 
+export async function listSupabaseBranches({
+  supabaseProjectId,
+}: {
+  supabaseProjectId: string;
+}): Promise<
+  Array<{
+    id: string;
+    name: string;
+    is_default: boolean;
+    project_ref: string;
+    parent_project_ref: string;
+  }>
+> {
+  if (IS_TEST_BUILD) {
+    return [
+      {
+        id: "default-branch-id",
+        name: "Default Branch",
+        is_default: true,
+        project_ref: "fake-project-id",
+        parent_project_ref: "fake-project-id",
+      },
+
+      {
+        id: "test-branch-id",
+        name: "Test Branch",
+        is_default: false,
+        project_ref: "test-branch-project-id",
+        parent_project_ref: "fake-project-id",
+      },
+    ];
+  }
+
+  logger.info(`Listing Supabase branches for project: ${supabaseProjectId}`);
+  const supabase = await getSupabaseClient();
+
+  const response = await fetch(
+    `https://api.supabase.com/v1/projects/${supabaseProjectId}/branches`,
+    {
+      method: "GET",
+      headers: {
+        Authorization: `Bearer ${(supabase as any).options.accessToken}`,
+      },
+    },
+  );
+
+  if (response.status !== 200) {
+    throw await createResponseError(response, "list branches");
+  }
+
+  logger.info(`Listed Supabase branches for project: ${supabaseProjectId}`);
+  const jsonResponse = await response.json();
+  return jsonResponse;
+}
+
 export async function deploySupabaseFunctions({
   supabaseProjectId,
   functionName,