cobble/src/functions/backend/db/extensions/grab-extensions-overview.ts

import BunSQLite from "@moduletrace/bun-sqlite";
import type {
    BUN_SQLITE_COBBLE_CUSTOM_TOOLS,
    BUN_SQLITE_COBBLE_KNOWLEDGEBASES,
    BUN_SQLITE_COBBLE_MCP_SERVERS,
    BUN_SQLITE_COBBLE_SKILLS,
} from "@/db/types/db";
import { CustomToolKinds } from "@/src/dict/custom-tool-dict";
import { McpTransports } from "@/src/dict/mcp-dict";
import type { BotExtension, BotExtensionsPayload } from "@/src/types";
import parseKnowledgebaseRow from "../knowledgebases/parse-knowledgebase-row";
import parseCustomToolRow from "./parse-custom-tool-row";
import parseMcpServerRow from "./parse-mcp-server-row";
import parseSkillRow from "./parse-skill-row";

type Params = {
    /**
     * Whether the viewer can give and take them: only then is what can be
     * added listed.
     */
    can_manage: boolean;
    /**
     * A bot's: given to it, and (with `workspace_id`) that workspace's
     * too, not removable from the bot. Without it, only the workspace's.
     */
    bot_id?: number | string | null;
    workspace_id?: number | string | null;
};

type Linked<T> = T & { source: "bot" | "workspace" };

/**
 * The MCP servers, tools, skills and knowledgebases a bot or a workspace has, and the rest of
 * them, which can be added (`grabBotExtensionsOverview`,
 * `grabWorkspaceExtensionsOverview`).
 */
export default async function grabExtensionsOverview({
    can_manage,
    bot_id,
    workspace_id,
}: Params): Promise<BotExtensionsPayload> {
    const bot = Number(bot_id || 0);
    const workspace = Number(workspace_id || 0);

    const linked = (column: string, bot_link: string, workspace_link: string) => ({
        sql: `SELECT x.*, CASE WHEN x.id IN (SELECT ${column} FROM ${bot_link} WHERE bot_id = ?) THEN 'bot' ELSE 'workspace' END AS source
            FROM %TABLE% x WHERE x.id IN (SELECT ${column} FROM ${bot_link} WHERE bot_id = ?)
            OR x.id IN (SELECT ${column} FROM ${workspace_link} WHERE workspace_id = ?)
            ORDER BY x.name COLLATE NOCASE ASC`,
        values: [bot, bot, workspace],
    });

    const mcp_query = linked("mcp_server_id", "bot_mcp_servers", "workspace_mcp_servers");
    const tool_query = linked("custom_tool_id", "bot_custom_tools", "workspace_custom_tools");
    const skill_query = linked("skill_id", "bot_skills", "workspace_skills");
    const knowledgebase_query = linked("knowledgebase_id", "bot_knowledgebases", "workspace_knowledgebases");

    const [servers, tools, skills, knowledgebases, documents, all_servers, all_tools, all_skills, all_knowledgebases] =
        await Promise.all([
        BunSQLite.sql<Linked<BUN_SQLITE_COBBLE_MCP_SERVERS>>({
            sql: mcp_query.sql.replace("%TABLE%", "mcp_servers"),
            values: mcp_query.values,
        }),
        BunSQLite.sql<Linked<BUN_SQLITE_COBBLE_CUSTOM_TOOLS>>({
            sql: tool_query.sql.replace("%TABLE%", "custom_tools"),
            values: tool_query.values,
        }),
        BunSQLite.sql<Linked<BUN_SQLITE_COBBLE_SKILLS>>({
            sql: skill_query.sql.replace("%TABLE%", "skills"),
            values: skill_query.values,
        }),
        BunSQLite.sql<Linked<BUN_SQLITE_COBBLE_KNOWLEDGEBASES>>({
            sql: knowledgebase_query.sql.replace("%TABLE%", "knowledgebases"),
            values: knowledgebase_query.values,
        }),
        BunSQLite.sql<{ knowledgebase_id: number; documents: number; failed: number }>({
            sql: `SELECT knowledgebase_id, COUNT(*) AS documents,
                    SUM(CASE WHEN ingest_error IS NOT NULL AND COALESCE(chunk_count, 0) = 0 THEN 1 ELSE 0 END) AS failed
                FROM knowledgebase_documents GROUP BY knowledgebase_id`,
        }),
        can_manage
            ? BunSQLite.sql<{ id: number; name: string }>({ sql: `SELECT id, name FROM mcp_servers ORDER BY name COLLATE NOCASE` })
            : null,
        can_manage
            ? BunSQLite.sql<{ id: number; name: string }>({ sql: `SELECT id, name FROM custom_tools ORDER BY name` })
            : null,
        can_manage
            ? BunSQLite.sql<{ id: number; name: string }>({
                  sql: `SELECT id, name FROM skills ORDER BY name COLLATE NOCASE`,
              })
            : null,
        can_manage
            ? BunSQLite.sql<{ id: number; name: string }>({
                  sql: `SELECT id, name FROM knowledgebases ORDER BY name COLLATE NOCASE`,
              })
            : null,
    ]);

    const document_counts = new Map(
        (documents.payload || []).map((row) => [Number(row.knowledgebase_id), row]),
    );

    const attached: BotExtension[] = [
        ...(servers.payload || []).map((row) => {
            const server = parseMcpServerRow({ row });
            const transport = McpTransports.find((t) => t.value == server.transport)?.value == "stdio" ? "Command" : "Remote";
            const count = server.tools.filter((tool) => !server.excluded_tools.includes(tool.name)).length;

            return {
                kind: "mcp" as const,
                id: server.id,
                name: server.name,
                description: server.description,
                detail: `MCP · ${transport}${server.status == "ok" ? ` · ${count} ${count == 1 ? "tool" : "tools"}` : ""}`,
                state: server.status,
                enabled: server.enabled,
                source: row.source,
            };
        }),
        ...(tools.payload || []).map((row) => {
            const tool = parseCustomToolRow({ row });

            return {
                kind: "tool" as const,
                id: tool.id,
                name: tool.name,
                description: tool.description,
                detail: `Tool · ${CustomToolKinds.find((k) => k.value == tool.kind)?.title || tool.kind}`,
                state: tool.status,
                enabled: tool.enabled,
                source: row.source,
            };
        }),
        ...(skills.payload || []).map((row) => {
            const skill = parseSkillRow({ row });

            return {
                kind: "skill" as const,
                id: skill.id,
                name: skill.name,
                description: skill.description,
                detail: `Skill · ${skill.slug}`,
                state: "ok" as const,
                enabled: skill.enabled,
                source: row.source,
            };
        }),
        ...(knowledgebases.payload || []).map((row) => {
            const knowledgebase = parseKnowledgebaseRow({ row });
            const count = document_counts.get(knowledgebase.id);
            const total = Number(count?.documents) || 0;

            return {
                kind: "knowledgebase" as const,
                id: knowledgebase.id,
                name: knowledgebase.name,
                description: knowledgebase.description,
                detail: `Knowledgebase · ${total} ${total == 1 ? "document" : "documents"}`,
                state: Number(count?.failed) ? ("error" as const) : ("ok" as const),
                enabled: knowledgebase.enabled,
                source: row.source,
            };
        }),
    ];

    /**
     * A bot can be given one its workspace has; a workspace only what it
     * hasn't.
     */
    const has = (kind: BotExtension["kind"], id: number) =>
        attached.some((item) => item.kind == kind && item.id == id && (!bot || item.source == "bot"));

    return {
        attached,
        addable: [
            ...(all_servers?.payload || [])
                .filter((row) => !has("mcp", Number(row.id)))
                .map((row) => ({ kind: "mcp" as const, id: Number(row.id), name: row.name })),
            ...(all_tools?.payload || [])
                .filter((row) => !has("tool", Number(row.id)))
                .map((row) => ({ kind: "tool" as const, id: Number(row.id), name: row.name })),
            ...(all_skills?.payload || [])
                .filter((row) => !has("skill", Number(row.id)))
                .map((row) => ({ kind: "skill" as const, id: Number(row.id), name: row.name })),
            ...(all_knowledgebases?.payload || [])
                .filter((row) => !has("knowledgebase", Number(row.id)))
                .map((row) => ({ kind: "knowledgebase" as const, id: Number(row.id), name: row.name })),
        ],
        can_manage,
    };
}