koc-loop/lib/mcp-tools.ts

import { type McpServer } from "@modelcontextprotocol/server";
import { z } from "zod/v4";
import {
  backfillResourceProfile,
  collectNow,
  recoveryExport,
  recoveryList,
  resourceExport,
  resourceGet,
  resourceSearch,
  retryFailed,
  runDueCollections,
  setCollectionPlan,
  taskGet,
  taskList,
  type McpOperationBindings,
} from "./mcp-operations";

type Options = {
  bindings: McpOperationBindings;
  origin: string;
};

function result(label: string, data: unknown) {
  return {
    content: [{ type: "text" as const, text: label }],
    structuredContent: data as Record<string, unknown>,
  };
}

function errorResult(error: unknown) {
  const message = error instanceof Error ? error.message : "操作失败，请稍后重试";
  return {
    isError: true as const,
    content: [{ type: "text" as const, text: message.slice(0, 240) }],
  };
}

function withError<T extends unknown[]>(handler: (...args: T) => Promise<ReturnType<typeof result>>) {
  return async (...args: T) => {
    try {
      return await handler(...args);
    } catch (error) {
      return errorResult(error);
    }
  };
}

const pagination = {
  limit: z.number().int().min(1).max(200).optional().describe("返回数量，默认50，最大200"),
  offset: z.number().int().min(0).optional().describe("分页偏移量，默认0"),
};

const resourceFilters = {
  query: z.string().max(100).optional().describe("账号名称、账号号、简介或标签，支持模糊搜索"),
  ip_location: z.string().max(100).optional().describe("IP地区关键词，支持模糊搜索"),
  cooperation_source: z.string().max(100).optional().describe("历史合作来源关键词，支持模糊搜索"),
  platform: z.string().max(30).optional().describe("平台，例如小红书；不传表示全部"),
};

export function registerMcpOperationTools(server: McpServer, options: Options) {
  server.registerTool(
    "task_list",
    {
      title: "查询任务及进度",
      description: "查询 KOC 分发任务列表、领取数、发布数、第7天回收数及领取链接。",
      inputSchema: z.object({
        query: z.string().max(100).optional().describe("任务名或品牌/项目关键词"),
        status: z.string().max(30).optional().describe("任务状态；不传或 all 表示全部"),
        ...pagination,
      }),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: true },
    },
    withError(async (input) => {
      const data = await taskList(input, options.bindings.KOC_PORTAL_URL?.trim() || "");
      return result(`共找到 ${data.total} 个任务。`, data);
    }),
  );

  server.registerTool(
    "task_get",
    {
      title: "查看任务完整情况",
      description: "查看指定任务、笔记、领取记录、发布回填和采集执行记录。",
      inputSchema: z.object({ task_id: z.string().min(1).describe("任务ID") }),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: true },
    },
    withError(async ({ task_id }) => {
      const data = await taskGet(task_id, options.bindings.KOC_PORTAL_URL?.trim() || "");
      return result(`已读取任务“${String(data.task.name)}”的完整情况。`, data);
    }),
  );

  server.registerTool(
    "recovery_list",
    {
      title: "查询数据回收队列",
      description: "查询已发布、未回填、待第7天数据或已到第7天仍未回填的笔记。",
      inputSchema: z.object({
        task_id: z.string().optional().describe("任务ID；不传则跨任务查询"),
        stage: z.enum(["all", "published", "unfilled", "waiting_day7", "day7_due"]).optional(),
        ...pagination,
      }),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: true },
    },
    withError(async ({ task_id, stage, limit, offset }) => {
      const data = await recoveryList({ taskId: task_id, stage, limit, offset });
      return result(`数据回收队列共 ${data.total} 条记录。`, data);
    }),
  );

  server.registerTool(
    "recovery_export",
    {
      title: "导出任务完整数据",
      description: "生成任务完整 Excel，包含笔记原图、发布截图、创作者截图和回收数据。",
      inputSchema: z.object({ task_id: z.string().min(1).describe("任务ID") }),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: false },
    },
    withError(async ({ task_id }) => {
      const data = await recoveryExport(task_id, options.origin);
      return result(`导出文件已生成，下载链接将在 ${data.expires_at} 失效。`, data);
    }),
  );

  server.registerTool(
    "collection_plan_set",
    {
      title: "设置自动采集计划",
      description: "为任务设置开始日期及第1至第7天的自动采集日，系统在北京时间09:00执行。",
      inputSchema: z.object({
        task_id: z.string().min(1),
        start_date: z.string().regex(/^\d{4}-\d{2}-\d{2}$/).describe("北京时间开始日期 YYYY-MM-DD"),
        days: z.array(z.number().int().min(1).max(7)).min(1).max(7).describe("需要采集的相对天数，例如 [2,5,7]"),
      }),
      annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: true },
    },
    withError(async ({ task_id, start_date, days }) => {
      const data = await setCollectionPlan(task_id, start_date, days, options.bindings);
      return result(`已为任务设置 ${data.days.length} 个自动采集日。`, data);
    }),
  );

  server.registerTool(
    "collection_run_due",
    {
      title: "执行到期采集",
      description: "立即执行今天或此前已到期但尚未成功的自动采集任务。",
      inputSchema: z.object({ task_id: z.string().optional().describe("任务ID；不传则执行所有到期任务") }),
      annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: true },
    },
    withError(async ({ task_id }) => {
      const data = await runDueCollections(task_id, options.bindings);
      return result(`到期采集完成：成功 ${data.succeeded}，失败 ${data.failed}。`, data);
    }),
  );

  server.registerTool(
    "collection_collect_now",
    {
      title: "立即采集指定笔记",
      description: "对指定分发记录立即采集互动数据；小红书为点赞/收藏/评论，抖音另含转发。",
      inputSchema: z.object({
        distribution_id: z.string().min(1).describe("分发记录ID，可从 task_get 或 recovery_list 获取"),
        schedule_day: z.number().int().min(1).max(7).optional().describe("标记为第几天采集，可不传"),
      }),
      annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: true },
    },
    withError(async ({ distribution_id, schedule_day }) => {
      const data = await collectNow(distribution_id, schedule_day, options.bindings);
      return result("指定笔记采集完成。", data);
    }),
  );

  server.registerTool(
    "collection_retry_failed",
    {
      title: "补采异常数据",
      description: "重试指定任务或全部任务中采集状态异常的已发布笔记。",
      inputSchema: z.object({ task_id: z.string().optional().describe("任务ID；不传则补采最近异常任务") }),
      annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: true },
    },
    withError(async ({ task_id }) => {
      const data = await retryFailed(task_id, options.bindings);
      return result(`补采完成：成功 ${data.succeeded}，失败 ${data.failed}。`, data);
    }),
  );

  server.registerTool(
    "resource_search",
    {
      title: "搜索 KOC 账号资源",
      description: "按账号名称/账号号/标签、IP地区、合作来源或平台搜索 KOC 资源。",
      inputSchema: z.object({ ...resourceFilters, ...pagination }),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: true },
    },
    withError(async ({ query, ip_location, cooperation_source, platform, limit, offset }) => {
      const data = await resourceSearch({ query, ipLocation: ip_location, cooperationSource: cooperation_source, platform, limit, offset });
      return result(`共找到 ${data.total} 个账号。`, data);
    }),
  );

  server.registerTool(
    "resource_get",
    {
      title: "查看 KOC 账号详情",
      description: "查看账号主页、账号号、粉丝数、性别、简介、标签、IP地区以及全部合作记录。",
      inputSchema: z.object({ account_id: z.string().min(1).describe("KOC LOOP 账号ID") }),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: true },
    },
    withError(async ({ account_id }) => {
      const data = await resourceGet(account_id);
      return result(`已读取账号“${String(data.account.nickname)}”的详情。`, data);
    }),
  );

  server.registerTool(
    "resource_backfill_profile",
    {
      title: "补全公开账号信息",
      description: "根据已回填的小红书或抖音作品链接补全账号主页、昵称、账号号、IP地区和粉丝数。",
      inputSchema: z.object({
        distribution_id: z.string().optional().describe("分发记录ID，和发布链接二选一"),
        publish_url: z.string().optional().describe("小红书或抖音作品链接，也可传包含链接的分享文案"),
      }).refine((value) => Boolean(value.distribution_id || value.publish_url), "请提供分发记录ID或发布链接"),
      annotations: { readOnlyHint: false, destructiveHint: false, idempotentHint: true, openWorldHint: true },
    },
    withError(async ({ distribution_id, publish_url }) => {
      const data = await backfillResourceProfile({ distributionId: distribution_id, publishUrl: publish_url }, options.bindings);
      return result(data.updated ? "账号公开信息已补全。" : "账号信息未发生变化。", data);
    }),
  );

  server.registerTool(
    "resource_export",
    {
      title: "导出 KOC 资源",
      description: "按账号、IP地区、合作来源或平台筛选并导出 KOC 资源 Excel。",
      inputSchema: z.object(resourceFilters),
      annotations: { readOnlyHint: true, destructiveHint: false, idempotentHint: false },
    },
    withError(async ({ query, ip_location, cooperation_source, platform }) => {
      const data = await resourceExport({ query, ipLocation: ip_location, cooperationSource: cooperation_source, platform }, options.origin);
      return result(`已生成 ${data.account_count} 个账号的导出文件。`, data);
    }),
  );
}