Add a token count linter

.remarkrc.mjs

@@ -9,6 +9,7 @@ import {
   getVersionRootPath,
 } from "./.remark-build/server/docs-helpers.mjs";
 import { remarkLintPageStructure } from "./.remark-build/server/lint-page-structure.mjs";
+import { remarkLintTokenCount } from "./.remark-build/server/lint-token-count.mjs";
 import { loadConfig } from "./.remark-build/server/config-docs.mjs";
 import { updatePathsInIncludes } from "./.remark-build/server/asset-path-helpers.mjs";
 import * as yaml from "yaml";
@@ -83,6 +84,7 @@ const configLint = {
     ["validate-links", { repository: false }],
     [remarkLintTeleportDocsLinks],
     [remarkLintPageStructure],
+    [remarkLintTokenCount],
     // Disabling the remarkLintFrontmatter check until we fix
     // gravitational/docs#80
     // [remarkLintFrontmatter, ["error"]],

server/lint-token-count.test.ts

@@ -0,0 +1,39 @@
+import { describe, expect, test } from "@jest/globals";
+import { remarkLintTokenCount } from "./lint-token-count";
+import { remark } from "remark";
+import mdx from "remark-mdx";
+
+const getReasons = (value: string) => {
+  return remark()
+    .use(mdx as any)
+    .use(remarkLintTokenCount as any)
+    .processSync(value)
+    .messages.map((m) => m.reason);
+};
+
+describe("server/lint-token-count", () => {
+  interface testCase {
+    name: string;
+    content: string;
+    expected: Array<string>;
+  }
+
+  const testCases: Array<testCase> = [
+    {
+      name: "should warn when token count exceeds limit",
+      content: "tokens".repeat(30001),
+      expected: [
+        "This page has an estimated token count of 30001, which exceeds the limit of 30000. Consider splitting the page into smaller sections. Disable this warning by adding {/* lint ignore token-count remark-lint */} in the end of the file.",
+      ],
+    },
+    {
+      name: "should not warn when token count is within limit",
+      content: "tokens".repeat(1000),
+      expected: [],
+    },
+  ];
+
+  test.each(testCases)("$name", (tc) => {
+    expect(getReasons(tc.content)).toEqual(tc.expected);
+  });
+});

server/lint-token-count.ts

@@ -0,0 +1,33 @@
+import { lintRule } from "unified-lint-rule";
+import { toMarkdown } from "mdast-util-to-markdown";
+import { mdxToMarkdown } from "mdast-util-mdx";
+import { frontmatterToMarkdown } from "mdast-util-frontmatter";
+import { gfmToMarkdown } from "mdast-util-gfm";
+import { estimateTokenCount } from "tokenx";
+import type { Node, Parent } from "unist";
+
+const tokenLimit = 30000;
+const messageSuffix = `Disable this warning by adding {/* lint ignore token-count remark-lint */} in the end of the file.`;
+
+export const remarkLintTokenCount = lintRule(
+  "remark-lint:token-count",
+  (root: Node, vfile) => {
+    const children = (root as Parent).children;
+    const last = children[children.length - 1];
+
+    const pageMarkdownContent = toMarkdown(root as any, {
+      extensions: [
+        mdxToMarkdown(),
+        frontmatterToMarkdown(["yaml"]),
+        gfmToMarkdown(),
+      ],
+    });
+    const tokenEstimation = estimateTokenCount(pageMarkdownContent);
+    if (tokenEstimation > tokenLimit) {
+      vfile.message(
+        `This page has an estimated token count of ${tokenEstimation}, which exceeds the limit of ${tokenLimit}. Consider splitting the page into smaller sections. ${messageSuffix}`,
+        last.position,
+      );
+    }
+  },
+);