Genarrative/apps/ai-game-creator-shell/tests/chatPromptPolish.test.tsx

// @vitest-environment jsdom
import {
  cleanup,
  fireEvent,
  render,
  screen,
  waitFor,
  within,
} from '@testing-library/react';
import userEvent from '@testing-library/user-event';
import { useState } from 'react';
import { afterEach, describe, expect, test, vi } from 'vitest';

import {
  CHAT_PROMPT_POLISH_REMINDER_MIN_TEXT_LENGTH,
  CHAT_PROMPT_POLISH_REMINDER_STORAGE_KEY,
  chatPromptDraftKey,
  readChatPromptPolishReminderDisabled,
  requestChatPromptPolish,
  shouldRemindChatPromptPolish,
  writeChatPromptPolishReminderDisabled,
} from '../src/features/project-workspace/chatPromptPolish';
import { ResourceReferenceInput } from '../src/features/project-workspace/ResourceReferenceInput';
import {
  type ChatComposerDraft,
  type ChatReference,
  resourceReferenceFromAsset,
} from '../src/features/project-workspace/resourceReferences';

type TauriInvoke = (
  command: string,
  args?: Record<string, unknown>,
) => Promise<unknown>;

const LONG_PROMPT =
  '做一个像素风横版动作小游戏，包含三段跳跃关卡、三个 Boss 和可以收集的金币与道具。';

function installTauriInvoke(invoke: TauriInvoke) {
  const mock = vi.fn(invoke);
  (
    window as unknown as {
      __TAURI__?: { core?: { invoke?: typeof mock } };
    }
  ).__TAURI__ = { core: { invoke: mock } };
  return mock;
}

function installPolishingInvoke(...results: string[]) {
  const queue = [...results];
  return installTauriInvoke(async (command) => {
    if (command !== 'polish_local_project_prompt') return undefined;
    const next = queue.shift();
    if (next === undefined) {
      throw new Error('fixture has no more polish results');
    }
    return next;
  });
}

// Lexical 的编辑器状态提交排在微任务里，读取输入区文本前先让 React 追平编辑器内容。
async function settleComposer() {
  await new Promise((resolve) => {
    setTimeout(resolve, 0);
  });
}

async function composerText() {
  await settleComposer();
  return screen.getByLabelText('创作想法').textContent ?? '';
}

function ControlledChatComposer({
  initialText,
  initialReferences = [],
  onSubmitDraft,
}: {
  initialText: string;
  initialReferences?: ChatReference[];
  onSubmitDraft: (draft: ChatComposerDraft) => void;
}) {
  const [draft, setDraft] = useState<ChatComposerDraft>({
    text: initialText,
    references: initialReferences,
  });
  return (
    <form
      onSubmit={(event) => {
        event.preventDefault();
        onSubmitDraft(draft);
      }}
    >
      <ResourceReferenceInput
        value={draft.text}
        references={draft.references}
        onChange={setDraft}
        assets={[]}
        projectPath="C:/project"
        ariaLabel="创作想法"
      />
      <button type="submit">发送</button>
    </form>
  );
}

function renderComposer({
  initialText,
  initialReferences,
  onSubmitDraft = vi.fn(),
}: {
  initialText: string;
  initialReferences?: ChatReference[];
  onSubmitDraft?: (draft: ChatComposerDraft) => void;
}) {
  render(
    <ControlledChatComposer
      initialText={initialText}
      initialReferences={initialReferences}
      onSubmitDraft={onSubmitDraft}
    />,
  );
  return onSubmitDraft;
}

function sendButton() {
  return screen.getByRole('button', { name: '发送' });
}

function reminderPanel() {
  return screen.getByRole('dialog', { name: '发送前提醒' });
}

afterEach(() => {
  cleanup();
  delete (
    window as unknown as { __TAURI__?: { core?: { invoke?: TauriInvoke } } }
  ).__TAURI__;
  window.localStorage.clear();
});

describe('发送前提醒判据', () => {
  test('only reminds for long plain prompts that were not acknowledged this round', () => {
    const draft: ChatComposerDraft = {
      text: '字'.repeat(CHAT_PROMPT_POLISH_REMINDER_MIN_TEXT_LENGTH),
      references: [],
    };
    expect(
      shouldRemindChatPromptPolish({
        draft,
        acknowledgedDraftKey: null,
        reminderDisabled: false,
      }),
    ).toBe(true);
    expect(
      shouldRemindChatPromptPolish({
        draft: { text: '短需求', references: [] },
        acknowledgedDraftKey: null,
        reminderDisabled: false,
      }),
    ).toBe(false);
    expect(
      shouldRemindChatPromptPolish({
        draft,
        acknowledgedDraftKey: null,
        reminderDisabled: true,
      }),
    ).toBe(false);
    expect(
      shouldRemindChatPromptPolish({
        draft,
        acknowledgedDraftKey: chatPromptDraftKey(draft),
        reminderDisabled: false,
      }),
    ).toBe(false);
    expect(
      shouldRemindChatPromptPolish({
        draft: { text: `/${'长'.repeat(60)}`, references: [] },
        acknowledgedDraftKey: null,
        reminderDisabled: false,
      }),
    ).toBe(false);
  });

  test('changes the draft key when the references change', () => {
    const reference = resourceReferenceFromAsset(
      {
        id: 'hero',
        kind: 'character',
        mediaType: 'image/png',
        localPath: 'assets/hero.png',
        source: { kind: 'uploaded' },
      },
      'asset-picker',
    );
    expect(chatPromptDraftKey({ text: '需求', references: [] })).not.toBe(
      chatPromptDraftKey({ text: '需求', references: [reference] }),
    );
  });

  test('persists the 不再提醒 preference on this machine only', () => {
    expect(readChatPromptPolishReminderDisabled()).toBe(false);
    writeChatPromptPolishReminderDisabled(true);
    expect(
      window.localStorage.getItem(CHAT_PROMPT_POLISH_REMINDER_STORAGE_KEY),
    ).toBe('true');
    expect(readChatPromptPolishReminderDisabled()).toBe(true);
    writeChatPromptPolishReminderDisabled(false);
    expect(
      window.localStorage.getItem(CHAT_PROMPT_POLISH_REMINDER_STORAGE_KEY),
    ).toBeNull();
    expect(readChatPromptPolishReminderDisabled()).toBe(false);
  });
});

describe('requestChatPromptPolish', () => {
  test('goes through the short-text Tauri command with trimmed input', async () => {
    const invoke = installPolishingInvoke('  润色后的需求  ');
    await expect(
      requestChatPromptPolish('  原始需求  ', '  项目上下文  '),
    ).resolves.toBe('润色后的需求');
    expect(invoke).toHaveBeenCalledWith('polish_local_project_prompt', {
      prompt: '原始需求',
      context: '项目上下文',
    });
  });

  test('omits blank context and keeps failures recoverable', async () => {
    const invoke = installPolishingInvoke('润色结果');
    await requestChatPromptPolish('原始需求', '   ');
    expect(invoke).toHaveBeenCalledWith('polish_local_project_prompt', {
      prompt: '原始需求',
      context: null,
    });

    const failing = vi.fn(async () => {
      throw new Error('platform llm unavailable');
    });
    installTauriInvoke(failing);
    await expect(requestChatPromptPolish('原始需求')).resolves.toBeNull();
    await expect(requestChatPromptPolish('   ')).resolves.toBeNull();
  });

  test('treats an empty model reply as a failure', async () => {
    installPolishingInvoke('   ');
    await expect(requestChatPromptPolish('原始需求')).resolves.toBeNull();
  });

  test('returns null when the native bridge is unavailable', async () => {
    delete (
      window as unknown as { __TAURI__?: { core?: { invoke?: TauriInvoke } } }
    ).__TAURI__;
    await expect(requestChatPromptPolish('原始需求')).resolves.toBeNull();
  });
});

describe('聊天输入区 AI 润色与发送前提醒', () => {
  test('fills back the polished result, can polish again, and always restores the first original', async () => {
    const user = userEvent.setup();
    const invoke = installPolishingInvoke('第一版润色结果', '第二版润色结果');
    renderComposer({ initialText: '原本的需求' });

    await user.click(screen.getByRole('button', { name: 'AI 润色' }));
    expect(await composerText()).toBe('第一版润色结果');
    expect(invoke).toHaveBeenCalledWith('polish_local_project_prompt', {
      prompt: '原本的需求',
      context: 'C:/project',
    });

    await user.click(screen.getByRole('button', { name: 'AI 润色' }));
    expect(await composerText()).toBe('第二版润色结果');
    expect(invoke.mock.calls.at(-1)?.[1]).toEqual({
      prompt: '第一版润色结果',
      context: 'C:/project',
    });

    await user.click(screen.getByRole('button', { name: '恢复原文' }));
    expect(await composerText()).toBe('原本的需求');
    expect(screen.queryByRole('button', { name: '恢复原文' })).toBeNull();
  });

  test('keeps the original text and shows a retry hint when polishing fails', async () => {
    const user = userEvent.setup();
    installTauriInvoke(async (command) => {
      if (command !== 'polish_local_project_prompt') return undefined;
      throw new Error('platform llm timeout');
    });
    renderComposer({ initialText: '原本的需求' });

    await user.click(screen.getByRole('button', { name: 'AI 润色' }));
    expect(await composerText()).toBe('原本的需求');
    expect(await screen.findByText('AI 润色失败，可重试')).not.toBeNull();
    expect(screen.queryByRole('button', { name: '恢复原文' })).toBeNull();
  });

  test('回包与原文相同时聊天输入区也给提示，且不冒出「恢复原文」', async () => {
    const user = userEvent.setup();
    installPolishingInvoke('原本的需求');
    renderComposer({ initialText: '原本的需求' });

    await user.click(screen.getByRole('button', { name: 'AI 润色' }));

    // 聊天输入区与资源侧共用 `usePromptPolish`：修好状态机还不够，共享输入区必须把
    // notice 渲染出来，否则这三个宿主里「点了没反应」看起来还是按钮坏了。
    expect(
      await screen.findByText('AI 润色结果与原文相同，未做修改'),
    ).not.toBeNull();
    expect(await composerText()).toBe('原本的需求');
    expect(screen.queryByRole('button', { name: '恢复原文' })).toBeNull();
  });

  test('holds a long prompt behind the reminder panel and sends the original on demand', async () => {
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    expect(reminderPanel()).not.toBeNull();
    expect(onSubmitDraft).not.toHaveBeenCalled();

    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: '使用原文提交' }),
    );
    await waitFor(() => {
      expect(onSubmitDraft).toHaveBeenCalledWith({
        text: LONG_PROMPT,
        references: [],
      });
    });
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).toBeNull();
  });

  test('closing the reminder cancels the send and keeps the reminder armed', async () => {
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: '关闭' }),
    );
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).toBeNull();
    expect(onSubmitDraft).not.toHaveBeenCalled();

    fireEvent.click(sendButton());
    expect(reminderPanel()).not.toBeNull();
    expect(onSubmitDraft).not.toHaveBeenCalled();
  });

  test('polishes first and then submits the polished prompt', async () => {
    installPolishingInvoke('润色后的长需求');
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: 'AI 润色' }),
    );
    await waitFor(() => {
      expect(onSubmitDraft).toHaveBeenCalledWith({
        text: '润色后的长需求',
        references: [],
      });
    });
    expect(await composerText()).toBe('润色后的长需求');
  });

  test('keeps the reminder open with the original text when the in-panel polish fails', async () => {
    installTauriInvoke(async () => {
      throw new Error('platform llm timeout');
    });
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: 'AI 润色' }),
    );

    expect(
      await screen.findByText('AI 润色失败，可重试或使用原文提交'),
    ).not.toBeNull();
    expect(onSubmitDraft).not.toHaveBeenCalled();
    expect(await composerText()).toBe(LONG_PROMPT);

    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: '使用原文提交' }),
    );
    await waitFor(() => {
      expect(onSubmitDraft).toHaveBeenCalledWith({
        text: LONG_PROMPT,
        references: [],
      });
    });
  });

  test('keeps the reminder open on Escape while the in-panel polish is in flight', async () => {
    let resolvePolish: (value: string) => void = () => {};
    installTauriInvoke(async (command) => {
      if (command !== 'polish_local_project_prompt') return undefined;
      return new Promise<string>((resolve) => {
        resolvePolish = resolve;
      });
    });
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: 'AI 润色' }),
    );
    await waitFor(() => {
      expect(
        within(reminderPanel())
          .getByRole('button', { name: 'AI 润色' })
          .getAttribute('aria-busy'),
      ).toBe('true');
    });

    // Escape 与「关闭」按钮、遮罩点击同一口径：在飞期间不许关面板，
    // 否则润色回来还会在用户已经取消之后继续把这一轮提交出去。
    fireEvent.keyDown(reminderPanel(), { key: 'Escape' });
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).not.toBeNull();
    expect(onSubmitDraft).not.toHaveBeenCalled();

    // 在飞请求照旧走完：润色回填后把这一轮发出去。
    resolvePolish('润色后的长需求');
    await waitFor(() => {
      expect(onSubmitDraft).toHaveBeenCalledWith({
        text: '润色后的长需求',
        references: [],
      });
    });
  });

  test('persists 不再提醒 locally and stops holding later sends', async () => {
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    fireEvent.click(
      within(reminderPanel()).getByRole('checkbox', { name: '不再提醒' }),
    );
    expect(
      window.localStorage.getItem(CHAT_PROMPT_POLISH_REMINDER_STORAGE_KEY),
    ).toBe('true');

    fireEvent.click(
      within(reminderPanel()).getByRole('button', { name: '关闭' }),
    );
    fireEvent.click(sendButton());
    await waitFor(() => {
      expect(onSubmitDraft).toHaveBeenCalledWith({
        text: LONG_PROMPT,
        references: [],
      });
    });
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).toBeNull();
  });

  test('reads the stored 不再提醒 preference when the composer mounts', () => {
    window.localStorage.setItem(
      CHAT_PROMPT_POLISH_REMINDER_STORAGE_KEY,
      'true',
    );
    const onSubmitDraft = renderComposer({ initialText: LONG_PROMPT });

    fireEvent.click(sendButton());
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).toBeNull();
    expect(onSubmitDraft).toHaveBeenCalledWith({
      text: LONG_PROMPT,
      references: [],
    });
  });

  test('does not hold short prompts or slash commands', () => {
    const shortSubmit = renderComposer({ initialText: '做个跳跃游戏' });
    fireEvent.click(sendButton());
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).toBeNull();
    expect(shortSubmit).toHaveBeenCalledWith({
      text: '做个跳跃游戏',
      references: [],
    });

    cleanup();
    const commandSubmit = renderComposer({
      initialText: `/${'命令'.repeat(40)}`,
    });
    fireEvent.click(sendButton());
    expect(screen.queryByRole('dialog', { name: '发送前提醒' })).toBeNull();
    expect(commandSubmit).toHaveBeenCalledWith({
      text: `/${'命令'.repeat(40)}`,
      references: [],
    });
  });
});