@checkstack/ai-backend 0.1.0

package/src/tools/rank-docs.test.ts

@@ -0,0 +1,153 @@
+import { describe, expect, test } from "bun:test";
+import type { DocsIndexEntry } from "../generated/docs-index";
+import {
+  SNIPPET_MAX_CHARS,
+  buildSnippet,
+  rankDocs,
+  tokenize,
+} from "./rank-docs";
+
+function entry(over: Partial<DocsIndexEntry> & { slug: string }): DocsIndexEntry {
+  return {
+    title: over.title ?? over.slug,
+    headings: over.headings ?? [],
+    content: over.content ?? "",
+    truncated: over.truncated ?? false,
+    ...(over.description ? { description: over.description } : {}),
+    slug: over.slug,
+  };
+}
+
+describe("tokenize", () => {
+  test("lowercases, splits on non-alphanumerics, drops <2-char tokens", () => {
+    expect(tokenize("Health-Check a HTTP 200!")).toEqual([
+      "health",
+      "check",
+      "http",
+      "200",
+    ]);
+  });
+});
+
+describe("buildSnippet", () => {
+  test("returns the whole content when under the budget", () => {
+    expect(buildSnippet({ content: "short text", queryTerms: ["text"] })).toBe(
+      "short text",
+    );
+  });
+
+  test("windows around the first query-term match with ellipses", () => {
+    const content = "x".repeat(600) + " NEEDLE " + "y".repeat(600);
+    const snippet = buildSnippet({
+      content,
+      queryTerms: ["needle"],
+      maxChars: 100,
+    });
+    expect(snippet.toLowerCase()).toContain("needle");
+    expect(snippet.length).toBeLessThanOrEqual(100 + 2); // + ellipses
+    expect(snippet.startsWith("…")).toBe(true);
+  });
+
+  test("falls back to the head when no term occurs inline", () => {
+    const content = "a".repeat(1000);
+    const snippet = buildSnippet({
+      content,
+      queryTerms: ["zzz"],
+      maxChars: 80,
+    });
+    expect(snippet.endsWith("…")).toBe(true);
+    expect(snippet.length).toBeLessThanOrEqual(81);
+  });
+});
+
+describe("rankDocs", () => {
+  const index: DocsIndexEntry[] = [
+    entry({
+      slug: "title-hit",
+      title: "Health checks",
+      content: "Unrelated body about widgets.",
+    }),
+    entry({
+      slug: "heading-hit",
+      title: "Operations",
+      headings: ["Configuring health checks"],
+      content: "Body about operating things.",
+    }),
+    entry({
+      slug: "content-hit",
+      title: "Misc",
+      content: "Somewhere in here we mention a health check probe.",
+    }),
+    entry({
+      slug: "no-hit",
+      title: "Billing",
+      content: "Invoices and payments.",
+    }),
+  ];
+
+  test("returns [] for an empty / whitespace-only query", () => {
+    expect(rankDocs({ index, query: "", limit: 5 })).toEqual([]);
+    expect(rankDocs({ index, query: "   ", limit: 5 })).toEqual([]);
+  });
+
+  test("returns [] when limit <= 0", () => {
+    expect(rankDocs({ index, query: "health", limit: 0 })).toEqual([]);
+  });
+
+  test("excludes pages with no matching term", () => {
+    const slugs = rankDocs({ index, query: "health check", limit: 10 }).map(
+      (h) => h.slug,
+    );
+    expect(slugs).not.toContain("no-hit");
+  });
+
+  test("title match outranks heading match outranks content match", () => {
+    const hits = rankDocs({ index, query: "health check", limit: 10 });
+    const order = hits.map((h) => h.slug);
+    expect(order.indexOf("title-hit")).toBeLessThan(
+      order.indexOf("heading-hit"),
+    );
+    expect(order.indexOf("heading-hit")).toBeLessThan(
+      order.indexOf("content-hit"),
+    );
+  });
+
+  test("honours the limit cap", () => {
+    const hits = rankDocs({ index, query: "health check", limit: 2 });
+    expect(hits.length).toBe(2);
+  });
+
+  test("attaches the matching heading when the hit is a sub-section", () => {
+    const hit = rankDocs({ index, query: "configuring", limit: 5 }).find(
+      (h) => h.slug === "heading-hit",
+    );
+    expect(hit?.heading).toBe("Configuring health checks");
+  });
+
+  test("snippets never exceed the budget", () => {
+    const big = entry({
+      slug: "big",
+      title: "Big",
+      content: ("health check ".repeat(500)).trim(),
+    });
+    const hit = rankDocs({ index: [big], query: "health", limit: 1 })[0];
+    expect(hit?.snippet.length).toBeLessThanOrEqual(SNIPPET_MAX_CHARS + 2);
+  });
+
+  test("ordering is stable / deterministic across runs", () => {
+    const a = rankDocs({ index, query: "health check", limit: 10 });
+    const b = rankDocs({ index, query: "health check", limit: 10 });
+    expect(a.map((h) => h.slug)).toEqual(b.map((h) => h.slug));
+  });
+
+  test("ties break by slug (deterministic ordering)", () => {
+    const tied: DocsIndexEntry[] = [
+      entry({ slug: "zebra", title: "Foo", content: "foo" }),
+      entry({ slug: "alpha", title: "Foo", content: "foo" }),
+    ];
+    const order = rankDocs({ index: tied, query: "foo", limit: 5 }).map(
+      (h) => h.slug,
+    );
+    expect(order).toEqual(["alpha", "zebra"]);
+  });
+});

package/src/tools/rank-docs.ts

@@ -0,0 +1,209 @@
+import type { DocsIndexEntry } from "../generated/docs-index";
+
+/**
+ * One ranked documentation hit produced by {@link rankDocs}. Mirrors the
+ * `DocHitSchema` wire shape in `@checkstack/ai-common` (plan §2.5) — kept as a
+ * plain interface here so the ranking module is a PURE, dependency-light,
+ * unit-testable function over the bundled index.
+ */
+export interface RankedDocHit {
+  slug: string;
+  title: string;
+  /** Section heading the snippet came from, when the best match is in a heading. */
+  heading?: string;
+  /** Bounded snippet around the best match, highlighting why the page matched. */
+  snippet: string;
+  /** BM25-ish relevance score (opaque ordering hint). */
+  score: number;
+}
+
+/** Max characters in a returned snippet (plan §3.4 size budget). */
+export const SNIPPET_MAX_CHARS = 500;
+
+/**
+ * Field weights for the BM25-ish term-frequency score. A query term in the
+ * title outranks the same term in a heading, which outranks one in the body.
+ */
+const TITLE_WEIGHT = 8;
+const HEADING_WEIGHT = 4;
+const DESCRIPTION_WEIGHT = 3;
+const CONTENT_WEIGHT = 1;
+
+/**
+ * BM25 term-frequency saturation constant. Diminishing returns past a few
+ * occurrences of a term, so a page that merely repeats a word does not dominate
+ * a page that genuinely covers the topic.
+ */
+const BM25_K1 = 1.5;
+
+/** Lowercases and splits text into alphanumeric tokens (length >= 2). */
+export function tokenize(text: string): string[] {
+  return text
+    .toLowerCase()
+    .split(/[^a-z0-9]+/)
+    .filter((t) => t.length >= 2);
+}
+
+/** Counts occurrences of each token in a token list. */
+function termCounts(tokens: string[]): Map<string, number> {
+  const counts = new Map<string, number>();
+  for (const t of tokens) counts.set(t, (counts.get(t) ?? 0) + 1);
+  return counts;
+}
+
+/**
+ * BM25-style term-frequency saturation: tf * (k1 + 1) / (tf + k1). Bounded and
+ * monotonic, so more occurrences help but with diminishing returns.
+ */
+function saturate(tf: number): number {
+  if (tf <= 0) return 0;
+  return (tf * (BM25_K1 + 1)) / (tf + BM25_K1);
+}
+
+interface ScorableEntry {
+  entry: DocsIndexEntry;
+  titleTokens: Map<string, number>;
+  headingTokensByHeading: { heading: string; counts: Map<string, number> }[];
+  descriptionTokens: Map<string, number>;
+  contentTokens: Map<string, number>;
+}
+
+function prepare(entry: DocsIndexEntry): ScorableEntry {
+  return {
+    entry,
+    titleTokens: termCounts(tokenize(entry.title)),
+    headingTokensByHeading: entry.headings.map((heading) => ({
+      heading,
+      counts: termCounts(tokenize(heading)),
+    })),
+    descriptionTokens: termCounts(tokenize(entry.description ?? "")),
+    contentTokens: termCounts(tokenize(entry.content)),
+  };
+}
+
+/**
+ * Picks the best-matching heading for a query (the heading containing the most
+ * query terms, by saturated term frequency). Returns `undefined` when no
+ * heading matches any query term.
+ */
+function bestHeading(
+  scorable: ScorableEntry,
+  queryTerms: string[],
+): string | undefined {
+  let best: string | undefined;
+  let bestScore = 0;
+  for (const { heading, counts } of scorable.headingTokensByHeading) {
+    let score = 0;
+    for (const term of queryTerms) score += saturate(counts.get(term) ?? 0);
+    if (score > bestScore) {
+      bestScore = score;
+      best = heading;
+    }
+  }
+  return best;
+}
+
+/**
+ * Builds a bounded snippet around the first occurrence of any query term in the
+ * content; falls back to the content head when no term is found inline (e.g. the
+ * match was purely in the title). Whitespace is collapsed for compactness.
+ */
+export function buildSnippet({
+  content,
+  queryTerms,
+  maxChars = SNIPPET_MAX_CHARS,
+}: {
+  content: string;
+  queryTerms: string[];
+  maxChars?: number;
+}): string {
+  const normalized = content.replaceAll(/\s+/g, " ").trim();
+  if (normalized.length <= maxChars) return normalized;
+
+  const lower = normalized.toLowerCase();
+  let matchAt = -1;
+  for (const term of queryTerms) {
+    const idx = lower.indexOf(term);
+    if (idx !== -1 && (matchAt === -1 || idx < matchAt)) matchAt = idx;
+  }
+
+  if (matchAt === -1) {
+    return normalized.slice(0, maxChars).trimEnd() + "…";
+  }
+
+  // Center the window on the match, then clamp to the content bounds.
+  const half = Math.floor(maxChars / 2);
+  let start = Math.max(0, matchAt - half);
+  const end = Math.min(normalized.length, start + maxChars);
+  start = Math.max(0, end - maxChars);
+
+  const prefix = start > 0 ? "…" : "";
+  const suffix = end < normalized.length ? "…" : "";
+  return prefix + normalized.slice(start, end).trim() + suffix;
+}
+
+/**
+ * Pure BM25-ish ranking over the bundled docs index. Tokenizes the query and
+ * scores each entry by saturated term frequency across title (boosted),
+ * headings (boosted), description, and content; returns the top-`limit` hits
+ * with a bounded snippet around the best match.
+ *
+ * Deterministic: ties break by score then slug, so the order is stable.
+ * Returns `[]` for an empty/whitespace query or when nothing matches.
+ */
+export function rankDocs({
+  index,
+  query,
+  limit,
+}: {
+  index: readonly DocsIndexEntry[];
+  query: string;
+  limit: number;
+}): RankedDocHit[] {
+  const queryTerms = [...new Set(tokenize(query))];
+  if (queryTerms.length === 0 || limit <= 0) return [];
+
+  const scored = index.map((entry) => {
+    const scorable = prepare(entry);
+    let score = 0;
+    for (const term of queryTerms) {
+      score += TITLE_WEIGHT * saturate(scorable.titleTokens.get(term) ?? 0);
+      score +=
+        DESCRIPTION_WEIGHT * saturate(scorable.descriptionTokens.get(term) ?? 0);
+      score += CONTENT_WEIGHT * saturate(scorable.contentTokens.get(term) ?? 0);
+      // Headings: sum the term's saturated tf across every heading.
+      let headingTf = 0;
+      for (const { counts } of scorable.headingTokensByHeading) {
+        headingTf += counts.get(term) ?? 0;
+      }
+      score += HEADING_WEIGHT * saturate(headingTf);
+    }
+    return { scorable, score };
+  });
+
+  return scored
+    .filter((s) => s.score > 0)
+    .toSorted((a, b) =>
+      b.score === a.score
+        ? a.scorable.entry.slug < b.scorable.entry.slug
+          ? -1
+          : a.scorable.entry.slug > b.scorable.entry.slug
+            ? 1
+            : 0
+        : b.score - a.score,
+    )
+    .slice(0, limit)
+    .map(({ scorable, score }) => {
+      const heading = bestHeading(scorable, queryTerms);
+      return {
+        slug: scorable.entry.slug,
+        title: scorable.entry.title,
+        ...(heading ? { heading } : {}),
+        snippet: buildSnippet({
+          content: scorable.entry.content,
+          queryTerms,
+        }),
+        score,
+      };
+    });
+}

package/src/tools/script-context-extract.test.ts

@@ -0,0 +1,93 @@
+import { describe, expect, test } from "bun:test";
+import { SDK_EDITOR_BUNDLE_DTS } from "@checkstack/sdk/editor-bundle";
+import {
+  resolveScriptContext,
+  extractDeclareModuleBlock,
+  ScriptContextExtractionError,
+  HEALTHCHECK_SHELL_ENV,
+  AUTOMATION_SHELL_ENV,
+} from "./script-context-extract";
+
+describe("extractDeclareModuleBlock (pure)", () => {
+  test("returns the full declare-module block for a known module", () => {
+    const block = extractDeclareModuleBlock({
+      bundle: SDK_EDITOR_BUNDLE_DTS,
+      moduleName: "@checkstack/sdk/healthcheck",
+    });
+    expect(block.startsWith('declare module "@checkstack/sdk/healthcheck"')).toBe(
+      true,
+    );
+    expect(block.trimEnd().endsWith("}")).toBe(true);
+    // It must contain the helper + context types, and NOT bleed into the next module.
+    expect(block).toContain("defineHealthCheck");
+    expect(block).toContain("HealthCheckScriptContext");
+    expect(block).not.toContain("defineIntegration");
+  });
+
+  test("brace-matches a nested block correctly (no early termination)", () => {
+    const bundle =
+      'declare module "x" {\n  interface A { readonly b: { readonly c: string } }\n}\ndeclare module "y" { export const z: number; }';
+    const block = extractDeclareModuleBlock({ bundle, moduleName: "x" });
+    expect(block).toContain("readonly c: string");
+    expect(block).not.toContain("export const z");
+  });
+
+  test("throws a clear error for a missing module", () => {
+    expect(() =>
+      extractDeclareModuleBlock({
+        bundle: SDK_EDITOR_BUNDLE_DTS,
+        moduleName: "@checkstack/sdk/does-not-exist",
+      }),
+    ).toThrow(ScriptContextExtractionError);
+  });
+});
+
+describe("resolveScriptContext (pure)", () => {
+  test("healthcheck-script extracts the healthcheck module + helper", () => {
+    const resolved = resolveScriptContext({
+      context: "healthcheck-script",
+      bundle: SDK_EDITOR_BUNDLE_DTS,
+    });
+    expect(resolved.language).toBe("typescript");
+    expect(resolved.sdkModule).toBe("@checkstack/sdk/healthcheck");
+    expect(resolved.helper).toBe("defineHealthCheck");
+    expect(resolved.declarations).toContain("HealthCheckScriptResult");
+    expect(resolved.shellEnv).toBeUndefined();
+    expect(resolved.allowsManagedPackages).toBe(true);
+    expect(resolved.starterExample).toContain("defineHealthCheck");
+  });
+
+  test("automation-action-script extracts the integration module + helper", () => {
+    const resolved = resolveScriptContext({
+      context: "automation-action-script",
+      bundle: SDK_EDITOR_BUNDLE_DTS,
+    });
+    expect(resolved.sdkModule).toBe("@checkstack/sdk/integration");
+    expect(resolved.helper).toBe("defineIntegration");
+    expect(resolved.declarations).toContain("IntegrationScriptContext");
+    expect(resolved.declarations).not.toContain("defineHealthCheck");
+  });
+
+  test("healthcheck-shell returns the env table, no SDK module", () => {
+    const resolved = resolveScriptContext({
+      context: "healthcheck-shell",
+      bundle: SDK_EDITOR_BUNDLE_DTS,
+    });
+    expect(resolved.language).toBe("shell");
+    expect(resolved.sdkModule).toBeUndefined();
+    expect(resolved.helper).toBeUndefined();
+    expect(resolved.shellEnv).toBe(HEALTHCHECK_SHELL_ENV);
+    expect(resolved.declarations).toContain("CHECKSTACK_CHECK_ID");
+    expect(resolved.allowsManagedPackages).toBe(false);
+  });
+
+  test("automation-action-shell returns the automation env table", () => {
+    const resolved = resolveScriptContext({
+      context: "automation-action-shell",
+      bundle: SDK_EDITOR_BUNDLE_DTS,
+    });
+    expect(resolved.language).toBe("shell");
+    expect(resolved.shellEnv).toBe(AUTOMATION_SHELL_ENV);
+    expect(resolved.declarations).toContain("CHECKSTACK_EVENT_ID");
+  });
+});