From 8db0dd324673b950b7017a74078c9659228b8e84 Mon Sep 17 00:00:00 2001
From: Ben Senescu <44480372+bensenescu@users.noreply.github.com>
Date: Wed, 26 Aug 2026 09:59:35 -0400
Subject: [PATCH] First site audit over MCP fails with a raw billing error
 (#525)

---
 .../services/AuditService.limitTier.test.ts   | 37 +++++++++++++++----
 .../features/audit/services/AuditService.ts   | 11 ++++--
 src/server/mcp/tools/site-audit-tools.ts      |  4 +-
 src/serverFunctions/audit.ts                  |  4 +-
 4 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/src/server/features/audit/services/AuditService.limitTier.test.ts b/src/server/features/audit/services/AuditService.limitTier.test.ts
index c58e2b4..ff1c918 100644
--- a/src/server/features/audit/services/AuditService.limitTier.test.ts
+++ b/src/server/features/audit/services/AuditService.limitTier.test.ts
@@ -1,12 +1,16 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
-const { isHostedMock, hasManagedAccessMock, hasPaidPlanMock } = vi.hoisted(
-  () => ({
-    isHostedMock: vi.fn(),
-    hasManagedAccessMock: vi.fn(),
-    hasPaidPlanMock: vi.fn(),
-  }),
-);
+const {
+  isHostedMock,
+  hasManagedAccessMock,
+  hasPaidPlanMock,
+  getOrCreateCustomerMock,
+} = vi.hoisted(() => ({
+  isHostedMock: vi.fn(),
+  hasManagedAccessMock: vi.fn(),
+  hasPaidPlanMock: vi.fn(),
+  getOrCreateCustomerMock: vi.fn(),
+}));
 
 vi.mock("cloudflare:workers", () => ({ env: {} }));
 vi.mock("@/server/lib/runtime-env", () => ({
@@ -15,6 +19,7 @@ vi.mock("@/server/lib/runtime-env", () => ({
 vi.mock("@/server/billing/subscription", () => ({
   customerHasManagedAccess: hasManagedAccessMock,
   customerHasPaidPlan: hasPaidPlanMock,
+  getOrCreateOrganizationCustomer: getOrCreateCustomerMock,
 }));
 vi.mock("@/server/features/audit/repositories/AuditRepository", () => ({
   AuditRepository: {},
@@ -26,6 +31,12 @@ vi.mock("@/server/lib/audit/progress-kv", () => ({ AuditProgressKV: {} }));
 
 import { AuditService } from "@/server/features/audit/services/AuditService";
 
+const customer = {
+  organizationId: "org-1",
+  userEmail: "user@example.com",
+  userId: "user-1",
+};
+
 describe("resolveAuditLimitTier", () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -36,10 +47,20 @@ describe("resolveAuditLimitTier", () => {
   it("uses the uncapped self-hosted tier without consulting billing", async () => {
     isHostedMock.mockResolvedValue(false);
 
-    await expect(AuditService.resolveAuditLimitTier("org-1")).resolves.toBe(
+    await expect(AuditService.resolveAuditLimitTier(customer)).resolves.toBe(
       "self_hosted",
     );
+    expect(getOrCreateCustomerMock).not.toHaveBeenCalled();
     expect(hasManagedAccessMock).not.toHaveBeenCalled();
     expect(hasPaidPlanMock).not.toHaveBeenCalled();
   });
+
+  it("ensures the Autumn customer before checking entitlements in hosted mode", async () => {
+    isHostedMock.mockResolvedValue(true);
+
+    await expect(AuditService.resolveAuditLimitTier(customer)).resolves.toBe(
+      "paid",
+    );
+    expect(getOrCreateCustomerMock).toHaveBeenCalledWith(customer);
+  });
 });
diff --git a/src/server/features/audit/services/AuditService.ts b/src/server/features/audit/services/AuditService.ts
index e880a20..1610b38 100644
--- a/src/server/features/audit/services/AuditService.ts
+++ b/src/server/features/audit/services/AuditService.ts
@@ -2,6 +2,7 @@ import { env } from "cloudflare:workers";
 import {
   customerHasManagedAccess,
   customerHasPaidPlan,
+  getOrCreateOrganizationCustomer,
   type BillingCustomerContext,
 } from "@/server/billing/subscription";
 import { AuditRepository } from "@/server/features/audit/repositories/AuditRepository";
@@ -29,12 +30,16 @@ import { isHostedServerAuthMode } from "@/server/lib/runtime-env";
 // small audit at a time, paid keeps the full limits, and customers with no
 // Autumn product at all are turned away. Self-hosted isn't gated.
 async function resolveAuditLimitTier(
-  organizationId: string,
+  customer: BillingCustomerContext,
 ): Promise<AuditLimitTier> {
   if (!(await isHostedServerAuthMode())) return "self_hosted";
+  // An org minted outside a billing path (better-auth hooks, MCP auth) has no
+  // Autumn customer yet, and `check` 404s instead of reporting no access — a
+  // brand-new MCP user's first audit failed with a raw billing error.
+  await getOrCreateOrganizationCustomer(customer);
   const [hasManagedAccess, hasPaidPlan] = await Promise.all([
-    customerHasManagedAccess(organizationId),
-    customerHasPaidPlan(organizationId),
+    customerHasManagedAccess(customer.organizationId),
+    customerHasPaidPlan(customer.organizationId),
   ]);
   if (!hasManagedAccess) {
     throw new AppError("PAYMENT_REQUIRED", "Subscribe to run site audits");
diff --git a/src/server/mcp/tools/site-audit-tools.ts b/src/server/mcp/tools/site-audit-tools.ts
index 4862dc4..8c58e1d 100644
--- a/src/server/mcp/tools/site-audit-tools.ts
+++ b/src/server/mcp/tools/site-audit-tools.ts
@@ -87,9 +87,7 @@ export const runSiteAuditTool = {
     // many-minute wait, which chat agents handle badly. The app UI passes its
     // own explicit lighthouseStrategy, so this default only governs agents.
     const lighthouseStrategy = (args.runLighthouse ?? false) ? "auto" : "none";
-    const limitTier = await AuditService.resolveAuditLimitTier(
-      context.auth.organizationId,
-    );
+    const limitTier = await AuditService.resolveAuditLimitTier(context.billing);
     let auditId: string;
     try {
       ({ auditId } = await AuditService.startAudit({
diff --git a/src/serverFunctions/audit.ts b/src/serverFunctions/audit.ts
index eb56fef..41eefcd 100644
--- a/src/serverFunctions/audit.ts
+++ b/src/serverFunctions/audit.ts
@@ -16,9 +16,7 @@ export const startAudit = createServerFn({ method: "POST" })
   .middleware(requireProjectContext)
   .validator(startAuditSchema)
   .handler(async ({ data, context }) => {
-    const limitTier = await AuditService.resolveAuditLimitTier(
-      context.organizationId,
-    );
+    const limitTier = await AuditService.resolveAuditLimitTier(context);
 
     const result = await AuditService.startAudit({
       actorUserId: context.userId,