ai-chatbot/lib/canvas/capabilities/ai-sdk-tools.ts

import { tool, type ToolSet } from "ai";
import { z } from "zod";

import { createCanvasCapabilityClientCommand } from "@/lib/canvas/capabilities/client-command";
import { CANVAS_CAPABILITY_REGISTRY } from "@/lib/canvas/capabilities/registry";
import {
  executeCanvasCapabilityRuntime,
  type CanvasCapabilityRuntimeResult,
} from "@/lib/canvas/capabilities/runtime";
import type { CanvasCapabilityDefinition } from "@/lib/canvas/capabilities/types";

type CanvasCapabilityToolCallbacks = {
  onClientCommand?: (
    command: NonNullable<CanvasCapabilityRuntimeResult["clientCommand"]>,
    capability: CanvasCapabilityDefinition
  ) => void;
  onResult?: (
    result: CanvasCapabilityRuntimeResult,
    capability: CanvasCapabilityDefinition
  ) => void;
};

function isCapabilityAvailableForSelection(params: {
  capability: CanvasCapabilityDefinition;
  selectedNodeIds: string[];
}) {
  const count = params.selectedNodeIds.length;
  return (
    params.capability.agent.exposed &&
    count >= params.capability.selection.min &&
    (!params.capability.selection.max || count <= params.capability.selection.max)
  );
}

function buildToolDescription(capability: CanvasCapabilityDefinition) {
  const modeDescription =
    capability.executionMode === "interactive"
      ? "该工具会在用户画布中打开交互面板。"
      : capability.executionMode === "headless"
        ? "该工具会在后台执行并把结果写回画布。"
        : "参数完整时后台执行，否则在用户画布中打开交互面板。";
  const aliases = capability.agent.aliases.length
    ? `常见表达：${capability.agent.aliases.join("、")}。`
    : "";
  return `${capability.description}${modeDescription}${aliases}默认作用于当前选中图片；连续操作时可把上一步返回的 createdNodeIds 传入 nodeIds。`;
}

export async function createCanvasCapabilityTools(params: {
  callbacks?: CanvasCapabilityToolCallbacks;
  defaultModelId?: string;
  destructiveConfirmed?: boolean;
  projectId: string;
  requestHeaders: Headers;
  requestUrl: string;
  selectedNodeIds: string[];
  userId: string;
}) {
  const selectedNodeIds = Array.from(
    new Set(
      params.selectedNodeIds
        .map((id) => String(id || "").trim())
        .filter(Boolean)
    )
  );
  const tools: ToolSet = {};
  const { getEnabledCanvasCapabilityIds } = await import(
    "@/lib/canvas/capabilities/availability"
  );
  const enabledCapabilityIds = await getEnabledCanvasCapabilityIds(
    CANVAS_CAPABILITY_REGISTRY
  );

  for (const capability of CANVAS_CAPABILITY_REGISTRY) {
    if (!enabledCapabilityIds.has(capability.id)) continue;
    if (!isCapabilityAvailableForSelection({ capability, selectedNodeIds })) continue;
    const toolName = capability.agent.toolName;
    if (!toolName) continue;

    tools[toolName] = tool({
      description: buildToolDescription(capability),
      inputSchema: z.object({
        nodeIds: z
          .array(z.string().trim().min(1).max(180))
          .min(1)
          .max(50)
          .optional()
          .describe("可选。传入上一个工具返回的 createdNodeIds 可链式处理新结果。"),
        args: capability.inputSchema.optional().describe("能力参数。"),
      }),
      execute: async (input) => {
        const rawArgs =
          input.args && typeof input.args === "object" && !Array.isArray(input.args)
            ? input.args
            : {};
        const args = {
          ...rawArgs,
          ...(capability.modelPolicy === "agent-selected" &&
          params.defaultModelId &&
          !("modelId" in rawArgs)
            ? { modelId: params.defaultModelId }
            : {}),
        };
        const command = createCanvasCapabilityClientCommand({
          capabilityId: capability.id,
          nodeIds: input.nodeIds?.length ? input.nodeIds : selectedNodeIds,
          projectId: params.projectId,
          args,
          source: "agent",
          confirmed:
            capability.risk === "destructive"
              ? params.destructiveConfirmed === true
              : undefined,
        });
        const result = await executeCanvasCapabilityRuntime({
          command,
          projectId: params.projectId,
          requestHeaders: params.requestHeaders,
          requestUrl: params.requestUrl,
          userId: params.userId,
        });

        if (result.clientCommand) {
          params.callbacks?.onClientCommand?.(result.clientCommand, capability);
        }
        params.callbacks?.onResult?.(result, capability);

        return {
          capabilityId: capability.id,
          title: capability.title,
          mode: result.mode,
          status: result.task.status,
          taskId: result.task.id,
          requestId: result.task.requestId,
          createdNodeIds: result.operation?.createdNodeIds || [],
          updatedNodeIds: result.operation?.updatedNodeIds || [],
          outputs: result.outputs || [],
          clientAction: result.clientCommand
            ? capability.clientAction
            : undefined,
        };
      },
    });
  }

  return tools;
}

export function getCanvasCapabilityToolNames(tools: ToolSet) {
  return Object.keys(tools);
}