writing-agent/src/preset/markers.ts

import type {
  PresetPackage,
  PresetPromptEntry,
  PresetPromptOrderItem,
  PresetPromptRole,
} from "../types/preset.js";

/** 酒馆角色卡 / WI marker：导入后保留位置；personaDescription 由当前用户角色填。 */
/** 酒馆的 worldInfoBefore / worldInfoAfter 现在由 pack 填，集合留空（保留判定入口）。 */
export const ST_UNFILLED_MARKERS = new Set<string>([]);

export const WORLD_INFO_BEFORE_MARKER = "worldInfoBefore";
export const WORLD_INFO_AFTER_MARKER = "worldInfoAfter";

export const PERSONA_DESCRIPTION_MARKER = "personaDescription";

/** 酒馆示例对话洞：new_example_chat_prompt + mes_example 切出来的多段。 */
export const DIALOGUE_EXAMPLES_MARKER = "dialogueExamples";

/** 角色卡三块：由直聊槽位填，不再当"酒馆未填洞"跳过。 */
export const CHAR_DESCRIPTION_MARKER = "charDescription";
export const CHAR_PERSONALITY_MARKER = "charPersonality";
export const SCENARIO_MARKER = "scenario";

export const CHAT_HISTORY_MARKER = "chatHistory";
export const WORLD_BOOK_BEFORE = "worldBookBefore";
export const WORLD_BOOK_AFTER = "worldBookAfter";
export const CURRENT_TURN = "currentTurn";
export const POST_TURN = "postTurn";

/** 运行时填洞的四个自有洞（不含共用的 chatHistory）；缺失时自动补进顺序。 */
export const APP_INSERTED_MARKERS = [
  WORLD_BOOK_BEFORE,
  WORLD_BOOK_AFTER,
  CURRENT_TURN,
  POST_TURN,
] as const;

type MarkerDef = {
  id: string;
  name: string;
  role: PresetPromptRole;
};

const MARKER_DEFS: MarkerDef[] = [
  { id: WORLD_BOOK_BEFORE, name: "世界书（历史前）", role: "system" },
  { id: CHAT_HISTORY_MARKER, name: "对话历史", role: "system" },
  { id: WORLD_BOOK_AFTER, name: "世界书（历史后）", role: "system" },
  { id: CURRENT_TURN, name: "本轮对话", role: "user" },
  { id: POST_TURN, name: "本轮之后", role: "system" },
];

function upsertMarkerEntry(
  prompts: PresetPromptEntry[],
  def: MarkerDef,
): PresetPromptEntry[] {
  if (prompts.some((p) => p.id === def.id)) return prompts;
  return [
    ...prompts,
    {
      id: def.id,
      name: def.name,
      enabled: true,
      role: def.role,
      content: "",
      marker: true,
      sourceIdentifier: def.id,
    },
  ];
}

function reindex(
  order: Omit<PresetPromptOrderItem, "orderIndex">[],
): PresetPromptOrderItem[] {
  return order.map((item, orderIndex) => ({ ...item, orderIndex }));
}

/**
 * 以 chatHistory 为锚点补齐四个自有 marker：**只补缺的，不动已有的**。
 * 已经排进 `prompt_order` 的洞保留作者给的顺序——工作台左轨就是真实发送顺序；
 * 缺的洞按酒馆的规范位置插进来：历史前 → 对话历史 → 历史后 → 本轮对话 → … → 本轮之后。
 * 可重复调用：补齐之后再调一次是无操作（之前会把这四个洞重新搅回锚点，覆盖作者排的顺序）。
 */
export function applyExtendedMarkers(preset: PresetPackage): PresetPackage {
  let prompts = [...preset.prompts];
  for (const def of MARKER_DEFS) {
    prompts = upsertMarkerEntry(prompts, def);
  }

  const enabledOf = new Map(
    preset.promptOrder.map((item) => [item.promptId, item.enabled]),
  );
  const next = [...preset.promptOrder].sort((a, b) => a.orderIndex - b.orderIndex);
  const has = (id: string) => next.some((item) => item.promptId === id);
  const indexOf = (id: string) => next.findIndex((item) => item.promptId === id);
  const insertAt = (index: number, id: string) => {
    const item: PresetPromptOrderItem = {
      promptId: id,
      enabled: enabledOf.get(id) ?? true,
      orderIndex: 0,
    };
    next.splice(Math.max(0, Math.min(index, next.length)), 0, item);
  };

  // 预设自带酒馆的 worldInfoBefore / worldInfoAfter（并启用）时，不插自有的世界书洞：
  // 内容直接填进预设那两条，位置上才对得上酒馆。
  const hasStBefore = next.some(
    (item) => item.promptId === WORLD_INFO_BEFORE_MARKER && item.enabled,
  );
  const hasStAfter = next.some(
    (item) => item.promptId === WORLD_INFO_AFTER_MARKER && item.enabled,
  );

  if (!has(CHAT_HISTORY_MARKER)) insertAt(next.length, CHAT_HISTORY_MARKER);
  if (!has(WORLD_BOOK_BEFORE) && !hasStBefore) insertAt(indexOf(CHAT_HISTORY_MARKER), WORLD_BOOK_BEFORE);
  if (!has(WORLD_BOOK_AFTER) && !hasStAfter) insertAt(indexOf(CHAT_HISTORY_MARKER) + 1, WORLD_BOOK_AFTER);
  if (!has(CURRENT_TURN)) {
    const afterWorldBook = indexOf(WORLD_BOOK_AFTER);
    insertAt(
      afterWorldBook >= 0 ? afterWorldBook + 1 : indexOf(CHAT_HISTORY_MARKER) + 1,
      CURRENT_TURN,
    );
  }
  if (!has(POST_TURN)) insertAt(next.length, POST_TURN);

  return {
    ...preset,
    prompts,
    promptOrder: reindex(next),
  };
}

export function isStUnfilledMarker(identifier: string): boolean {
  return ST_UNFILLED_MARKERS.has(identifier);
}