248 lines
10 KiB
TypeScript
248 lines
10 KiB
TypeScript
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天的自动采集日，系统在北京时间10: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);
|
||
}),
|
||
);
|
||
}
|