Agent: DB tools (#2122)

src/components/chat/DyadMarkdownParser.tsx

@@ -29,6 +29,8 @@ import { DyadCodeSearch } from "./DyadCodeSearch";
 import { DyadRead } from "./DyadRead";
 import { DyadListFiles } from "./DyadListFiles";
 import { DyadDatabaseSchema } from "./DyadDatabaseSchema";
+import { DyadSupabaseTableSchema } from "./DyadSupabaseTableSchema";
+import { DyadSupabaseProjectInfo } from "./DyadSupabaseProjectInfo";
 import { mapActionToButton } from "./ChatInput";
 import { SuggestedAction } from "@/lib/schemas";
 import { FixAllErrorsButton } from "./FixAllErrorsButton";
@@ -59,6 +61,8 @@ const DYAD_CUSTOM_TAGS = [
   "dyad-mcp-tool-result",
   "dyad-list-files",
   "dyad-database-schema",
+  "dyad-supabase-table-schema",
+  "dyad-supabase-project-info",
 ];
 
 interface DyadMarkdownParserProps {
@@ -634,6 +638,33 @@ function renderCustomTag(
         </DyadDatabaseSchema>
       );
 
+    case "dyad-supabase-table-schema":
+      return (
+        <DyadSupabaseTableSchema
+          node={{
+            properties: {
+              table: attributes.table || "",
+              state: getState({ isStreaming, inProgress }),
+            },
+          }}
+        >
+          {content}
+        </DyadSupabaseTableSchema>
+      );
+
+    case "dyad-supabase-project-info":
+      return (
+        <DyadSupabaseProjectInfo
+          node={{
+            properties: {
+              state: getState({ isStreaming, inProgress }),
+            },
+          }}
+        >
+          {content}
+        </DyadSupabaseProjectInfo>
+      );
+
     default:
       return null;
   }

src/components/chat/DyadSupabaseProjectInfo.tsx

+import React, { useState } from "react";
+import { CustomTagState } from "./stateTypes";
+import {
+  Database,
+  Loader2,
+  CircleX,
+  ChevronsDownUp,
+  ChevronsUpDown,
+} from "lucide-react";
+
+interface DyadSupabaseProjectInfoProps {
+  node: {
+    properties: {
+      state?: CustomTagState;
+    };
+  };
+  children: React.ReactNode;
+}
+
+export function DyadSupabaseProjectInfo({
+  node,
+  children,
+}: DyadSupabaseProjectInfoProps) {
+  const [isContentVisible, setIsContentVisible] = useState(false);
+  const { state } = node.properties;
+  const isLoading = state === "pending";
+  const isAborted = state === "aborted";
+  const content = typeof children === "string" ? children : "";
+
+  return (
+    <div
+      className={`bg-(--background-lightest) hover:bg-(--background-lighter) rounded-lg px-4 py-2 border my-2 cursor-pointer ${
+        isLoading
+          ? "border-amber-500"
+          : isAborted
+            ? "border-red-500"
+            : "border-border"
+      }`}
+      onClick={() => setIsContentVisible(!isContentVisible)}
+    >
+      <div className="flex items-center justify-between">
+        <div className="flex items-center gap-2">
+          {isLoading ? (
+            <Loader2 className="size-4 animate-spin text-amber-600" />
+          ) : isAborted ? (
+            <CircleX className="size-4 text-red-500" />
+          ) : (
+            <Database className="size-4 text-muted-foreground" />
+          )}
+          <span className="text-gray-700 dark:text-gray-300 font-medium text-sm">
+            Supabase Project Info
+          </span>
+          {isLoading && (
+            <span className="text-xs text-amber-600">Fetching...</span>
+          )}
+          {isAborted && (
+            <span className="text-xs text-red-500">Did not finish</span>
+          )}
+        </div>
+        <div className="flex items-center">
+          {isContentVisible ? (
+            <ChevronsDownUp
+              size={20}
+              className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
+            />
+          ) : (
+            <ChevronsUpDown
+              size={20}
+              className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
+            />
+          )}
+        </div>
+      </div>
+      {isContentVisible && content && (
+        <div className="mt-2 p-3 text-xs font-mono whitespace-pre-wrap max-h-80 overflow-y-auto bg-muted/30 rounded-md">
+          {content}
+        </div>
+      )}
+    </div>
+  );
+}

src/components/chat/DyadSupabaseTableSchema.tsx

+import React, { useState } from "react";
+import { CustomTagState } from "./stateTypes";
+import {
+  Table2,
+  Loader2,
+  CircleX,
+  ChevronsDownUp,
+  ChevronsUpDown,
+} from "lucide-react";
+
+interface DyadSupabaseTableSchemaProps {
+  node: {
+    properties: {
+      table?: string;
+      state?: CustomTagState;
+    };
+  };
+  children: React.ReactNode;
+}
+
+export function DyadSupabaseTableSchema({
+  node,
+  children,
+}: DyadSupabaseTableSchemaProps) {
+  const [isContentVisible, setIsContentVisible] = useState(false);
+  const { table, state } = node.properties;
+  const isLoading = state === "pending";
+  const isAborted = state === "aborted";
+  const content = typeof children === "string" ? children : "";
+
+  return (
+    <div
+      className={`bg-(--background-lightest) hover:bg-(--background-lighter) rounded-lg px-4 py-2 border my-2 cursor-pointer ${
+        isLoading
+          ? "border-amber-500"
+          : isAborted
+            ? "border-red-500"
+            : "border-border"
+      }`}
+      onClick={() => setIsContentVisible(!isContentVisible)}
+    >
+      <div className="flex items-center justify-between">
+        <div className="flex items-center gap-2">
+          {isLoading ? (
+            <Loader2 className="size-4 animate-spin text-amber-600" />
+          ) : isAborted ? (
+            <CircleX className="size-4 text-red-500" />
+          ) : (
+            <Table2 className="size-4 text-muted-foreground" />
+          )}
+          <span className="text-gray-700 dark:text-gray-300 font-medium text-sm">
+            {table ? `Table Schema: ${table}` : "Supabase Table Schema"}
+          </span>
+          {isLoading && (
+            <span className="text-xs text-amber-600">Fetching...</span>
+          )}
+          {isAborted && (
+            <span className="text-xs text-red-500">Did not finish</span>
+          )}
+        </div>
+        <div className="flex items-center">
+          {isContentVisible ? (
+            <ChevronsDownUp
+              size={20}
+              className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
+            />
+          ) : (
+            <ChevronsUpDown
+              size={20}
+              className="text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
+            />
+          )}
+        </div>
+      </div>
+      {isContentVisible && content && (
+        <div className="mt-2 p-3 text-xs font-mono whitespace-pre-wrap max-h-80 overflow-y-auto bg-muted/30 rounded-md">
+          {content}
+        </div>
+      )}
+    </div>
+  );
+}

src/pro/main/ipc/handlers/local_agent/tool_definitions.ts

@@ -14,7 +14,8 @@ import { executeSqlTool } from "./tools/execute_sql";
 
 import { readFileTool } from "./tools/read_file";
 import { listFilesTool } from "./tools/list_files";
-import { getDatabaseSchemaTool } from "./tools/get_database_schema";
+import { getSupabaseProjectInfoTool } from "./tools/get_supabase_project_info";
+import { getSupabaseTableSchemaTool } from "./tools/get_supabase_table_schema";
 import { setChatSummaryTool } from "./tools/set_chat_summary";
 import { addIntegrationTool } from "./tools/add_integration";
 import { readLogsTool } from "./tools/read_logs";
@@ -43,7 +44,8 @@ export const TOOL_DEFINITIONS: readonly ToolDefinition[] = [
   // searchReplaceTool,
   readFileTool,
   listFilesTool,
-  getDatabaseSchemaTool,
+  getSupabaseProjectInfoTool,
+  getSupabaseTableSchemaTool,
   setChatSummaryTool,
   addIntegrationTool,
   readLogsTool,

src/pro/main/ipc/handlers/local_agent/tools/get_database_schema.ts

-import { z } from "zod";
-import { ToolDefinition, AgentContext } from "./types";
-import { getSupabaseContext } from "../../../../../../supabase_admin/supabase_context";
-
-const getDatabaseSchemaSchema = z.object({});
-
-const XML_TAG = "<dyad-database-schema></dyad-database-schema>";
-
-export const getDatabaseSchemaTool: ToolDefinition<
-  z.infer<typeof getDatabaseSchemaSchema>
-> = {
-  name: "get_database_schema",
-  description: "Fetch the database schema from Supabase",
-  inputSchema: getDatabaseSchemaSchema,
-  defaultConsent: "always",
-  isEnabled: (ctx) => !!ctx.supabaseProjectId,
-
-  getConsentPreview: () => "Get Supabase schema",
-
-  buildXml: (_args, _isComplete) => {
-    // This tool has no inputs, so always return the same XML
-    return XML_TAG;
-  },
-
-  execute: async (_args, ctx: AgentContext) => {
-    if (!ctx.supabaseProjectId) {
-      throw new Error("Supabase is not connected to this app");
-    }
-
-    const schema = await getSupabaseContext({
-      supabaseProjectId: ctx.supabaseProjectId,
-      organizationSlug: ctx.supabaseOrganizationSlug ?? null,
-    });
-
-    return schema || "";
-  },
-};

src/pro/main/ipc/handlers/local_agent/tools/get_supabase_project_info.ts

+import { z } from "zod";
+import { ToolDefinition, AgentContext, escapeXmlContent } from "./types";
+import { getSupabaseProjectInfo } from "../../../../../../supabase_admin/supabase_context";
+
+const getSupabaseProjectInfoSchema = z.object({
+  includeDbFunctions: z
+    .boolean()
+    .optional()
+    .describe(
+      "When true, includes database functions in the response. Defaults to false.",
+    ),
+});
+
+export const getSupabaseProjectInfoTool: ToolDefinition<
+  z.infer<typeof getSupabaseProjectInfoSchema>
+> = {
+  name: "get_supabase_project_info",
+  description:
+    "Get Supabase project overview: project ID, publishable key, secret names, and table names. Use this to discover what tables exist before fetching detailed schemas. Optionally include database functions.",
+  inputSchema: getSupabaseProjectInfoSchema,
+  defaultConsent: "always",
+  isEnabled: (ctx) => !!ctx.supabaseProjectId,
+
+  getConsentPreview: () => "Get Supabase project info",
+
+  execute: async (args, ctx: AgentContext) => {
+    if (!ctx.supabaseProjectId) {
+      throw new Error("Supabase is not connected to this app");
+    }
+
+    ctx.onXmlStream(
+      "<dyad-supabase-project-info></dyad-supabase-project-info>",
+    );
+
+    const info = await getSupabaseProjectInfo({
+      supabaseProjectId: ctx.supabaseProjectId,
+      organizationSlug: ctx.supabaseOrganizationSlug ?? null,
+      includeDbFunctions: args.includeDbFunctions,
+    });
+
+    ctx.onXmlComplete(
+      `<dyad-supabase-project-info>\n${escapeXmlContent(info)}\n</dyad-supabase-project-info>`,
+    );
+
+    return info;
+  },
+};

src/pro/main/ipc/handlers/local_agent/tools/get_supabase_table_schema.ts

+import { z } from "zod";
+import {
+  ToolDefinition,
+  AgentContext,
+  escapeXmlAttr,
+  escapeXmlContent,
+} from "./types";
+import { getSupabaseTableSchema } from "../../../../../../supabase_admin/supabase_context";
+
+const getSupabaseTableSchemaSchema = z.object({
+  tableName: z
+    .string()
+    .optional()
+    .describe(
+      "Optional table name to get schema for. If omitted, returns schema for all tables.",
+    ),
+});
+
+export const getSupabaseTableSchemaTool: ToolDefinition<
+  z.infer<typeof getSupabaseTableSchemaSchema>
+> = {
+  name: "get_supabase_table_schema",
+  description:
+    "Get database table schema from Supabase. If tableName is provided, returns schema for that specific table (columns, policies, triggers). If omitted, returns schema for all tables.",
+  inputSchema: getSupabaseTableSchemaSchema,
+  defaultConsent: "always",
+  isEnabled: (ctx) => !!ctx.supabaseProjectId,
+
+  getConsentPreview: (args) =>
+    args.tableName
+      ? `Get schema for table "${args.tableName}"`
+      : "Get schema for all tables",
+
+  execute: async (args, ctx: AgentContext) => {
+    if (!ctx.supabaseProjectId) {
+      throw new Error("Supabase is not connected to this app");
+    }
+
+    const tableAttr = args.tableName
+      ? ` table="${escapeXmlAttr(args.tableName)}"`
+      : "";
+    ctx.onXmlStream(
+      `<dyad-supabase-table-schema${tableAttr}></dyad-supabase-table-schema>`,
+    );
+
+    const schema = await getSupabaseTableSchema({
+      supabaseProjectId: ctx.supabaseProjectId,
+      organizationSlug: ctx.supabaseOrganizationSlug ?? null,
+      tableName: args.tableName,
+    });
+
+    ctx.onXmlComplete(
+      `<dyad-supabase-table-schema${tableAttr}>\n${escapeXmlContent(schema)}\n</dyad-supabase-table-schema>`,
+    );
+
+    return schema;
+  },
+};

src/supabase_admin/supabase_context.ts

 import { IS_TEST_BUILD } from "@/ipc/utils/test_utils";
 import { getSupabaseClient } from "./supabase_management_client";
-import { SUPABASE_SCHEMA_QUERY } from "./supabase_schema_query";
+import {
+  SUPABASE_SCHEMA_QUERY,
+  SUPABASE_FUNCTIONS_QUERY,
+  buildSupabaseSchemaQuery,
+} from "./supabase_schema_query";
 
 async function getPublishableKey({
   projectId,
@@ -108,3 +112,120 @@ export async function getSupabaseContext({
 
   return context;
 }
+
+// Query to get just table names (lightweight)
+const TABLE_NAMES_QUERY = `
+  SELECT table_name 
+  FROM information_schema.tables 
+  WHERE table_schema = 'public'
+  ORDER BY table_name;
+`;
+
+/**
+ * Get high-level Supabase project info: project ID, publishable key, secret names, and table names.
+ * This is a lightweight call that doesn't fetch full schema details.
+ * Optionally includes database functions when include_db_functions is true.
+ */
+export async function getSupabaseProjectInfo({
+  supabaseProjectId,
+  organizationSlug,
+  includeDbFunctions,
+}: {
+  supabaseProjectId: string;
+  organizationSlug: string | null;
+  includeDbFunctions?: boolean;
+}): Promise<string> {
+  if (IS_TEST_BUILD) {
+    let result = `# Supabase Project Info
+
+## Project ID
+${supabaseProjectId}
+
+## Publishable Key
+test-publishable-key
+
+## Secret Names
+["TEST_SECRET_1", "TEST_SECRET_2"]
+
+## Table Names
+["users", "posts", "comments"]
+`;
+    if (includeDbFunctions) {
+      result += `
+## Database Functions
+[{"name": "test_function", "arguments": "", "return_type": "void", "language": "plpgsql"}]
+`;
+    }
+    return result;
+  }
+
+  const supabase = await getSupabaseClient({ organizationSlug });
+  const publishableKey = await getPublishableKey({
+    projectId: supabaseProjectId,
+    organizationSlug,
+  });
+
+  const secrets = await supabase.getSecrets(supabaseProjectId);
+  const secretNames = secrets?.map((secret) => secret.name) ?? [];
+
+  const tableResult = await supabase.runQuery(
+    supabaseProjectId,
+    TABLE_NAMES_QUERY,
+  );
+  const tableNames =
+    (tableResult as unknown as { table_name: string }[] | undefined)?.map(
+      (row) => row.table_name,
+    ) ?? [];
+
+  let result = `# Supabase Project Info
+
+## Project ID
+${supabaseProjectId}
+
+## Publishable Key
+${publishableKey}
+
+## Secret Names
+${JSON.stringify(secretNames)}
+
+## Table Names
+${JSON.stringify(tableNames)}
+`;
+
+  if (includeDbFunctions) {
+    const functionsResult = await supabase.runQuery(
+      supabaseProjectId,
+      SUPABASE_FUNCTIONS_QUERY,
+    );
+    result += `
+## Database Functions
+${JSON.stringify(functionsResult)}
+`;
+  }
+
+  return result;
+}
+
+/**
+ * Get database table schema. If tableName is provided, returns schema for that specific table.
+ * If tableName is omitted, returns schema for all tables.
+ */
+export async function getSupabaseTableSchema({
+  supabaseProjectId,
+  organizationSlug,
+  tableName,
+}: {
+  supabaseProjectId: string;
+  organizationSlug: string | null;
+  tableName?: string;
+}): Promise<string> {
+  if (IS_TEST_BUILD) {
+    return `[[TEST_TABLE_SCHEMA${tableName ? `:${tableName}` : ""}]]`;
+  }
+
+  const supabase = await getSupabaseClient({ organizationSlug });
+  const query = buildSupabaseSchemaQuery(tableName);
+  const schemaResult = await supabase.runQuery(supabaseProjectId, query);
+
+  return JSON.stringify(schemaResult);
+}

src/supabase_admin/supabase_schema_query.ts

@@ -2,7 +2,27 @@
 // which is Apache 2.0 licensed and copyrighted to Jijun Leng
 // https://github.com/jjleng/code-panda/blob/61f1fa514c647de1a8d2ad7f85102d49c6db2086/LICENSE
 
-export const SUPABASE_SCHEMA_QUERY = `
+/**
+ * Build schema query with optional table name filter.
+ * When tableName is provided, only fetches schema for that specific table.
+ */
+export function buildSupabaseSchemaQuery(tableName?: string): string {
+  // Escape single quotes in table name to prevent SQL injection
+  const escapedTableName = tableName?.replace(/'/g, "''");
+  const tableFilter = escapedTableName
+    ? ` AND tables.table_name = '${escapedTableName}'`
+    : "";
+  const columnFilter = escapedTableName
+    ? ` AND c.table_name = '${escapedTableName}'`
+    : "";
+  const policyFilter = escapedTableName
+    ? ` AND cls.relname = '${escapedTableName}'`
+    : "";
+  const triggerFilter = escapedTableName
+    ? ` AND t.event_object_table = '${escapedTableName}'`
+    : "";
+
+  return `
         WITH table_info AS (
             SELECT
                 tables.table_name,
@@ -12,7 +32,7 @@ export const SUPABASE_SCHEMA_QUERY = `
             LEFT JOIN pg_stat_user_tables psut ON tables.table_name = psut.relname
             LEFT JOIN pg_class cls ON psut.relid = cls.oid
             LEFT JOIN pg_description pd ON psut.relid = pd.objoid AND pd.objsubid = 0
-            WHERE tables.table_schema = 'public'
+            WHERE tables.table_schema = 'public'${tableFilter}
         ),
         column_info AS (
             SELECT
@@ -26,7 +46,7 @@ export const SUPABASE_SCHEMA_QUERY = `
                     ) ORDER BY c.ordinal_position
                 ) as columns
             FROM information_schema.columns c
-            WHERE c.table_schema = 'public'
+            WHERE c.table_schema = 'public'${columnFilter}
             GROUP BY c.table_name
         ),
         tables_result AS (
@@ -60,9 +80,12 @@ export const SUPABASE_SCHEMA_QUERY = `
                 )::text as data
             FROM pg_policy pol
             JOIN pg_class cls ON pol.polrelid = cls.oid
-            WHERE cls.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')
+            WHERE cls.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')${policyFilter}
         ),
-        functions_result AS (
+        ${
+          tableName
+            ? ""
+            : `functions_result AS (
             SELECT
                 'functions' as result_type,
                 jsonb_build_object(
@@ -81,8 +104,9 @@ export const SUPABASE_SCHEMA_QUERY = `
             FROM pg_proc p
             LEFT JOIN pg_description d ON p.oid = d.objoid
             LEFT JOIN pg_language l ON p.prolang = l.oid
-            WHERE p.pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public') AND p.prokind = 'f' -- 'f' = normal function (otherwise source code fetch fails)
-        ),
+            WHERE p.pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public') AND p.prokind = 'f'
+        ),` // -- 'f' = normal function (otherwise source code fetch fails)
+        }
         triggers_result AS (
             SELECT
                 'triggers' as result_type,
@@ -97,14 +121,42 @@ export const SUPABASE_SCHEMA_QUERY = `
             FROM information_schema.triggers t
             LEFT JOIN pg_trigger pg_t ON t.trigger_name = pg_t.tgname
             LEFT JOIN pg_proc p ON pg_t.tgfoid = p.oid
-            WHERE t.trigger_schema = 'public'
+            WHERE t.trigger_schema = 'public'${triggerFilter}
         )
         SELECT result_type, data
         FROM (
             SELECT * FROM tables_result
             UNION ALL SELECT * FROM policies_result
-            UNION ALL SELECT * FROM functions_result
+            ${tableName ? "" : "UNION ALL SELECT * FROM functions_result"}
             UNION ALL SELECT * FROM triggers_result
         ) combined_results
         ORDER BY result_type;
 `;
+}
+
+export const SUPABASE_SCHEMA_QUERY = buildSupabaseSchemaQuery();
+
+/**
+ * Query to fetch only database functions from the public schema.
+ */
+export const SUPABASE_FUNCTIONS_QUERY = `
+  SELECT
+    jsonb_build_object(
+      'name', p.proname::text,
+      'description', d.description::text,
+      'arguments', pg_get_function_arguments(p.oid)::text,
+      'return_type', pg_get_function_result(p.oid)::text,
+      'language', l.lanname::text,
+      'volatility', CASE p.provolatile
+        WHEN 'i' THEN 'IMMUTABLE'
+        WHEN 's' THEN 'STABLE'
+        WHEN 'v' THEN 'VOLATILE'
+      END,
+      'source_code', pg_get_functiondef(p.oid)::text
+    )::text as data
+  FROM pg_proc p
+  LEFT JOIN pg_description d ON p.oid = d.objoid
+  LEFT JOIN pg_language l ON p.prolang = l.oid
+  WHERE p.pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')
+    AND p.prokind = 'f';
+`; // 'f' = normal function (otherwise source code fetch fails)