diff --git a/.remarkrc.mjs b/.remarkrc.mjs index cf5f096e..9f130b2f 100644 --- a/.remarkrc.mjs +++ b/.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"]], diff --git a/server/lint-token-count.test.ts b/server/lint-token-count.test.ts new file mode 100644 index 00000000..ae1514aa --- /dev/null +++ b/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; + } + + const testCases: Array = [ + { + 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); + }); +}); diff --git a/server/lint-token-count.ts b/server/lint-token-count.ts new file mode 100644 index 00000000..ee64a7f0 --- /dev/null +++ b/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, + ); + } + }, +);