ai-chatbot/lib/ai/runtime/direct-agent-image-edit-executor.ts

import {
  BoxCutoutServiceError,
  executeBoxCutoutJob,
} from "@/lib/ai/box-cutout/service";
import {
  executeExtendJob,
  ExtendServiceError,
  inferExtendOptionsFromText,
} from "@/lib/ai/extend/service";
import {
  executeResizeJob,
  inferResizeOptionsFromText,
  ResizeServiceError,
} from "@/lib/ai/resize/service";
import {
  executeRmbgJob,
  fetchImageBufferFromUrl,
  RmbgServiceError,
} from "@/lib/ai/rmbg/service";
import { generateCanvasImage } from "@/lib/canvas/generate-image";
import { requestLongTaskResult } from "@/lib/client/long-task";
import { getMediaAssetsByUserAndUrls } from "@/lib/db/queries";
import type {
  AgentExecutionPlan,
  AgentPreparedExecution,
} from "@/lib/ai/auto-agent-graph";
import {
  findAgentToolCallArgs,
  type DirectAgentToolCapability,
} from "@/lib/ai/runtime/direct-agent-execution";
import { buildDirectAgentExecutionPrompt } from "@/lib/ai/runtime/optimized-prompt";
import { buildAgentToolImagesMarkdown } from "@/lib/ai/runtime/agent-media-markdown";

interface DirectAgentImageParamsSnapshot {
  negativePrompt?: string;
  aspectRatio?: string;
  imageSize?: string;
}

type SaveGeneratedImages = (
  imageMarkdown: string,
  userId: string,
  chatId: string | null,
  messageId: string | null,
  prompt: string,
  modelId: string,
  options?: {
    sourceAssetIds?: string[];
    sourceImageUrls?: string[];
    relationType?: "generated_from" | "edited_from";
    sourceType?: string;
    projectId?: string | null;
    userPrompt?: string;
    optimizedPrompt?: string;
    requestedImageCount?: number;
    aspectRatio?: string;
    imageSize?: string;
    skipBilling?: boolean;
  },
) => Promise<{
  savedImageCount: number;
  savedImageIds: string[];
  savedAssetIds: string[];
  savedOriginalUrls: string[];
}>;

export interface DirectAgentImageEditExecutionResult {
  originalUrl: string;
  thumbnailUrl: string | null;
  imageSize: string;
  mediaAssetId: string | null;
}

export interface ExecuteDirectAgentImageEditToolParams {
  capability: DirectAgentToolCapability;
  requestUrl: string;
  requestHeaders: Headers;
  userId: string;
  chatId: string;
  messageId: string;
  projectId?: string | null;
  targetAssetIds?: string[] | null;
  sourceImageUrl: string;
  sourceImageUrls: string[];
  requestTextContent: string;
  optimizedPrompt?: string;
  hiddenPromptPrefix?: string;
  finalModelId: string;
  imageParams?: DirectAgentImageParamsSnapshot | null;
  defaultImageQualityParams?: DirectAgentImageParamsSnapshot | null;
  smartAutoExecution?: AgentPreparedExecution;
  smartAutoPlan?: AgentExecutionPlan;
  saveGeneratedImages: SaveGeneratedImages;
}

function toFiniteNumber(value: unknown): number | undefined {
  const normalized = Number(value);
  return Number.isFinite(normalized) ? normalized : undefined;
}

async function executeDirectAgentEnhanceJob(params: {
  requestUrl: string;
  requestHeaders: Headers;
  imageUrl: string;
  prompt: string;
  chatId: string;
  messageId: string;
}) {
  const endpoint = new URL("/api/enhance", params.requestUrl).toString();
  const cookieHeader = params.requestHeaders.get("cookie") || "";
  const result = await requestLongTaskResult({
    endpoint,
    payload: {
      imageUrl: params.imageUrl,
      prompt: params.prompt,
      chatId: params.chatId,
      messageId: params.messageId,
    },
    pollingParams: {
      prompt: params.prompt,
      chatId: params.chatId,
      messageId: params.messageId,
      sourceImageUrl: params.imageUrl,
    },
    timeoutMs: 10 * 60 * 1000,
    headers: cookieHeader
      ? {
          cookie: cookieHeader,
        }
      : undefined,
  });

  if (!result?.success || !result?.original_url) {
    throw new Error(String(result?.error || result?.detail || "图像增强失败"));
  }

  return {
    originalUrl: String(result.original_url),
    thumbnailUrl: String(result.thumbnail_url || result.original_url),
    imageSize: "2K",
  };
}

export function isDirectAgentImageEditServiceError(error: unknown) {
  return (
    error instanceof RmbgServiceError ||
    error instanceof ResizeServiceError ||
    error instanceof ExtendServiceError ||
    error instanceof BoxCutoutServiceError
  );
}

async function resolveDirectAgentResultAssetId(params: {
  userId: string;
  originalUrl: string;
  mediaAssetId?: string | null;
  capability: DirectAgentToolCapability;
}) {
  if (params.mediaAssetId || !params.originalUrl) {
    return params.mediaAssetId || null;
  }

  try {
    const [resultAsset] = await getMediaAssetsByUserAndUrls({
      userId: params.userId,
      urls: [params.originalUrl],
    });
    return resultAsset?.id || null;
  } catch (assetLookupError) {
    console.warn(
      `[Direct Agent Image Edit] Failed to lookup result asset after ${params.capability}:`,
      assetLookupError,
    );
    return null;
  }
}

async function finalizeDirectAgentImageEditResult(params: {
  userId: string;
  capability: DirectAgentToolCapability;
  originalUrl: string;
  thumbnailUrl?: string | null;
  imageSize: string;
  mediaAssetId?: string | null;
}): Promise<DirectAgentImageEditExecutionResult> {
  return {
    originalUrl: params.originalUrl,
    thumbnailUrl: params.thumbnailUrl || null,
    imageSize: params.imageSize,
    mediaAssetId: await resolveDirectAgentResultAssetId({
      userId: params.userId,
      originalUrl: params.originalUrl,
      mediaAssetId: params.mediaAssetId,
      capability: params.capability,
    }),
  };
}

export async function executeDirectAgentImageEditTool(
  params: ExecuteDirectAgentImageEditToolParams,
): Promise<DirectAgentImageEditExecutionResult> {
  if (params.capability === "generic_edit") {
    const directEditPrompt =
      buildDirectAgentExecutionPrompt({
        originalText: params.requestTextContent || "",
        optimizedPrompt: params.optimizedPrompt,
        hiddenPromptPrefix: params.hiddenPromptPrefix,
        fallback: "在保持主体和整体构图稳定的前提下完成本轮图片编辑。",
      }) || "在保持主体和整体构图稳定的前提下完成本轮图片编辑。";
    const aspectRatio =
      params.imageParams?.aspectRatio || params.defaultImageQualityParams?.aspectRatio;
    const imageSize =
      params.imageParams?.imageSize || params.defaultImageQualityParams?.imageSize || "2K";
    const result = await generateCanvasImage({
      prompt: directEditPrompt,
      negativePrompt: params.imageParams?.negativePrompt,
      model: params.finalModelId,
      ownerKey: params.userId,
      aspectRatio,
      imageSize,
      imageCount: 1,
      referenceImages: params.sourceImageUrls,
      debugLabel: "agent-direct-edit",
      debugMeta: {
        chatId: params.chatId,
        projectId: params.projectId || null,
        subIntent: params.smartAutoExecution?.routeSubIntent || null,
      },
    });

    const firstImage = result.images[0];
    if (!firstImage?.originalUrl) {
      throw new Error("图片编辑完成但未返回结果图");
    }

    const originalUrl = firstImage.originalUrl;
    const thumbnailUrl = firstImage.thumbnailUrl || null;
    const editResultMarkdown = buildAgentToolImagesMarkdown([
      {
        originalUrl,
        thumbnailUrl,
        imageSize,
      },
    ]);

    await params.saveGeneratedImages(
      editResultMarkdown,
      params.userId,
      params.chatId,
      params.messageId,
      directEditPrompt,
      result.modelRef || params.finalModelId,
      {
        sourceAssetIds: params.targetAssetIds || undefined,
        sourceImageUrls: params.sourceImageUrls,
        relationType: "edited_from",
        sourceType: "edited",
        projectId: params.projectId || null,
        userPrompt: params.requestTextContent,
        optimizedPrompt: directEditPrompt,
        requestedImageCount: 1,
        aspectRatio,
        imageSize,
      },
    );

    return finalizeDirectAgentImageEditResult({
      userId: params.userId,
      capability: params.capability,
      originalUrl,
      thumbnailUrl,
      imageSize,
      mediaAssetId: null,
    });
  }

  if (params.capability === "enhance") {
    const result = await executeDirectAgentEnhanceJob({
      requestUrl: params.requestUrl,
      requestHeaders: params.requestHeaders,
      imageUrl: params.sourceImageUrl,
      prompt: params.requestTextContent || "图像增强处理",
      chatId: params.chatId,
      messageId: params.messageId,
    });

    return finalizeDirectAgentImageEditResult({
      userId: params.userId,
      capability: params.capability,
      originalUrl: result.originalUrl,
      thumbnailUrl: result.thumbnailUrl,
      imageSize: result.imageSize,
      mediaAssetId: null,
    });
  }

  if (params.capability === "remove_background") {
    const sourceBuffer = await fetchImageBufferFromUrl(params.sourceImageUrl);
    const result = await executeRmbgJob({
      userId: params.userId,
      imageBuffer: sourceBuffer,
      prompt: params.requestTextContent || "抠图处理",
      chatId: params.chatId,
      messageId: params.messageId,
      sourceImageUrl: params.sourceImageUrl,
    });

    return finalizeDirectAgentImageEditResult({
      userId: params.userId,
      capability: params.capability,
      originalUrl: result.original_url,
      thumbnailUrl: result.thumbnail_url,
      imageSize: "2K",
      mediaAssetId: result.mediaAssetId || null,
    });
  }

  if (params.capability === "resize") {
    const plannedArgs = findAgentToolCallArgs(
      "resizeImage",
      params.smartAutoExecution,
      params.smartAutoPlan,
    );
    const inferredOptions = inferResizeOptionsFromText(params.requestTextContent);
    const canvasWidth = toFiniteNumber(
      plannedArgs?.canvasWidth ?? plannedArgs?.width ?? inferredOptions.canvasWidth,
    );
    const canvasHeight = toFiniteNumber(
      plannedArgs?.canvasHeight ?? plannedArgs?.height ?? inferredOptions.canvasHeight,
    );

    if (!canvasWidth || !canvasHeight) {
      throw new ResizeServiceError(
        "请告诉我目标尺寸或比例，例如 1024x1024、4:5、16:9。",
        400,
      );
    }

    const transparentBackground =
      plannedArgs?.transparentBackground === true ||
      plannedArgs?.transparentBackground === "true" ||
      inferredOptions.transparentBackground === true;
    const backgroundColor =
      typeof plannedArgs?.backgroundColor === "string"
        ? plannedArgs.backgroundColor
        : inferredOptions.backgroundColor;
    const mimeType =
      typeof plannedArgs?.mimeType === "string" ? plannedArgs.mimeType : undefined;
    const result = await executeResizeJob({
      userId: params.userId,
      imageUrl: params.sourceImageUrl,
      prompt: params.requestTextContent || "尺寸调整处理",
      chatId: params.chatId,
      messageId: params.messageId,
      sourceImageUrl: params.sourceImageUrl,
      canvasWidth,
      canvasHeight,
      transparentBackground,
      backgroundColor,
      mimeType,
    });

    return finalizeDirectAgentImageEditResult({
      userId: params.userId,
      capability: params.capability,
      originalUrl: result.original_url,
      thumbnailUrl: result.thumbnail_url,
      imageSize: `${result.canvas_width}x${result.canvas_height}`,
      mediaAssetId: result.mediaAssetId || null,
    });
  }

  if (params.capability === "extend") {
    const plannedArgs = findAgentToolCallArgs(
      "extendImage",
      params.smartAutoExecution,
      params.smartAutoPlan,
    );
    const inferredOptions = inferExtendOptionsFromText(params.requestTextContent);
    const xScale = toFiniteNumber(plannedArgs?.xScale ?? inferredOptions.xScale);
    const yScale = toFiniteNumber(plannedArgs?.yScale ?? inferredOptions.yScale);
    const result = await executeExtendJob({
      userId: params.userId,
      imageUrl: params.sourceImageUrl,
      prompt: params.requestTextContent || "图像拓展处理",
      chatId: params.chatId,
      messageId: params.messageId,
      sourceImageUrl: params.sourceImageUrl,
      ...(xScale ? { xScale } : {}),
      ...(yScale ? { yScale } : {}),
    });

    return finalizeDirectAgentImageEditResult({
      userId: params.userId,
      capability: params.capability,
      originalUrl: result.original_url,
      thumbnailUrl: result.thumbnail_url,
      imageSize: "2K",
      mediaAssetId: result.mediaAssetId || null,
    });
  }

  if (params.capability === "box_cutout") {
    const plannedArgs = findAgentToolCallArgs(
      "boxCutout",
      params.smartAutoExecution,
      params.smartAutoPlan,
    );
    const cropRect =
      plannedArgs?.cropRect && typeof plannedArgs.cropRect === "object"
        ? (plannedArgs.cropRect as Record<string, unknown>)
        : undefined;

    if (!cropRect) {
      throw new BoxCutoutServiceError(
        "缺少框选区域，暂时无法执行框选抠图。",
        400,
      );
    }

    const result = await executeBoxCutoutJob({
      userId: params.userId,
      imageUrl: params.sourceImageUrl,
      cropRect: {
        x: toFiniteNumber(cropRect.x),
        y: toFiniteNumber(cropRect.y),
        width: toFiniteNumber(cropRect.width),
        height: toFiniteNumber(cropRect.height),
      },
      paddingRatio: toFiniteNumber(plannedArgs?.paddingRatio),
      prompt: params.requestTextContent || "框选抠图处理",
      chatId: params.chatId,
      messageId: params.messageId,
      sourceNodeId:
        typeof plannedArgs?.sourceNodeId === "string"
          ? plannedArgs.sourceNodeId
          : undefined,
    });

    return finalizeDirectAgentImageEditResult({
      userId: params.userId,
      capability: params.capability,
      originalUrl: result.original_url,
      thumbnailUrl: result.thumbnail_url || null,
      imageSize: "2K",
      mediaAssetId: result.mediaAssetId || null,
    });
  }

  throw new Error(`不支持的图片编辑工具能力：${params.capability}`);
}