ai-chatbot/lib/ai/runtime/task-boundary-context-pruner.ts

import "server-only";

import type { DBMessage } from "@/lib/db/schema";
import type {
  CanvasImageNodeContext,
  ContextPayload,
  WorkingStateImageMemory,
} from "@/lib/ai/context-memory/types";
import type { TaskBoundaryDecision } from "@/lib/ai/runtime/task-boundary-classifier";

function buildSelectedCanvasSubset(contextPayload?: ContextPayload) {
  const selectedNodeIds = new Set(contextPayload?.selectedNodeIds || []);
  const selectedAssetIds = new Set(contextPayload?.selectedAssetIds || []);
  const selectedNodes = (contextPayload?.canvasImageNodes || []).filter((node) => {
    if (!node) return false;
    if (node.isSelected) return true;
    if (node.nodeId && selectedNodeIds.has(node.nodeId)) return true;
    if (node.assetId && selectedAssetIds.has(node.assetId)) return true;
    return false;
  });

  return {
    selectedNodes,
    selectedNodeIds: selectedNodes.map((node) => node.nodeId).filter(Boolean),
    selectedAssetIds: selectedNodes.map((node) => node.assetId).filter(Boolean) as string[],
  };
}

function buildSelectedCanvasSummary(selectedNodes: CanvasImageNodeContext[]) {
  if (selectedNodes.length === 0) return undefined;
  const previewLines = selectedNodes.slice(0, 3).map((node, index) => {
    const title = String(node.title || node.prompt || node.nodeId || `图片${index + 1}`).trim();
    return `选中图片${index + 1}: ${title}`;
  });
  return [`用户已选中 ${selectedNodes.length} 张图片，优先处理这些`, ...previewLines].join("\n");
}

export function pruneTaskBoundaryContext(params: {
  decision: TaskBoundaryDecision;
  contextPayload?: ContextPayload;
  autoContextText?: string;
  workingState?: WorkingStateImageMemory | null;
  hasCurrentImageInput: boolean;
}) {
  if (!params.contextPayload && !params.autoContextText) {
    return {
      contextPayload: params.contextPayload,
      autoContextText: params.autoContextText,
      workingState: params.decision.shouldUseWorkflowMemory ? params.workingState || null : null,
    };
  }

  const selectedSubset = buildSelectedCanvasSubset(params.contextPayload);
  const shouldPreferSelectedOnly =
    selectedSubset.selectedNodes.length > 0 &&
    (params.decision.taskType === "prompt_only" ||
      params.decision.taskType === "analysis_only" ||
      params.decision.taskType === "execution_image" ||
      params.decision.taskType === "execution_video" ||
      params.hasCurrentImageInput);

  let nextContextPayload = params.contextPayload;
  let nextAutoContextText = params.autoContextText;

  if (shouldPreferSelectedOnly) {
    nextContextPayload = params.contextPayload
      ? {
          ...params.contextPayload,
          focusPolicy: "selected",
          selectedNodeIds: selectedSubset.selectedNodeIds,
          selectedAssetIds: selectedSubset.selectedAssetIds,
          canvasImageNodes: selectedSubset.selectedNodes,
        }
      : undefined;
    nextAutoContextText = buildSelectedCanvasSummary(selectedSubset.selectedNodes);
  } else if (!params.decision.shouldUseCanvasContext) {
    nextContextPayload = params.contextPayload
      ? {
          source: params.contextPayload.source,
          projectId: params.contextPayload.projectId,
          focusPolicy: params.contextPayload.focusPolicy,
          selectedNodeIds: [],
          selectedAssetIds: [],
          canvasImageNodes: [],
        }
      : undefined;
    nextAutoContextText = undefined;
  }

  return {
    contextPayload: nextContextPayload,
    autoContextText: nextAutoContextText,
    workingState: params.decision.shouldUseWorkflowMemory ? params.workingState || null : null,
  };
}

function sanitizeHistoryText(text: string) {
  return String(text || "")
    .replace(/!\[[^\]]*]\([^)]+\)/g, "")
    .replace(/<video[^>]*>[\s\S]*?<\/video>/gi, "")
    .replace(/<video[^>]+src="[^"]+"[^>]*>/gi, "")
    .replace(/<!--[\s\S]*?-->/g, "")
    .replace(/◈\s*[^\n]+\n*/g, "")
    .replace(/\n{3,}/g, "\n\n")
    .trim();
}

function toTextOnlyHistoryMessages(messages: DBMessage[], maxMessages: number): DBMessage[] {
  const sliced = messages.slice(-Math.max(1, maxMessages));
  const result: DBMessage[] = [];

  for (const message of sliced) {
    const parts = Array.isArray(message.parts) ? (message.parts as Array<Record<string, unknown>>) : [];
    const textParts = parts
      .filter((part) => part?.type === "text" && typeof part?.text === "string")
      .map((part) => ({
        type: "text",
        text: sanitizeHistoryText(String(part.text || "")),
      }))
      .filter((part) => part.text);

    if (textParts.length === 0) continue;

    result.push({
      ...message,
      parts: textParts,
      attachments: [],
    } satisfies DBMessage);
  }

  return result;
}

export function pruneTaskBoundaryHistoryMessages(params: {
  decision: TaskBoundaryDecision;
  messages: DBMessage[];
  purpose: "media" | "model";
  hasCurrentImageInput: boolean;
}): DBMessage[] {
  const messages = Array.isArray(params.messages) ? params.messages : [];
  if (messages.length === 0) return [] as DBMessage[];

  switch (params.decision.taskType) {
    case "greeting":
      return [];
    case "prompt_only":
      return params.purpose === "media"
        ? []
        : toTextOnlyHistoryMessages(messages, params.hasCurrentImageInput ? 2 : 4);
    case "analysis_only":
      return params.purpose === "media"
        ? []
        : toTextOnlyHistoryMessages(messages, 4);
    case "chat":
      return params.purpose === "media"
        ? []
        : toTextOnlyHistoryMessages(messages, 6);
    case "execution_image":
    case "execution_video":
      if (params.decision.shouldUseHistoryImages) {
        return messages.slice(-8);
      }
      return params.purpose === "media" ? [] : toTextOnlyHistoryMessages(messages, 4);
    case "workflow_followup":
      return messages.slice(-(params.purpose === "media" ? 8 : 12));
    default:
      return params.purpose === "media" ? [] : toTextOnlyHistoryMessages(messages, 4);
  }
}