index.js

'use strict';

const log4js = require('ep_etherpad-lite/node_modules/log4js');
const authorManager = require('ep_etherpad-lite/node/db/AuthorManager');
const padManager = require('ep_etherpad-lite/node/db/PadManager');
const padMessageHandler = require('ep_etherpad-lite/node/handler/PadMessageHandler');
const {ChatMessage} = require('ep_etherpad-lite/static/js/ChatMessage');
const epAiCore = require('ep_ai_core/index');
const {t} = require('./i18n');

const {extractMention} = require('./chatHandler');
const {buildContext} = require('./contextBuilder');
const {applyEdit, announceAiAuthor} = require('./padEditor');
const {suggestEdit} = require('./suggestEdit');
const {resolveSuggestionMode} = require('./suggestionMode');

const logger = log4js.getLogger('ep_ai_chat');

const conversations = {};
const conversationLastAccess = {};
const MAX_TRACKED_PADS = 1000;
const CONVERSATION_TTL_MS = 60 * 60 * 1000; // 1 hour

// Evict stale conversations periodically
const evictStaleConversations = () => {
  const now = Date.now();
  const padIds = Object.keys(conversationLastAccess);
  // Evict expired entries
  for (const padId of padIds) {
    if (now - conversationLastAccess[padId] > CONVERSATION_TTL_MS) {
      delete conversations[padId];
      delete conversationLastAccess[padId];
    }
  }
  // If still over limit, evict oldest
  const remaining = Object.entries(conversationLastAccess);
  if (remaining.length > MAX_TRACKED_PADS) {
    remaining.sort((a, b) => a[1] - b[1]);
    const toEvict = remaining.length - MAX_TRACKED_PADS;
    for (let i = 0; i < toEvict; i++) {
      delete conversations[remaining[i][0]];
      delete conversationLastAccess[remaining[i][0]];
    }
  }
};
const evictionTimer = setInterval(evictStaleConversations, 5 * 60 * 1000);
evictionTimer.unref();

// Rate limiting: track last request time per pad. Window defaults to 5s
// in production; tests and dev can shrink it via chat.rateLimitMs.
const rateLimits = {};
const DEFAULT_RATE_LIMIT_MS = 5000;

const isRateLimited = (padId) => {
  const now = Date.now();
  const lastRequest = rateLimits[padId] || 0;
  const window = typeof chatSettings.rateLimitMs === 'number'
      ? chatSettings.rateLimitMs : DEFAULT_RATE_LIMIT_MS;
  if (now - lastRequest < window) return true;
  rateLimits[padId] = now;
  return false;
};

let chatSettings = {
  trigger: '@ai',
  authorName: 'AI Assistant',
  authorColor: '#7c4dff',
  systemPrompt: null,
  maxContextChars: 50000,
  chatHistoryLength: 20,
  conversationBufferSize: 10,
};

let aiAuthorId = null;
let commentsModulesCache = {commentManager: null, shared: null, depAvailable: false};

const probeCommentsModule = () => {
  try {
    const commentManager = require('ep_comments_page/commentManager');
    const shared = require('ep_comments_page/static/js/shared');
    commentsModulesCache = {commentManager, shared, depAvailable: true};
    logger.info('ep_comments_page detected; suggestion mode available');
  } catch {
    commentsModulesCache = {commentManager: null, shared: null, depAvailable: false};
    logger.info('ep_comments_page not installed; suggestions disabled');
  }
};

const getAiAuthorId = async () => {
  if (aiAuthorId) return aiAuthorId;
  const result = await authorManager.createAuthor(chatSettings.authorName);
  aiAuthorId = result.authorID;
  if (chatSettings.authorColor) await authorManager.setAuthorColorId(aiAuthorId, chatSettings.authorColor);
  return aiAuthorId;
};

const getRequesterDisplayName = async (authorId) => {
  if (!authorId) return 'Anonymous';
  try {
    const name = await authorManager.getAuthorName(authorId);
    return name || 'Anonymous';
  } catch {
    return 'Anonymous';
  }
};

const sendChatReply = async (padId, text) => {
  const authorId = await getAiAuthorId();
  const msg = new ChatMessage(text, authorId, Date.now());
  await padMessageHandler.sendChatMessageToPadClients(msg, padId);
};

const getConversation = (padId) => {
  if (!conversations[padId]) conversations[padId] = [];
  conversationLastAccess[padId] = Date.now();
  return conversations[padId];
};

const addToConversation = (padId, role, content) => {
  const conv = getConversation(padId);
  conv.push({role, content});
  const maxSize = chatSettings.conversationBufferSize * 2;
  while (conv.length > maxSize) conv.shift();
};

exports.loadSettings = async (hookName, {settings}) => {
  const aiSettings = settings.ep_ai_core || {};
  const chat = aiSettings.chat || {};
  chatSettings = {...chatSettings, ...chat};
  logger.info(`ep_ai_chat loaded. Trigger: "${chatSettings.trigger}"`);
  aiAuthorId = null;
  probeCommentsModule();
  const wantsSuggest = (chat.suggestionMode || '').toLowerCase() === 'suggest';
  if (!commentsModulesCache.depAvailable && wantsSuggest) {
    logger.warn(
        "ep_ai_chat: suggestionMode is set to 'suggest' but ep_comments_page is " +
        "not installed; falling back to 'apply'");
  }
};

/**
 * socketio: capture Etherpad's socket.io server reference so we can
 * broadcast USER_NEWINFO (for the AI author chip) and pushAddComment
 * (for AI suggestions) without needing an originating client socket.
 * Etherpad doesn't expose `io` as a module property — this hook is the
 * supported entry point.
 */
let socketioRef = null;
exports.socketio = (hookName, {io}) => {
  socketioRef = io;
  logger.info('ep_ai_chat: socket.io reference captured');
};
exports.getSocketIo = () => socketioRef;

/**
 * clientVars: expose the trigger string and the AI author's identity to
 * the client. The chatPrefillFromUser client hook uses this to recognise
 * the AI's row in the user list and substitute "@ai " for the otherwise
 * useless "@AI_Assistant " prefill.
 */
exports.clientVars = async (hookName, context) => {
  const aiId = await getAiAuthorId();
  return {
    ep_ai_chat: {
      trigger: chatSettings.trigger,
      authorName: chatSettings.authorName,
      authorId: aiId,
    },
  };
};

/**
 * userJoin: when a client joins a pad, push the AI's author info to the
 * room so the AI shows up in the editors list immediately — without
 * waiting for a first @ai request to trigger announceAiAuthor via an
 * applied edit. Best-effort; never let this hook throw.
 */
exports.userJoin = async (hookName, {padId}) => {
  logger.info(`userJoin hook fired for pad=${padId}`);
  try {
    if (!padId) return;
    const aiId = await getAiAuthorId();
    await announceAiAuthor(padId, aiId, socketioRef);
    logger.info(`userJoin announce sent for pad=${padId} ai=${aiId} io=${!!socketioRef}`);
  } catch (err) {
    logger.warn(`userJoin announce failed: ${err.message}`);
  }
};

exports.handleMessage = async (hookName, context) => {
  const {message} = context;
  if (!message || !message.data) return;
  if (message.type !== 'COLLABROOM' || message.data.type !== 'CHAT_MESSAGE') return;

  const chatMsg = message.data.message;
  const chatText = chatMsg?.text || message.data.text;
  if (!chatText) return;

  // Extract selection info if the client attached it
  const selection = chatMsg?.customMetadata?.selection || null;

  const {mentioned, query, override} = extractMention(chatText, chatSettings.trigger);
  if (!mentioned) return;

  const padId = context.sessionInfo?.padId;
  if (!padId) return;

  // Rate limit: prevent API cost abuse
  if (isRateLimited(padId)) {
    logger.info(`Rate limited @ai request on pad ${padId}`);
    return;
  }

  const aiSettings = epAiCore.getSettings();
  const accessMode = epAiCore.accessControl.getAccessMode(padId, aiSettings);
  if (accessMode === 'none') {
    await sendChatReply(padId, t('ep_ai_chat.no_access'));
    return;
  }

  // Audit log
  const requestAuthor = context.sessionInfo?.authorId || 'unknown';
  logger.info(`AI request: pad=${padId} author=${requestAuthor} query="${query.substring(0, 100)}"`);

  const {mode: effectiveMode, fellBackFromSuggest} =
      resolveSuggestionMode(padId, override, aiSettings, commentsModulesCache.depAvailable);
  if (override === 'suggest' && fellBackFromSuggest) {
    await sendChatReply(padId,
        'Suggestions need ep_comments_page. Falling back to direct edit. ' +
        'Install the plugin to enable review-before-apply.');
  }

  await sendChatReply(padId, '\u2728 Thinking...');

  setImmediate(async () => {
    try {
      const pad = await padManager.getPad(padId);
      const currentText = pad.text();
      const conversation = getConversation(padId);

      const llmConfig = {
        apiBaseUrl: aiSettings.apiBaseUrl,
        apiKey: aiSettings.apiKey,
        model: aiSettings.model,
        maxTokens: aiSettings.maxTokens,
        provider: aiSettings.provider,
      };
      const client = epAiCore.llmClient.create(llmConfig);

      const requesterName = await getRequesterDisplayName(requestAuthor);
      const requester = {
        authorId: requestAuthor || 'unknown',
        name: requesterName,
      };

      // Step 1: Ask the AI to decide — respond with JSON that either
      // contains an edit action or just a chat reply
      const decideMessages = await buildContext(
          pad, padId, query, conversation, chatSettings, accessMode, selection, requester,
      );

      // Augment system prompt to request structured decision
      const canEdit = accessMode === 'full';
      const editInstructions = canEdit
        ? `

When the user asks you to change, improve, edit, rewrite, fix, or modify the document in any way, you MUST respond with a JSON block containing your edit. Use this exact format:

\`\`\`json
{"action": "edit", "findText": "exact text from the document to replace", "replaceText": "the improved replacement text", "explanation": "brief explanation of what you changed"}
\`\`\`

The findText MUST be an exact substring from the current document. Be precise.

If the user is NOT asking for an edit (just asking a question, discussing content, etc.), respond normally with plain text — no JSON block.`
        : '\n\nYou have READ-ONLY access. You cannot edit the pad. Just answer questions.';

      decideMessages[0].content += editInstructions;

      const response = await client.complete(decideMessages);

      // Step 2: Check if the response contains an edit JSON block
      const jsonMatch = response.content.match(/```json\s*\n([\s\S]*?)\n```/);
      let applied = false;

      if (jsonMatch && canEdit) {
        try {
          const editData = JSON.parse(jsonMatch[1]);
          if (editData.action === 'edit' && editData.findText && editData.replaceText !== undefined) {
            editData.authorId = await getAiAuthorId();
            editData.requesterAuthorId = requestAuthor;

            const useSuggest =
                effectiveMode === 'suggest' && commentsModulesCache.depAvailable;

            let editResult;
            if (useSuggest) {
              editResult = await suggestEdit(pad, editData, {
                requesterAuthorId: requestAuthor,
                aiAuthorId: editData.authorId,
                aiAuthorName: chatSettings.authorName,
                commentManager: commentsModulesCache.commentManager,
                shared: commentsModulesCache.shared,
                io: socketioRef,
              });
              if (editResult.success) {
                applied = true;
                logger.info(
                    `AI suggestion created: pad=${padId} commentId=${editResult.commentId}`);
                await sendChatReply(padId,
                    '\ud83d\udca1 Suggestion ready \u2014 review in the comments sidebar.');
              } else {
                logger.warn(`Suggest failed, falling back to apply: ${editResult.error}`);
                editResult = await applyEdit(pad, editData, socketioRef);
                if (editResult.success) {
                  applied = true;
                  const explanation = editData.explanation || 'Edit applied.';
                  await sendChatReply(padId,
                      `\u2705 ${explanation} (suggestion failed; applied directly)`);
                }
              }
            } else {
              editResult = await applyEdit(pad, editData, socketioRef);
              if (editResult.success) {
                applied = true;
                const explanation = editData.explanation || 'Edit applied.';
                logger.info(`AI edit applied: pad=${padId} find="${editData.findText.substring(0, 50)}" replace="${editData.replaceText.substring(0, 50)}"`);
                await sendChatReply(padId, `\u2705 ${explanation}`);
              } else {
                logger.warn(`Edit failed: ${editResult.error}`);
                // Fall through to send the raw response
              }
            }
          }
        } catch (e) {
          logger.warn(`Failed to parse edit JSON: ${e.message}`);
          // Fall through to send the raw response
        }
      }

      if (!applied) {
        // Send the response as-is (strip any failed JSON blocks)
        const cleanResponse = response.content
            .replace(/```json[\s\S]*?```/g, '')
            .trim();
        await sendChatReply(padId, cleanResponse || response.content);
      }

      addToConversation(padId, 'user', query);
      addToConversation(padId, 'assistant', response.content);
    } catch (err) {
      logger.error(`AI chat error: ${err.message}`);
      let msg = t('ep_ai_chat.error_generic');
      if (err.message.includes('429')) msg = t('ep_ai_chat.error_rate_limit');
      else if (err.message.includes('401') || err.message.includes('403')) msg = t('ep_ai_chat.error_auth');
      try { await sendChatReply(padId, msg); } catch { logger.error('Failed to send error to chat'); }
    }
  });
};