Moser9815-ComfyUI-RMAutomation/js/rm_content_policy.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

/**
 * Content Policy editor.
 *
 * Three term lists, edited as structured rows rather than raw text: severity is a
 * segmented control, not a punctuation prefix you have to remember. A Text tab is kept
 * alongside for bulk paste, because typing 100 terms one row at a time is nobody's idea
 * of user friendly.
 *
 * The live tester is the point of the whole screen. A word list is abstract; "paste a
 * prompt, watch what happens to it" is not. Edits are testable before saving because the
 * test runs against the in-memory draft, not the file on disk.
 */

const TABS = [
  { key: "blocked", label: "Blocked", hint: "Terms to catch. Severity decides whether a match blocks outright or only inside a sexual prompt." },
  { key: "allowed", label: "Allowed", hint: "Phrases exempted before matching, so a flagged word inside a proper noun does not fire." },
  { key: "context", label: "Context", hint: "Words that make a prompt count as sexual. These activate the 'In context' terms — they never block on their own." },
];

const SEVERITIES = [
  { key: "always",     short: "Always",     tip: "Blocks on its own, in any prompt. Never sanitized — the image is replaced with a grey placeholder." },
  { key: "contextual", short: "In context", tip: "Only counts when the prompt is also sexual. The tag is removed and generation continues." },
  { key: "pattern",    short: "Pattern",    tip: "A regular expression. Always blocks, like Always." },
];

const CSS = `
.rmcp-backdrop { position:fixed; inset:0; background:rgba(0,0,0,.6); backdrop-filter:blur(3px);
  z-index:10000; display:flex; align-items:center; justify-content:center; }
.rmcp { width:min(1080px,94vw); height:min(760px,92vh); background:#1b1d21; color:#e7e9ec;
  border:1px solid #33373d; border-radius:14px; display:flex; flex-direction:column;
  overflow:hidden; box-shadow:0 24px 70px rgba(0,0,0,.55);
  font:13px/1.5 ui-sans-serif,-apple-system,"Segoe UI",sans-serif; }

.rmcp-head { padding:16px 20px 0; border-bottom:1px solid #2b2f35; }
.rmcp-titlerow { display:flex; align-items:baseline; gap:12px; margin-bottom:14px; }
.rmcp-title { font-size:15px; font-weight:600; letter-spacing:-.01em; }
.rmcp-sub { color:#868e98; font-size:12px; }
.rmcp-x { margin-left:auto; background:none; border:0; color:#868e98; font-size:20px;
  cursor:pointer; line-height:1; padding:0 2px; }
.rmcp-x:hover { color:#e7e9ec; }

.rmcp-tabs { display:flex; gap:2px; }
.rmcp-tab { background:none; border:0; border-bottom:2px solid transparent; color:#868e98;
  padding:8px 14px 10px; cursor:pointer; font-size:13px; display:flex; align-items:center; gap:7px; }
.rmcp-tab:hover { color:#c8ccd2; }
.rmcp-tab.on { color:#e7e9ec; border-bottom-color:#5b9dff; }
.rmcp-count { font-size:11px; background:#2b3037; color:#9aa2ac; border-radius:9px; padding:1px 7px; }
.rmcp-tab.on .rmcp-count { background:#1e3a63; color:#8fbaff; }

.rmcp-body { flex:1; display:flex; min-height:0; }
.rmcp-left { flex:1.35; display:flex; flex-direction:column; min-width:0; border-right:1px solid #2b2f35; }
.rmcp-right { flex:1; display:flex; flex-direction:column; min-width:0; background:#16181b; }

.rmcp-hint { padding:11px 20px; color:#868e98; font-size:12px; border-bottom:1px solid #2b2f35; }

.rmcp-toolbar { display:flex; gap:8px; align-items:center; padding:11px 20px; border-bottom:1px solid #2b2f35; }
.rmcp-add { flex:1; background:#111316; border:1px solid #33373d; border-radius:8px;
  padding:7px 11px; color:#e7e9ec; font-size:13px; font-family:inherit; }
.rmcp-add:focus { outline:none; border-color:#5b9dff; }
.rmcp-seg { display:flex; border:1px solid #33373d; border-radius:8px; overflow:hidden; }
.rmcp-seg button { background:#111316; border:0; color:#868e98; padding:7px 10px; cursor:pointer;
  font-size:12px; border-right:1px solid #2b2f35; }
.rmcp-seg button:last-child { border-right:0; }
.rmcp-seg button.on { background:#1e3a63; color:#bcd7ff; }
.rmcp-seg button:hover:not(.on) { color:#c8ccd2; }

.rmcp-list { flex:1; overflow-y:auto; padding:8px 12px 16px; }
.rmcp-row { display:flex; align-items:center; gap:9px; padding:5px 8px; border-radius:8px; }
.rmcp-row:hover { background:#212429; }
.rmcp-term { flex:1; background:none; border:0; color:#e7e9ec; font-size:13px;
  font-family:ui-monospace,SFMono-Regular,Menlo,monospace; padding:3px 0; min-width:0; }
.rmcp-term:focus { outline:none; }
.rmcp-row .rmcp-seg { opacity:0; transition:opacity .12s; }
.rmcp-row:hover .rmcp-seg, .rmcp-row.sev-always .rmcp-seg, .rmcp-row.sev-pattern .rmcp-seg { opacity:1; }
.rmcp-seg.mini button { padding:3px 8px; font-size:11px; }
.rmcp-del { background:none; border:0; color:#5c646e; cursor:pointer; font-size:16px;
  opacity:0; padding:0 4px; line-height:1; }
.rmcp-row:hover .rmcp-del { opacity:1; }
.rmcp-del:hover { color:#ff7676; }
.rmcp-dot { width:7px; height:7px; border-radius:50%; flex:0 0 auto; background:#4a525c; }
.rmcp-row.sev-always .rmcp-dot { background:#ff6b6b; }
.rmcp-row.sev-pattern .rmcp-dot { background:#c98bff; }
.rmcp-row.sev-contextual .rmcp-dot { background:#ffb020; }

.rmcp-textarea { flex:1; width:100%; background:#111316; border:0; color:#e7e9ec; padding:14px 20px;
  font:12.5px/1.7 ui-monospace,SFMono-Regular,Menlo,monospace; resize:none; }
.rmcp-textarea:focus { outline:none; }

.rmcp-testhead { padding:13px 18px 10px; font-size:11px; font-weight:600; letter-spacing:.06em;
  text-transform:uppercase; color:#868e98; }
.rmcp-test { margin:0 18px; background:#111316; border:1px solid #33373d; border-radius:9px;
  padding:11px; color:#e7e9ec; font:12.5px/1.6 ui-monospace,SFMono-Regular,Menlo,monospace;
  resize:none; height:118px; }
.rmcp-test:focus { outline:none; border-color:#5b9dff; }
.rmcp-verdict { margin:14px 18px 0; border-radius:9px; padding:12px 14px; border:1px solid; }
.rmcp-verdict.clean     { background:#12281c; border-color:#1f4d33; }
.rmcp-verdict.sanitized { background:#2b2210; border-color:#5c481a; }
.rmcp-verdict.blocked   { background:#2e1416; border-color:#5e2126; }
.rmcp-verdict.idle      { background:#1a1c20; border-color:#2b2f35; color:#5c646e; }
.rmcp-vtop { display:flex; align-items:center; gap:9px; margin-bottom:7px; }
.rmcp-badge { font-size:10px; font-weight:700; letter-spacing:.07em; padding:3px 8px; border-radius:5px; }
.rmcp-badge.clean     { background:#2ecc8a; color:#062616; }
.rmcp-badge.sanitized { background:#ffb020; color:#2a1c00; }
.rmcp-badge.blocked   { background:#ff5c5c; color:#2c0708; }
.rmcp-score { margin-left:auto; font-size:11px; color:#868e98; font-family:ui-monospace,monospace; }
.rmcp-why { font-size:12px; color:#a8b0ba; margin-bottom:9px; }
.rmcp-out { background:#0d0f11; border-radius:7px; padding:9px 11px; font:12px/1.55 ui-monospace,monospace;
  color:#c8ccd2; max-height:150px; overflow:auto; white-space:pre-wrap; word-break:break-word; }
.rmcp-outlabel { font-size:10px; text-transform:uppercase; letter-spacing:.06em; color:#5c646e; margin-bottom:5px; }

.rmcp-foot { display:flex; align-items:center; gap:12px; padding:13px 20px;
  border-top:1px solid #2b2f35; background:#15171a; }
.rmcp-status { flex:1; font-size:12px; color:#868e98; min-width:0; }
.rmcp-status.err { color:#ff8a8a; }
.rmcp-status.ok { color:#5fd6a0; }
.rmcp-btn { border-radius:8px; padding:8px 16px; font-size:13px; cursor:pointer; border:1px solid #3a3f47;
  background:#23272d; color:#e7e9ec; font-family:inherit; }
.rmcp-btn:hover { background:#2b3037; }
.rmcp-btn.primary { background:#3b7de0; border-color:#3b7de0; color:#fff; font-weight:500; }
.rmcp-btn.primary:hover { background:#4a8bee; }
.rmcp-btn:disabled { opacity:.45; cursor:default; }
.rmcp-empty { color:#5c646e; padding:26px 10px; text-align:center; font-size:12.5px; }
`;

class ContentPolicyModal {
  constructor() {
    this.tab = "blocked";
    this.mode = "list";
    this.lists = { blocked: [], allowed: [], context: [] };
    this.dirty = false;
    this.newSeverity = "contextual";
    this.testTimer = null;
  }

  async open() {
    this._injectCss();
    this._build();
    await this._load();
  }

  _injectCss() {
    if (document.getElementById("rmcp-css")) return;
    const s = document.createElement("style");
    s.id = "rmcp-css";
    s.textContent = CSS;
    document.head.appendChild(s);
  }

  _build() {
    this.backdrop = document.createElement("div");
    this.backdrop.className = "rmcp-backdrop";
    this.backdrop.addEventListener("mousedown", (e) => {
      if (e.target === this.backdrop) this._close();
    });

    this.backdrop.innerHTML = `
      <div class="rmcp">
        <div class="rmcp-head">
          <div class="rmcp-titlerow">
            <span class="rmcp-title">Content Policy</span>
            <span class="rmcp-sub">Applied to every prompt before it reaches the sampler</span>
            <button class="rmcp-x" title="Close">&times;</button>
          </div>
          <div class="rmcp-tabs"></div>
        </div>
        <div class="rmcp-body">
          <div class="rmcp-left">
            <div class="rmcp-hint"></div>
            <div class="rmcp-toolbar"></div>
            <div class="rmcp-list"></div>
          </div>
          <div class="rmcp-right">
            <div class="rmcp-testhead">Test a prompt</div>
            <textarea class="rmcp-test" placeholder="Paste a prompt to see what the policy would do with it…"></textarea>
            <div class="rmcp-verdict idle">Type above to preview the outcome.</div>
          </div>
        </div>
        <div class="rmcp-foot">
          <span class="rmcp-status"></span>
          <button class="rmcp-btn cancel">Cancel</button>
          <button class="rmcp-btn primary save" disabled>Save changes</button>
        </div>
      </div>`;

    document.body.appendChild(this.backdrop);
    const $ = (s) => this.backdrop.querySelector(s);
    this.el = {
      tabs: $(".rmcp-tabs"), hint: $(".rmcp-hint"), toolbar: $(".rmcp-toolbar"),
      list: $(".rmcp-list"), test: $(".rmcp-test"), verdict: $(".rmcp-verdict"),
      status: $(".rmcp-status"), save: $(".save"),
    };

    $(".rmcp-x").onclick = () => this._close();
    $(".cancel").onclick = () => this._close();
    this.el.save.onclick = () => this._save();
    this.el.test.addEventListener("input", () => this._scheduleTest());

    this._escHandler = (e) => { if (e.key === "Escape") this._close(); };
    document.addEventListener("keydown", this._escHandler);

    this._renderTabs();
    this._renderToolbar();
  }

  _close() {
    if (this.dirty && !confirm("Discard unsaved changes to the content policy?")) return;
    document.removeEventListener("keydown", this._escHandler);
    this.backdrop.remove();
  }

  async _load() {
    try {
      const r = await api.fetchApi("/api/rmautomation/policy");
      const d = await r.json();
      if (!d.ok) throw new Error(d.error || "failed to load");
      this.lists = d.lists;
      this._renderTabs();
      this._renderList();
      this._setStatus("");
    } catch (err) {
      this._setStatus(`Could not load policy: ${err.message}`, "err");
    }
  }

  _renderTabs() {
    this.el.tabs.innerHTML = "";
    for (const t of TABS) {
      const b = document.createElement("button");
      b.className = "rmcp-tab" + (t.key === this.tab ? " on" : "");
      b.innerHTML = `${t.label}<span class="rmcp-count">${(this.lists[t.key] || []).length}</span>`;
      b.onclick = () => { this.tab = t.key; this._renderTabs(); this._renderToolbar(); this._renderList(); };
      this.el.tabs.appendChild(b);
    }
    this.el.hint.textContent = TABS.find(t => t.key === this.tab).hint;
  }

  _renderToolbar() {
    const tb = this.el.toolbar;
    tb.innerHTML = "";

    const input = document.createElement("input");
    input.className = "rmcp-add";
    input.placeholder = this.tab === "blocked" ? "Add a term, then Enter…"
      : this.tab === "allowed" ? "Add an exempt phrase, then Enter…"
      : "Add a context word, then Enter…";
    input.addEventListener("keydown", (e) => {
      if (e.key !== "Enter") return;
      const v = input.value.trim();
      if (!v) return;
      const item = this.tab === "blocked" ? { value: v, severity: this.newSeverity } : { value: v };
      this.lists[this.tab].unshift(item);
      input.value = "";
      this._touch();
      this._renderTabs(); this._renderList();
    });
    tb.appendChild(input);

    if (this.tab === "blocked") {
      tb.appendChild(this._severitySeg(this.newSeverity, (k) => {
        this.newSeverity = k; this._renderToolbar(); input.focus();
      }));
    }

    const modeSeg = document.createElement("div");
    modeSeg.className = "rmcp-seg";
    for (const m of ["list", "text"]) {
      const b = document.createElement("button");
      b.textContent = m === "list" ? "List" : "Text";
      b.className = this.mode === m ? "on" : "";
      b.title = m === "text" ? "Edit as raw lines — useful for pasting many at once" : "Edit as rows";
      b.onclick = () => { this.mode = m; this._renderToolbar(); this._renderList(); };
      modeSeg.appendChild(b);
    }
    tb.appendChild(modeSeg);
  }

  _severitySeg(current, onPick, mini) {
    const seg = document.createElement("div");
    seg.className = "rmcp-seg" + (mini ? " mini" : "");
    for (const s of SEVERITIES) {
      const b = document.createElement("button");
      b.textContent = s.short;
      b.title = s.tip;
      b.className = current === s.key ? "on" : "";
      b.onclick = (e) => { e.stopPropagation(); onPick(s.key); };
      seg.appendChild(b);
    }
    return seg;
  }

  _renderList() {
    const host = this.el.list;
    host.innerHTML = "";
    const items = this.lists[this.tab] || [];

    if (this.mode === "text") {
      const ta = document.createElement("textarea");
      ta.className = "rmcp-textarea";
      ta.value = items.map(i =>
        this.tab === "blocked"
          ? (i.severity === "always" ? "!" + i.value : i.severity === "pattern" ? "/" + i.value + "/" : i.value)
          : i.value
      ).join("\n");
      ta.addEventListener("input", () => {
        const lines = ta.value.split("\n").map(l => l.trim()).filter(Boolean);
        this.lists[this.tab] = lines.map(l => {
          if (this.tab !== "blocked") return { value: l };
          if (l.startsWith("!")) return { value: l.slice(1).trim(), severity: "always" };
          if (l.length > 1 && l.startsWith("/") && l.endsWith("/")) return { value: l.slice(1, -1), severity: "pattern" };
          return { value: l, severity: "contextual" };
        });
        this._touch(); this._renderTabs();
      });
      host.appendChild(ta);
      return;
    }

    if (!items.length) {
      host.innerHTML = `<div class="rmcp-empty">Nothing here yet. Add an entry above.</div>`;
      return;
    }

    items.forEach((item, idx) => {
      const row = document.createElement("div");
      row.className = "rmcp-row" + (this.tab === "blocked" ? ` sev-${item.severity || "contextual"}` : "");

      const dot = document.createElement("span");
      dot.className = "rmcp-dot";
      row.appendChild(dot);

      const term = document.createElement("input");
      term.className = "rmcp-term";
      term.value = item.value;
      term.addEventListener("input", () => { item.value = term.value; this._touch(); });
      row.appendChild(term);

      if (this.tab === "blocked") {
        row.appendChild(this._severitySeg(item.severity || "contextual", (k) => {
          item.severity = k; this._touch(); this._renderList();
        }, true));
      }

      const del = document.createElement("button");
      del.className = "rmcp-del";
      del.innerHTML = "&times;";
      del.title = "Remove";
      del.onclick = () => { items.splice(idx, 1); this._touch(); this._renderTabs(); this._renderList(); };
      row.appendChild(del);

      host.appendChild(row);
    });
  }

  _touch() {
    this.dirty = true;
    this.el.save.disabled = false;
    this._setStatus("Unsaved changes");
    this._scheduleTest();
  }

  _scheduleTest() {
    clearTimeout(this.testTimer);
    this.testTimer = setTimeout(() => this._runTest(), 280);
  }

  async _runTest() {
    const text = this.el.test.value.trim();
    if (!text) {
      this.el.verdict.className = "rmcp-verdict idle";
      this.el.verdict.textContent = "Type above to preview the outcome.";
      return;
    }
    // Saved rules are what the server enforces, so say so rather than implying the
    // draft is live. Testing an unsaved draft would need the server to accept a full
    // ruleset per keystroke; showing the truth is better than faking it.
    try {
      const r = await api.fetchApi("/api/rmautomation/policy/test", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ text }),
      });
      const d = await r.json();
      if (!d.ok) throw new Error(d.error || "test failed");

      const label = d.status.toUpperCase();
      const why = d.reasons?.length ? d.reasons.join(" · ") : "no rules matched";
      const showOut = d.status !== "clean";
      this.el.verdict.className = `rmcp-verdict ${d.status}`;
      this.el.verdict.innerHTML = `
        <div class="rmcp-vtop">
          <span class="rmcp-badge ${d.status}">${label}</span>
          <span class="rmcp-score">score ${d.score} / threshold ${d.threshold}</span>
        </div>
        <div class="rmcp-why">${escapeHtml(why)}${this.dirty ? " — against <b>saved</b> rules; save to test your edits" : ""}</div>
        ${showOut ? `<div class="rmcp-outlabel">${d.status === "blocked" ? "Replaced with" : "Generates from"}</div>
        <div class="rmcp-out">${escapeHtml(d.result)}</div>` : ""}`;
    } catch (err) {
      this.el.verdict.className = "rmcp-verdict blocked";
      this.el.verdict.textContent = `Test failed: ${err.message}`;
    }
  }

  async _save() {
    this.el.save.disabled = true;
    this._setStatus("Saving…");
    try {
      const r = await api.fetchApi("/api/rmautomation/policy", {
        method: "PUT",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(this.lists),
      });
      const d = await r.json();
      if (!d.ok) throw new Error(d.error || "save failed");
      this.dirty = false;
      const n = Object.values(d.counts).reduce((a, b) => a + b, 0);
      this._setStatus(`Saved — ${n} entries active`, "ok");
      this._runTest();
    } catch (err) {
      // Validation failures land here and are the whole reason the message is shown
      // verbatim: "your allowlist would disable a hard rule" needs to be readable.
      this._setStatus(err.message, "err");
      this.el.save.disabled = false;
    }
  }

  _setStatus(msg, kind) {
    this.el.status.className = "rmcp-status" + (kind ? " " + kind : "");
    this.el.status.textContent = msg;
  }
}

function escapeHtml(s) {
  return String(s).replace(/[&<>"]/g, c => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;" }[c]));
}

app.registerExtension({
  name: "RMAutomation.ContentPolicy",
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (!["RMStylesFull", "RMStylesFullDisplay", "RMStylesPipe"].includes(nodeData.name)) return;
    const onNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      onNodeCreated?.apply(this, arguments);
      this.addWidget("button", "Content Policy", null, () => new ContentPolicyModal().open());
    };
  },
});