diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 35a78b9..91c4678 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -68,6 +68,8 @@ Rules for test-only code (see README, "Test-only switches"): is off without `CLEANSTACK_TEST_MODE=1` and when `VERCEL=1`. - Server code reads the user through `@/lib/auth`; eslint rejects `@clerk/nextjs/server` imports elsewhere, so the auth bypass keeps a single entry point. +- Every AI call is preceded by `checkAiBudget(teamId, calls)` and followed by an awaited + `meterBedrockCall(...)` (`src/lib/bedrock-meter.ts`), so the guest and team caps see it. - AI calls take their model from `languageModel()` (`src/lib/ai-model.ts`), never `bedrock()` directly, so CI never reaches Bedrock. - Playwright's `webServer` inherits your shell's environment, and shell variables win over diff --git a/README.md b/README.md index 76fa8a3..103c03d 100644 --- a/README.md +++ b/README.md @@ -238,6 +238,7 @@ aws lambda update-function-code --function-name cleanstack-executor \ | Startup log `invalid or missing environment` | Set the listed variables in Vercel and redeploy. | | Rotating `WEBHOOK_SECRET` | Update the Vercel env **and** the profiler and ai-trigger Lambdas together. Requests fail with 401 while they differ. | | AI spend | `GET /api/admin/ai-spend` with `x-admin-secret` lists the current month's estimated spend per team. | +| Guests see "Guest AI capacity is used up for today" | All guests together reached $5 of estimated Bedrock spend since 00:00 UTC. It resets at midnight UTC; the number is `GUEST_LIMITS.aiSpendAllGuestsPerDayUsd`. | --- @@ -253,8 +254,22 @@ What the code does today. Settings that live only in AWS are listed as such. - After a successful execution, the executor deletes every object version under the run's raw prefix (`{user}/{pipeline}/{run}/`), including `extracted_text.txt` (`auto_delete_raw`, on by default). - Raw files of **failed** runs are not deleted, and there are no S3 lifecycle rules in the repo. - `DELETE /api/account?confirm=true` purges every version and delete marker under the user's raw prefix and their pipelines' processed prefixes, then deletes the DB rows. If S3 fails it aborts before touching the DB. -- **AI.** Uploaded content is sent to Amazon Bedrock and wrapped in `` tags with an instruction to treat it as data. This reduces prompt-injection risk but does not prevent it, which is why rules need approval (or the committee in auto mode). The monthly AI spend cap is checked in `suggest-transforms` only. +- **AI.** Uploaded content is sent to Amazon Bedrock and wrapped in `` tags with an instruction to treat it as data. This reduces prompt-injection risk but does not prevent it, which is why rules need approval (or the committee in auto mode). `checkAiBudget` (`src/lib/bedrock-meter.ts`) runs in Postgres before every Bedrock call (suggest-transforms on every pass, the auto-validate committee, chat-builder, generate-data), and every call is metered in `bedrock_usage`. - **Guest access** (`src/lib/guest.ts`, off unless `GUEST_COOKIE_SECRET` is set). `POST /api/guest` issues an httpOnly, `SameSite=Lax`, `Secure` cookie `cs_guest` = `guest_<22 chars>..`, valid for 24 h; `auth()` returns that `guest_…` id as the user id (it is the `team_id`), and only when no Clerk user is signed in. Sessions are recorded in `guest_sessions` with an HMAC of the client IP (never the raw IP); at most 5 per IP and 200 overall per 24 h. Turnstile is checked when `TURNSTILE_SECRET_KEY` is set. Guest quotas and blocked features are listed in `src/lib/guest-limits.ts`. +- **Guest limits** (all counted in Postgres, so they hold when Upstash is down): + + | Limit | Value | Where | + |---|---|---| + | File size | 2 MB | presigned POST policy, `/api/upload` 413, profiler | + | Rows per run | 5,000 | `suggest-transforms`, before templates or AI; the run fails | + | Rows per guest | 10,000 (all passes; runs refused for size do not count) | `checkQuota` "guest" plan | + | Uploads | 3 per guest, 10 per IP hash per 24 h | the capped `INSERT` in `src/lib/guest-quota.ts` (also refuses expired sessions) | + | Pipelines | 5 per guest | `POST /api/pipelines` | + | AI calls | 10 per rolling hour | `checkAiBudget` | + | AI spend | $0.25 per guest; $5 per UTC day for all guests together | `checkAiBudget` | + | Sessions | 5 per IP hash and 200 overall per 24 h | `POST /api/guest` | + + Guests cannot use the chat builder or synthetic data, auto-clean, training export, Slack alerts, account deletion, templates or `/api/admin`: the middleware answers 403 (or redirects `/templates` to the dashboard) and each route also refuses guest ids (`forbidGuest`). When auto-validate is over budget, the run falls back to manual review. - **Abuse limits.** Upstash sliding-window limits per user: 20 uploads/h, 50 AI calls/h, 30 chat messages/h. They fail open when Redis is not configured. - **Headers.** A CSP is set in `next.config.ts`. It allows `'unsafe-inline'` and `'unsafe-eval'` for scripts (needed by Next.js and Clerk without nonces). - **Secrets.** Nothing secret is committed. `.env.example` has placeholders only, and `run-migration.mjs` reads the database host from the environment. AWS account ids and ARNs from older commits remain in git history. @@ -271,7 +286,7 @@ Report vulnerabilities as described in [SECURITY.md](SECURITY.md). - **No retention enforcement.** Raw files of failed runs stay in S3; `data_retention_days` is not enforced. - **Downloads are served as stored.** Deliverables produced before the sidecar split may still contain `__orig_*` columns. **Export As** and the training export strip them; the native download does not. Re-run those pipelines to regenerate them. - **Excel blanks are not counted as nulls** by the profiler (cells are read as empty strings). Before/after scores are still computed the same way. -- **AI spend cap coverage.** The cap is checked before `suggest-transforms` only. Committee calls are metered but not blocked, and chat-builder calls are neither metered nor blocked. +- **AI budget is checked before a call and recorded after it.** Concurrent requests can each pass the check and overshoot a cap by the cost of the calls in flight (the committee's 3 calls are reserved together). Costs are estimates from `src/lib/ai-config.ts`, not the AWS bill. - **Upload size limit is per file only.** 100 MB (`MAX_UPLOAD_MB`) for users, 2 MB for guests. A file under the limit can still exhaust the executor Lambda's memory or time for expensive rules. `semantic_deduplicate` (MinHash + LSH, a few seconds for 50k short texts) is skipped with a reason instead of running past the Lambda deadline, and above `SEMANTIC_DEDUP_MAX_ROWS` rows; other rules have no such guard. - **Schema drift alerts** store their diff inside `column_definitions`, so the next alert reports a phantom `_diff` column. - **Templates** copy approved rules from all recent runs of a pipeline, not only the latest run. diff --git a/src/app/api/account/route.ts b/src/app/api/account/route.ts index 6d6bd61..d19c5cb 100644 --- a/src/app/api/account/route.ts +++ b/src/app/api/account/route.ts @@ -1,4 +1,5 @@ import { auth } from "@/lib/auth"; +import { forbidGuest } from "@/lib/guest-guard"; import { NextRequest, NextResponse } from "next/server"; import { S3Client } from "@aws-sdk/client-s3"; import { query, queryOne, withTransaction } from "@/lib/db"; @@ -29,6 +30,8 @@ const SAFE_ID = /^[A-Za-z0-9_-]+$/; export async function DELETE(req: NextRequest) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; // Require explicit confirmation to prevent accidental or CSRF-triggered deletion const confirm = req.nextUrl.searchParams.get("confirm"); diff --git a/src/app/api/alerts/configure/route.ts b/src/app/api/alerts/configure/route.ts index 8ad59fa..36a8b4b 100644 --- a/src/app/api/alerts/configure/route.ts +++ b/src/app/api/alerts/configure/route.ts @@ -1,10 +1,13 @@ import { NextRequest, NextResponse } from "next/server"; import { auth } from "@/lib/auth"; +import { forbidGuest } from "@/lib/guest-guard"; import { queryOne, queryOneWithTeam } from "@/lib/db"; export async function POST(req: NextRequest) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; const { pipeline_id, slack_webhook_url } = await req.json(); if (!pipeline_id) { @@ -50,6 +53,8 @@ export async function POST(req: NextRequest) { export async function GET(req: NextRequest) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; const { searchParams } = new URL(req.url); const pipeline_id = searchParams.get("pipeline_id"); diff --git a/src/app/api/auto-validate/[runId]/route.ts b/src/app/api/auto-validate/[runId]/route.ts index 8e1675e..c23d0eb 100644 --- a/src/app/api/auto-validate/[runId]/route.ts +++ b/src/app/api/auto-validate/[runId]/route.ts @@ -6,7 +6,7 @@ import { languageModel } from "@/lib/ai-model"; import { BEDROCK_MODEL_ID } from "@/lib/ai-config"; import { SQSClient, SendMessageCommand } from "@aws-sdk/client-sqs"; import { query, queryOne } from "@/lib/db"; -import { meterBedrockCall } from "@/lib/bedrock-meter"; +import { checkAiBudget, meterBedrockCall } from "@/lib/bedrock-meter"; import type { TransformRule, DataProfile } from "@/lib/types"; import { requireEnv, optionalEnv, awsRegion } from "@/lib/env"; import { logger } from "@/lib/logger"; @@ -169,6 +169,20 @@ Include a vote for every rule_id listed. No extra text.`; [runId] ); + // AI budget for the 3 committee calls (src/lib/bedrock-meter.ts). Over budget, + // the run falls back to manual review instead of failing: the rules exist. + if (pipelineRow?.team_id) { + const budget = await checkAiBudget(pipelineRow.team_id, 3); + if (!budget.ok) { + log.warn("AI budget reached; auto-validate falls back to manual review", { run_id: runId, scope: budget.scope }); + await queryOne( + "UPDATE pipeline_runs SET status = 'awaiting_approval', error_message = $2, updated_at = now() WHERE id = $1", + [runId, `Auto-review skipped: ${budget.error} Review the rules manually.`] + ); + return NextResponse.json({ ok: true, fallback: "manual_review", scope: budget.scope }); + } + } + const [auditorResult, statResult, domainResult] = await Promise.all([ runConsultant( "SafetyAuditor", @@ -202,13 +216,13 @@ ${responseFormat}`, // Meter all 3 Bedrock calls if (pipelineRow?.team_id) { - [ + await Promise.all([ { r: auditorResult, type: "auto_validate_auditor" }, { r: statResult, type: "auto_validate_stat" }, { r: domainResult, type: "auto_validate_domain" }, - ].forEach(({ r, type }) => + ].map(({ r, type }) => meterBedrockCall({ teamId: pipelineRow.team_id, runId, callType: type, model: BEDROCK_MODEL_ID, usage: r.usage }) - ); + )); } const [auditorVotes, statVotes, domainVotes] = [auditorResult.votes, statResult.votes, domainResult.votes]; diff --git a/src/app/api/chat-builder/generate-data/route.ts b/src/app/api/chat-builder/generate-data/route.ts index 4fca658..ada111e 100644 --- a/src/app/api/chat-builder/generate-data/route.ts +++ b/src/app/api/chat-builder/generate-data/route.ts @@ -4,6 +4,9 @@ import { generateText } from "ai"; import { languageModel } from "@/lib/ai-model"; import { aiLimiter, checkRateLimit } from "@/lib/rate-limit"; import { logger } from "@/lib/logger"; +import { forbidGuest } from "@/lib/guest-guard"; +import { checkAiBudget, meterBedrockCall } from "@/lib/bedrock-meter"; +import { BEDROCK_MODEL_ID } from "@/lib/ai-config"; const log = logger.child({ route: "POST /api/chat-builder/generate-data" }); @@ -59,9 +62,15 @@ export async function POST(req: NextRequest) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; + const rateLimitRes = await checkRateLimit(aiLimiter, userId); if (rateLimitRes) return rateLimitRes; + const budget = await checkAiBudget(userId); + if (!budget.ok) return NextResponse.json({ error: budget.error }, { status: budget.status }); + const body: GenerateRequest = await req.json(); const { description: rawDescription, config, format = "csv", row_count = 20 } = body; const description = typeof rawDescription === "string" ? rawDescription.slice(0, 2000) : ""; @@ -114,11 +123,12 @@ Requirements: let rows: Record[]; try { - const { text } = await generateText({ + const { text, usage } = await generateText({ model: languageModel(), prompt, maxOutputTokens: 4000, }); + await meterBedrockCall({ teamId: userId, runId: null, callType: "generate_data", model: BEDROCK_MODEL_ID, usage }); // Strip any accidental markdown fences const cleaned = text.replace(/```json\s*/g, "").replace(/```\s*/g, "").trim(); diff --git a/src/app/api/chat-builder/route.test.ts b/src/app/api/chat-builder/route.test.ts new file mode 100644 index 0000000..4ad8912 --- /dev/null +++ b/src/app/api/chat-builder/route.test.ts @@ -0,0 +1,54 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const h = vi.hoisted(() => ({ budget: vi.fn(), meter: vi.fn(), streamText: vi.fn(), generateText: vi.fn() })); +vi.mock("@/lib/auth", () => ({ auth: async () => ({ userId: "user_1" }) })); +vi.mock("@/lib/bedrock-meter", () => ({ checkAiBudget: (...a: unknown[]) => h.budget(...a), meterBedrockCall: (...a: unknown[]) => h.meter(...a) })); +vi.mock("@/lib/ai-model", () => ({ languageModel: () => "model" })); +vi.mock("ai", async (orig) => ({ + ...(await orig()), + streamText: (...a: unknown[]) => h.streamText(...a), + generateText: (...a: unknown[]) => h.generateText(...a), +})); + +import { POST as chat } from "./route"; +import { POST as generate } from "./generate-data/route"; + +const chatReq = () => new Request("http://x/api/chat-builder", { method: "POST", body: JSON.stringify({ messages: [] }) }); +const genReq = () => new Request("http://x/api/chat-builder/generate-data", { + method: "POST", body: JSON.stringify({ description: "orders", config: { name: "n", description: "d", rules: [] } }), +}); +const OVER = { ok: false, status: 402, scope: "team_month", error: "AI spend cap reached ($200.00 / $200 this month). Contact support." }; + +beforeEach(() => { + vi.clearAllMocks(); + h.budget.mockResolvedValue({ ok: true }); + h.streamText.mockReturnValue({ toUIMessageStreamResponse: () => new Response("stream") }); + h.generateText.mockResolvedValue({ text: '[{"a":1}]', usage: { inputTokens: 10, outputTokens: 5 } }); +}); + +describe("chat-builder AI metering for signed-in users", () => { + it("checks the spend cap before streaming and meters the call on finish", async () => { + expect((await chat(chatReq())).status).toBe(200); + expect(h.budget).toHaveBeenCalledWith("user_1"); + const { onFinish } = h.streamText.mock.calls[0][0]; + await onFinish({ totalUsage: { inputTokens: 100, outputTokens: 20 } }); + expect(h.meter).toHaveBeenCalledWith(expect.objectContaining({ teamId: "user_1", runId: null, callType: "chat_builder", usage: { inputTokens: 100, outputTokens: 20 } })); + }); + + it("refuses with the cap message and never calls the model", async () => { + h.budget.mockResolvedValue(OVER); + const res = await chat(chatReq()); + expect(res.status).toBe(402); + expect((await res.json()).error).toBe(OVER.error); + expect(h.streamText).not.toHaveBeenCalled(); + }); + + it("generate-data is capped and metered the same way", async () => { + expect((await generate(genReq() as never)).status).toBe(200); + expect(h.meter).toHaveBeenCalledWith(expect.objectContaining({ callType: "generate_data", usage: { inputTokens: 10, outputTokens: 5 } })); + vi.clearAllMocks(); + h.budget.mockResolvedValue(OVER); + expect((await generate(genReq() as never)).status).toBe(402); + expect(h.generateText).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/chat-builder/route.ts b/src/app/api/chat-builder/route.ts index 311cc2b..b21bcd9 100644 --- a/src/app/api/chat-builder/route.ts +++ b/src/app/api/chat-builder/route.ts @@ -3,6 +3,9 @@ import { languageModel } from "@/lib/ai-model"; import { auth } from "@/lib/auth"; import { NextResponse } from "next/server"; import { chatLimiter, checkRateLimit } from "@/lib/rate-limit"; +import { forbidGuest } from "@/lib/guest-guard"; +import { checkAiBudget, meterBedrockCall } from "@/lib/bedrock-meter"; +import { BEDROCK_MODEL_ID } from "@/lib/ai-config"; export const maxDuration = 60; @@ -48,15 +51,24 @@ export async function POST(req: Request) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; + const rateLimitRes = await checkRateLimit(chatLimiter, userId); if (rateLimitRes) return rateLimitRes; + // Monthly AI spend cap, and meter the call when the stream finishes. + const budget = await checkAiBudget(userId); + if (!budget.ok) return NextResponse.json({ error: budget.error }, { status: budget.status }); + const { messages }: { messages: UIMessage[] } = await req.json(); const result = streamText({ model: languageModel(), system: SYSTEM, messages: await convertToModelMessages(messages), + onFinish: ({ totalUsage }) => + meterBedrockCall({ teamId: userId, runId: null, callType: "chat_builder", model: BEDROCK_MODEL_ID, usage: totalUsage }), }); return result.toUIMessageStreamResponse(); diff --git a/src/app/api/export-training/[runId]/route.ts b/src/app/api/export-training/[runId]/route.ts index 519f825..22bc8bb 100644 --- a/src/app/api/export-training/[runId]/route.ts +++ b/src/app/api/export-training/[runId]/route.ts @@ -1,4 +1,5 @@ import { auth } from "@/lib/auth"; +import { forbidGuest } from "@/lib/guest-guard"; import { NextRequest, NextResponse } from "next/server"; import { S3Client, GetObjectCommand } from "@aws-sdk/client-s3"; import { queryOneWithTeam } from "@/lib/db"; @@ -41,6 +42,8 @@ export async function GET( ) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; const rateLimitRes = await checkRateLimit(aiLimiter, userId); if (rateLimitRes) return rateLimitRes; diff --git a/src/app/api/guest-blocked.test.ts b/src/app/api/guest-blocked.test.ts new file mode 100644 index 0000000..ad185bc --- /dev/null +++ b/src/app/api/guest-blocked.test.ts @@ -0,0 +1,44 @@ +/** + * Route-level guest guards (defense in depth behind the middleware): every + * blocked route answers 403 for a guest id before touching the database, S3 or Bedrock. + */ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { NextRequest } from "next/server"; + +const h = vi.hoisted(() => ({ userId: "guest_AAAAAAAAAAAAAAAAAAAAAA" as string | null, db: vi.fn() })); +vi.mock("@/lib/auth", () => ({ auth: async () => ({ userId: h.userId }), currentUserEmail: async () => null, userEmailById: async () => null })); +vi.mock("@/lib/db", () => { + const fail = (...a: unknown[]) => { h.db(...a); throw new Error("db must not be reached"); }; + return { query: fail, queryOne: fail, queryWithTeam: fail, queryOneWithTeam: fail, withTransaction: fail }; +}); +vi.mock("@/lib/ai-model", () => ({ languageModel: () => { throw new Error("model must not be reached"); } })); + +const ID = "11111111-2222-3333-4444-555555555555"; +const runCtx = { params: Promise.resolve({ runId: ID }) }; +const idCtx = { params: Promise.resolve({ id: ID }) }; +const req = (method = "POST", body: unknown = {}) => + new NextRequest("http://localhost/x", { method, headers: { "content-type": "application/json" }, body: method === "GET" ? undefined : JSON.stringify(body) }); + +const cases: Array<[string, () => Promise]> = [ + ["POST /api/chat-builder", async () => (await import("./chat-builder/route")).POST(req("POST", { messages: [] }))], + ["POST /api/chat-builder/generate-data", async () => (await import("./chat-builder/generate-data/route")).POST(req())], + ["POST /api/runs/:id/auto-clean", async () => (await import("./runs/[runId]/auto-clean/route")).POST(req(), runCtx)], + ["GET /api/export-training/:id", async () => (await import("./export-training/[runId]/route")).GET(req("GET"), runCtx)], + ["POST /api/alerts/configure", async () => (await import("./alerts/configure/route")).POST(req())], + ["GET /api/alerts/configure", async () => (await import("./alerts/configure/route")).GET(req("GET"))], + ["DELETE /api/account", async () => (await import("./account/route")).DELETE(req("DELETE"))], + ["GET /api/templates", async () => (await import("./templates/route")).GET(req("GET"))], + ["POST /api/templates", async () => (await import("./templates/route")).POST(req())], + ["POST /api/templates/:id/use", async () => (await import("./templates/[id]/use/route")).POST(req(), idCtx)], +]; + +beforeEach(() => { h.db.mockClear(); }); + +describe("guest-blocked routes", () => { + it.each(cases)("%s answers 403 for a guest", async (_name, call) => { + const res = await call(); + expect(res.status).toBe(403); + expect(await res.json()).toMatchObject({ guest: true }); + expect(h.db).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/pipelines/route.test.ts b/src/app/api/pipelines/route.test.ts new file mode 100644 index 0000000..a750481 --- /dev/null +++ b/src/app/api/pipelines/route.test.ts @@ -0,0 +1,33 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { NextRequest } from "next/server"; + +const h = vi.hoisted(() => ({ userId: "user_1" as string | null, queryOne: vi.fn() })); +vi.mock("@/lib/auth", () => ({ auth: async () => ({ userId: h.userId }) })); +vi.mock("@/lib/db", () => ({ query: vi.fn(), queryWithTeam: vi.fn(), queryOne: (...a: unknown[]) => h.queryOne(...a) })); + +import { POST } from "./route"; + +const GUEST = "guest_AAAAAAAAAAAAAAAAAAAAAA"; +const create = () => POST(new NextRequest("http://x/api/pipelines", { method: "POST", body: JSON.stringify({ name: " p " }) })); +beforeEach(() => { vi.clearAllMocks(); h.userId = "user_1"; }); + +describe("POST /api/pipelines", () => { + it("users: plain insert", async () => { + h.queryOne.mockResolvedValue({ id: "p1" }); + expect((await create()).status).toBe(201); + expect(h.queryOne.mock.calls[0][0]).not.toContain("count(*)"); + }); + + it("guests: insert capped at 5 pipelines, 429 past it", async () => { + h.userId = GUEST; + h.queryOne.mockResolvedValue({ id: "p1" }); + expect((await create()).status).toBe(201); + const [sql, params] = h.queryOne.mock.calls[0]; + expect(sql).toContain("(SELECT count(*) FROM pipelines WHERE team_id = $3) < $4"); + expect(params).toEqual(["p", null, GUEST, 5]); + h.queryOne.mockResolvedValue(null); + const res = await create(); + expect(res.status).toBe(429); + expect((await res.json()).error).toBe("Guests can create 5 pipelines. Sign up to keep going."); + }); +}); diff --git a/src/app/api/pipelines/route.ts b/src/app/api/pipelines/route.ts index 624fb57..fb003b5 100644 --- a/src/app/api/pipelines/route.ts +++ b/src/app/api/pipelines/route.ts @@ -3,6 +3,9 @@ import { NextRequest, NextResponse } from "next/server"; import { query, queryOne, queryWithTeam } from "@/lib/db"; import type { Pipeline } from "@/lib/types"; import { logger } from "@/lib/logger"; +import { isGuestId } from "@/lib/guest"; +import { GUEST_LIMITS } from "@/lib/guest-limits"; +import { GUEST_PIPELINE_INSERT_SQL } from "@/lib/guest-quota"; const log = logger.child({ route: "/api/pipelines" }); @@ -35,6 +38,17 @@ export async function POST(req: NextRequest) { } try { + if (isGuestId(userId)) { + const pipeline = await queryOne(GUEST_PIPELINE_INSERT_SQL, + [name.trim(), description?.trim() || null, userId, GUEST_LIMITS.pipelinesPerGuest]); + if (!pipeline) { + return NextResponse.json( + { error: `Guests can create ${GUEST_LIMITS.pipelinesPerGuest} pipelines. Sign up to keep going.`, guest: true }, + { status: 429 } + ); + } + return NextResponse.json({ pipeline }, { status: 201 }); + } const pipeline = await queryOne( `INSERT INTO pipelines (name, description, owner_id, team_id) VALUES ($1, $2, $3, $4) diff --git a/src/app/api/runs/[runId]/auto-clean/route.ts b/src/app/api/runs/[runId]/auto-clean/route.ts index f895fb3..c7497fe 100644 --- a/src/app/api/runs/[runId]/auto-clean/route.ts +++ b/src/app/api/runs/[runId]/auto-clean/route.ts @@ -1,4 +1,5 @@ import { auth } from "@/lib/auth"; +import { forbidGuest } from "@/lib/guest-guard"; import { NextRequest, NextResponse } from "next/server"; import { S3Client, CopyObjectCommand } from "@aws-sdk/client-s3"; import { randomUUID } from "crypto"; @@ -17,6 +18,8 @@ export async function POST( ) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; const { runId } = await params; diff --git a/src/app/api/suggest-transforms/route.ts b/src/app/api/suggest-transforms/route.ts index ffeb2b5..c8da9f4 100644 --- a/src/app/api/suggest-transforms/route.ts +++ b/src/app/api/suggest-transforms/route.ts @@ -4,9 +4,11 @@ import { generateText, Output } from "ai"; import { languageModel } from "@/lib/ai-model"; import { BEDROCK_MODEL_ID } from "@/lib/ai-config"; -import { checkQuota } from "@/lib/billing"; +import { checkQuota, quotaBlockedMessage } from "@/lib/billing"; import { userEmailById } from "@/lib/auth"; -import { meterBedrockCall, checkAiSpendCap } from "@/lib/bedrock-meter"; +import { meterBedrockCall, checkAiBudget } from "@/lib/bedrock-meter"; +import { isGuestId } from "@/lib/guest"; +import { GUEST_LIMITS } from "@/lib/guest-limits"; import { aiLimiter, checkRateLimit } from "@/lib/rate-limit"; export const maxDuration = 300; @@ -112,8 +114,20 @@ export async function POST(req: NextRequest) { return rateLimitRes; } - // AI-spend guard for pass 1 only — prevent Bedrock cost if Free-tier quota exceeded - if ((run.iteration ?? 1) === 1) { + const guest = isGuestId(run.team_id); + const failRun = (message: string) => + queryOne("UPDATE pipeline_runs SET status = 'failed', error_message = $2, updated_at = now() WHERE id = $1", [run_id, message]); + + // Guests: one file may have at most GUEST_LIMITS.rowsPerRun rows (checked + // before templates and AI, so an oversized file costs nothing). + if (guest && Number(run.row_count_raw ?? 0) > GUEST_LIMITS.rowsPerRun) { + await failRun(`Guest runs are limited to ${GUEST_LIMITS.rowsPerRun.toLocaleString()} rows; this file has ${Number(run.row_count_raw).toLocaleString()}. Sign up to clean larger files.`); + return NextResponse.json({ error: "Guest row limit exceeded" }, { status: 402 }); + } + + // Row quota: pass 1 for users (later passes re-read the same rows); every pass + // for guests, whose allowance is all-time. + if ((run.iteration ?? 1) === 1 || guest) { const pipelineRow = await queryOne<{ team_id: string }>( "SELECT team_id FROM pipelines WHERE id = $1", [run.pipeline_id] @@ -124,17 +138,10 @@ export async function POST(req: NextRequest) { if (quota.blocked) { await queryOne( "UPDATE pipeline_runs SET status = 'failed', error_message = $2, updated_at = now() WHERE id = $1", - [run_id, `Monthly row limit reached (${quota.used.toLocaleString()} / ${quota.includedRows.toLocaleString()} rows on ${quota.plan} plan). Upgrade to continue.`] + [run_id, quotaBlockedMessage(quota)] ); return NextResponse.json({ error: "Quota exceeded" }, { status: 402 }); } - // AI spend cap check - const aiSpend = await checkAiSpendCap(pipelineRow.team_id); - if (aiSpend.blocked) { - await queryOne("UPDATE pipeline_runs SET status = 'failed', error_message = $2, updated_at = now() WHERE id = $1", - [run_id, `AI spend cap reached ($${aiSpend.currentSpendUsd.toFixed(2)} / $${aiSpend.hardCapUsd} this month). Contact support.`]); - return NextResponse.json({ error: "AI spend cap exceeded" }, { status: 402 }); - } } } @@ -185,6 +192,15 @@ export async function POST(req: NextRequest) { return NextResponse.json({ error: "Profile not found" }, { status: 404 }); } + // AI budget before the Bedrock call, every pass (src/lib/bedrock-meter.ts): + // the monthly hard cap for users; hourly calls, per-guest and all-guests + // daily spend for guests. Template runs above never reach Bedrock. + const budget = await checkAiBudget(run.team_id); + if (!budget.ok) { + await failRun(budget.error); + return NextResponse.json({ error: budget.error, scope: budget.scope }, { status: budget.status }); + } + // Document mode — separate prompt + schema if (run.mode === "document") { const dp = profile.column_stats as unknown as { @@ -311,7 +327,7 @@ IMPORTANT RULES: prompt: docPrompt, }); docOutput = result.output; - meterBedrockCall({ teamId: run.team_id, runId: run_id, callType: "suggest_transforms_doc", model: BEDROCK_MODEL_ID, usage: result.usage }); + await meterBedrockCall({ teamId: run.team_id, runId: run_id, callType: "suggest_transforms_doc", model: BEDROCK_MODEL_ID, usage: result.usage }); } catch (aiErr) { log.error("Bedrock document suggestion failed", { run_id, err: aiErr }); await queryOne("UPDATE pipeline_runs SET status = 'failed', error_message = $2, updated_at = now() WHERE id = $1", @@ -654,7 +670,7 @@ For each rule, write ai_reasoning as one precise sentence that references the sp prompt, }); output = result.output; - meterBedrockCall({ teamId: run.team_id, runId: run_id, callType: "suggest_transforms", model: BEDROCK_MODEL_ID, usage: result.usage }); + await meterBedrockCall({ teamId: run.team_id, runId: run_id, callType: "suggest_transforms", model: BEDROCK_MODEL_ID, usage: result.usage }); } catch (aiErr) { log.error("Bedrock tabular suggestion failed", { run_id, err: aiErr }); await queryOne( diff --git a/src/app/api/templates/[id]/use/route.ts b/src/app/api/templates/[id]/use/route.ts index 989a630..ed35fd1 100644 --- a/src/app/api/templates/[id]/use/route.ts +++ b/src/app/api/templates/[id]/use/route.ts @@ -1,5 +1,6 @@ import { NextRequest, NextResponse } from "next/server"; import { auth } from "@/lib/auth"; +import { forbidGuest } from "@/lib/guest-guard"; import { queryOne } from "@/lib/db"; import type { PipelineTemplate } from "@/lib/types"; @@ -9,6 +10,8 @@ export async function POST( ) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; const { id } = await params; diff --git a/src/app/api/templates/route.ts b/src/app/api/templates/route.ts index 7da7879..9f766dc 100644 --- a/src/app/api/templates/route.ts +++ b/src/app/api/templates/route.ts @@ -1,9 +1,13 @@ import { NextRequest, NextResponse } from "next/server"; import { auth } from "@/lib/auth"; +import { forbidGuest } from "@/lib/guest-guard"; import { query, queryOne } from "@/lib/db"; import type { PipelineTemplate, TemplateRule } from "@/lib/types"; export async function GET(req: NextRequest) { + const guestRes = forbidGuest((await auth()).userId); + if (guestRes) return guestRes; + const { searchParams } = new URL(req.url); const category = searchParams.get("category"); @@ -21,6 +25,8 @@ export async function GET(req: NextRequest) { export async function POST(req: NextRequest) { const { userId } = await auth(); if (!userId) return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); + const guestRes = forbidGuest(userId); + if (guestRes) return guestRes; const { pipeline_id, name, description, category } = await req.json(); if (!pipeline_id || !name) { diff --git a/src/app/api/upload/route.test.ts b/src/app/api/upload/route.test.ts index 8e7c019..8d62ea3 100644 --- a/src/app/api/upload/route.test.ts +++ b/src/app/api/upload/route.test.ts @@ -8,6 +8,12 @@ const h = vi.hoisted(() => ({ queryOneWithTeam: vi.fn(), queryOne: vi.fn(), createPresignedPost: vi.fn(), + insertGuestRun: vi.fn(), + guestUploadRefusal: vi.fn(), +})); +vi.mock("@/lib/guest-quota", () => ({ + insertGuestRun: (...a: unknown[]) => h.insertGuestRun(...a), + guestUploadRefusal: (...a: unknown[]) => h.guestUploadRefusal(...a), })); vi.mock("@/lib/auth", () => ({ auth: () => h.auth(), currentUserEmail: () => h.email() })); @@ -70,13 +76,26 @@ describe("POST /api/upload", () => { }); }); - it("guests get a 2 MB policy", async () => { + it("guests get a 2 MB policy, and their run is created by the capped insert", async () => { h.auth.mockResolvedValue({ userId: "guest_" + "a".repeat(22) }); + h.insertGuestRun.mockResolvedValue({ id: "run_g" }); const res = await call(); expect(res.status).toBe(200); + expect(h.queryOne).not.toHaveBeenCalled(); + expect(h.insertGuestRun.mock.calls[0][0]).toMatchObject({ guestId: "guest_" + "a".repeat(22), fileFormat: "csv" }); expect(h.createPresignedPost.mock.calls[0][1].Conditions[0]).toEqual(["content-length-range", 1, 2 * 1024 * 1024]); }); + it("answers 429 with the reason when a guest is over the upload caps, without presigning", async () => { + h.auth.mockResolvedValue({ userId: "guest_" + "a".repeat(22) }); + h.insertGuestRun.mockResolvedValue(null); + h.guestUploadRefusal.mockResolvedValue("Guests can upload 3 files. Sign up to keep going."); + const res = await call(); + expect(res.status).toBe(429); + expect(await res.json()).toEqual({ error: "Guests can upload 3 files. Sign up to keep going.", guest: true }); + expect(h.createPresignedPost).not.toHaveBeenCalled(); + }); + it("refuses a declared size over the limit with 413 before creating a run", async () => { h.auth.mockResolvedValue({ userId: "guest_" + "a".repeat(22) }); const res = await call({ pipeline_id: "p", filename: "big.csv", size: 2 * 1024 * 1024 + 1 }); diff --git a/src/app/api/upload/route.ts b/src/app/api/upload/route.ts index a5ff8e4..91a559e 100644 --- a/src/app/api/upload/route.ts +++ b/src/app/api/upload/route.ts @@ -5,11 +5,14 @@ import { S3Client } from "@aws-sdk/client-s3"; import { createPresignedPost } from "@aws-sdk/s3-presigned-post"; import { queryOne, queryOneWithTeam } from "@/lib/db"; import { getCachedQuota } from "@/lib/quota-cache"; +import { quotaBlockedMessage } from "@/lib/billing"; import { uploadLimiter, checkRateLimit } from "@/lib/rate-limit"; import type { PipelineRun } from "@/lib/types"; import { requireEnv, awsRegion } from "@/lib/env"; import { logger } from "@/lib/logger"; import { formatMb, maxUploadBytesFor } from "@/lib/upload-limits"; +import { isGuestId } from "@/lib/guest"; +import { guestUploadRefusal, insertGuestRun } from "@/lib/guest-quota"; const log = logger.child({ route: "POST /api/upload" }); @@ -102,7 +105,7 @@ export async function POST(req: NextRequest) { if (quota.blocked) { return NextResponse.json( { - error: `Monthly row limit reached (${quota.used.toLocaleString()} / ${quota.includedRows.toLocaleString()} rows on ${quota.plan} plan). Upgrade at /pricing to continue.`, + error: quotaBlockedMessage(quota), }, { status: 402 } ); @@ -146,12 +149,19 @@ export async function POST(req: NextRequest) { const runId = randomUUID(); const s3Key = `${userId}/${pipeline_id}/${runId}/raw.${ext}`; - const run = await queryOne( - `INSERT INTO pipeline_runs (id, pipeline_id, status, file_format, raw_s3_key, started_at) - VALUES ($1, $2, 'pending', $3, $4, now()) - RETURNING *`, - [runId, pipeline_id, ext, s3Key] - ); + // Guests: the insert itself enforces the per-guest and per-IP upload caps. + const guest = isGuestId(userId); + const run = guest + ? await insertGuestRun({ runId, pipelineId: pipeline_id, fileFormat: ext, s3Key, guestId: userId }) + : await queryOne( + `INSERT INTO pipeline_runs (id, pipeline_id, status, file_format, raw_s3_key, started_at) + VALUES ($1, $2, 'pending', $3, $4, now()) + RETURNING *`, + [runId, pipeline_id, ext, s3Key] + ); + if (!run && guest) { + return NextResponse.json({ error: await guestUploadRefusal(userId), guest: true }, { status: 429 }); + } if (!run) return NextResponse.json({ error: "Failed to create run" }, { status: 500 }); diff --git a/src/lib/bedrock-meter.test.ts b/src/lib/bedrock-meter.test.ts new file mode 100644 index 0000000..b88fd93 --- /dev/null +++ b/src/lib/bedrock-meter.test.ts @@ -0,0 +1,75 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const h = vi.hoisted(() => ({ query: vi.fn(), queryOne: vi.fn() })); +vi.mock("@/lib/db", () => ({ query: (...a: unknown[]) => h.query(...a), queryOne: (...a: unknown[]) => h.queryOne(...a) })); + +import { checkAiBudget, meterBedrockCall } from "./bedrock-meter"; +import { BEDROCK_MODEL_ID } from "./ai-config"; + +const GUEST = "guest_AAAAAAAAAAAAAAAAAAAAAA"; +const guestRow = (calls_hour: number, guest_total: number, guests_today: number) => + h.queryOne.mockResolvedValue({ calls_hour: String(calls_hour), guest_total: String(guest_total), guests_today: String(guests_today) }); + +beforeEach(() => { vi.clearAllMocks(); h.query.mockResolvedValue([]); }); + +describe("checkAiBudget: signed-in users", () => { + it("uses the monthly hard cap", async () => { + h.queryOne.mockImplementation(async (sql: string) => + sql.includes("ai_spend_limits") ? { soft_cap_usd: "50", hard_cap_usd: "200" } : { total: "199.99" }); + expect(await checkAiBudget("user_1")).toEqual({ ok: true }); + h.queryOne.mockImplementation(async (sql: string) => + sql.includes("ai_spend_limits") ? { soft_cap_usd: "50", hard_cap_usd: "200" } : { total: "200" }); + expect(await checkAiBudget("user_1")).toMatchObject({ ok: false, status: 402, scope: "team_month" }); + }); +}); + +describe("checkAiBudget: guests", () => { + it("allows a fresh guest and counts in one query scoped to the guest and to all guests today", async () => { + guestRow(0, 0, 0); + expect(await checkAiBudget(GUEST)).toEqual({ ok: true }); + const [sql, params] = h.queryOne.mock.calls[0]; + expect(params).toEqual([GUEST]); + expect(sql).toContain("interval '1 hour'"); + expect(sql).toContain("team_id LIKE 'guest\\_%'"); + expect(sql).toContain("date_trunc('day', now() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'"); + }); + + it("refuses the 11th call in an hour, and a 3-call committee at 8", async () => { + guestRow(9, 0.01, 0); + expect(await checkAiBudget(GUEST)).toEqual({ ok: true }); + guestRow(10, 0.01, 0); + expect(await checkAiBudget(GUEST)).toMatchObject({ ok: false, status: 429, scope: "guest_hourly" }); + guestRow(8, 0.01, 0); + expect(await checkAiBudget(GUEST, 3)).toMatchObject({ ok: false, scope: "guest_hourly" }); + guestRow(7, 0.01, 0); + expect(await checkAiBudget(GUEST, 3)).toEqual({ ok: true }); + }); + + it("refuses at $0.25 for the guest and at $5 for all guests today", async () => { + guestRow(0, 0.2499, 4.99); + expect(await checkAiBudget(GUEST)).toEqual({ ok: true }); + guestRow(0, 0.25, 0.25); + expect(await checkAiBudget(GUEST)).toMatchObject({ ok: false, status: 402, scope: "guest_total" }); + guestRow(0, 0, 5); + expect(await checkAiBudget(GUEST)).toMatchObject({ ok: false, status: 429, scope: "guests_daily" }); + }); +}); + +describe("meterBedrockCall", () => { + it("records the call and resolves only after the insert (so the next check sees it)", async () => { + let resolveInsert!: () => void; + h.query.mockReturnValue(new Promise((r) => { resolveInsert = r; })); + let done = false; + const p = meterBedrockCall({ teamId: GUEST, runId: null, callType: "chat_builder", model: BEDROCK_MODEL_ID, usage: { inputTokens: 1000, outputTokens: 100 } }).then(() => { done = true; }); + await Promise.resolve(); + expect(done).toBe(false); + resolveInsert(); + await p; + expect(h.query.mock.calls[0][1].slice(0, 6)).toEqual([GUEST, null, BEDROCK_MODEL_ID, "chat_builder", 1000, 100]); + }); + + it("never throws when the insert fails", async () => { + h.query.mockRejectedValue(new Error("db down")); + await expect(meterBedrockCall({ teamId: GUEST, runId: null, callType: "x", model: "m", usage: {} })).resolves.toBeUndefined(); + }); +}); diff --git a/src/lib/bedrock-meter.ts b/src/lib/bedrock-meter.ts index f3ca652..39f7bb1 100644 --- a/src/lib/bedrock-meter.ts +++ b/src/lib/bedrock-meter.ts @@ -1,6 +1,8 @@ import { query, queryOne } from "@/lib/db"; import { logger } from "@/lib/logger"; import { estimateBedrockCostUsd } from "@/lib/ai-config"; +import { isGuestId } from "@/lib/guest"; +import { GUEST_LIMITS } from "@/lib/guest-limits"; const log = logger.child({ component: "bedrock-meter" }); @@ -16,8 +18,9 @@ export async function meterBedrockCall(opts: { const outputTokens = usage.completionTokens ?? usage.outputTokens ?? 0; const cost = estimateBedrockCostUsd(model, inputTokens, outputTokens); - // Fire-and-forget — never block the main flow - query( + // Never throws (a failed insert is logged). Await it so the next budget check + // sees this call; callers that must not wait can drop the promise. + await query( `INSERT INTO bedrock_usage (team_id, run_id, model, call_type, input_tokens, output_tokens, estimated_cost_usd) VALUES ($1, $2, $3, $4, $5, $6, $7)`, [teamId, runId, model, callType, inputTokens, outputTokens, cost] @@ -59,3 +62,56 @@ export async function checkAiSpendCap(teamId: string): Promise<{ softCapUsd: softCap, }; } + +export type AiBudget = + | { ok: true } + | { ok: false; status: 402 | 429; error: string; scope: "team_month" | "guest_hourly" | "guest_total" | "guests_daily" }; + +/** + * Checked before EVERY Bedrock call (suggest-transforms, auto-validate, + * chat-builder, generate-data). `calls` is how many calls the caller is about + * to make (the committee makes 3 in parallel). + * + * - Signed-in users: the monthly hard cap in ai_spend_limits (default $200). + * - Guests (GUEST_LIMITS): at most aiCallsPerHour calls in the last hour, + * aiSpendPerGuestUsd in total, and all guests together under + * aiSpendAllGuestsPerDayUsd since 00:00 UTC. One query, counted from + * bedrock_usage, so it holds without Upstash. + * + * Usage is recorded after a call, so concurrent requests can overshoot by the + * cost of the calls in flight; see README "AI spend". + */ +export async function checkAiBudget(teamId: string, calls = 1): Promise { + if (!isGuestId(teamId)) { + const cap = await checkAiSpendCap(teamId); + return cap.blocked + ? { ok: false, status: 402, scope: "team_month", error: `AI spend cap reached ($${cap.currentSpendUsd.toFixed(2)} / $${cap.hardCapUsd} this month). Contact support.` } + : { ok: true }; + } + + const row = await queryOne<{ calls_hour: string; guest_total: string; guests_today: string }>( + `SELECT + (SELECT count(*) FROM bedrock_usage + WHERE team_id = $1 AND created_at > now() - interval '1 hour') AS calls_hour, + (SELECT COALESCE(SUM(estimated_cost_usd), 0) FROM bedrock_usage + WHERE team_id = $1) AS guest_total, + (SELECT COALESCE(SUM(estimated_cost_usd), 0) FROM bedrock_usage + WHERE team_id LIKE 'guest\\_%' AND created_at >= date_trunc('day', now() AT TIME ZONE 'UTC') AT TIME ZONE 'UTC') AS guests_today`, + [teamId] + ); + const callsHour = Number(row?.calls_hour ?? 0); + const guestTotal = Number(row?.guest_total ?? 0); + const guestsToday = Number(row?.guests_today ?? 0); + + if (guestsToday >= GUEST_LIMITS.aiSpendAllGuestsPerDayUsd) { + log.warn("daily AI budget for all guests reached", { guests_today_usd: guestsToday }); + return { ok: false, status: 429, scope: "guests_daily", error: "Guest AI capacity is used up for today. Sign up to continue, or try again tomorrow." }; + } + if (guestTotal >= GUEST_LIMITS.aiSpendPerGuestUsd) { + return { ok: false, status: 402, scope: "guest_total", error: "This guest session has used its AI allowance. Sign up to continue." }; + } + if (callsHour + calls > GUEST_LIMITS.aiCallsPerHour) { + return { ok: false, status: 429, scope: "guest_hourly", error: `Guests can make ${GUEST_LIMITS.aiCallsPerHour} AI calls per hour. Try again later or sign up.` }; + } + return { ok: true }; +} diff --git a/src/lib/billing.test.ts b/src/lib/billing.test.ts new file mode 100644 index 0000000..7ae40a7 --- /dev/null +++ b/src/lib/billing.test.ts @@ -0,0 +1,41 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const h = vi.hoisted(() => ({ queryOne: vi.fn() })); +vi.mock("@/lib/db", () => ({ queryOne: (...a: unknown[]) => h.queryOne(...a) })); + +import { checkQuota, isAdmin, quotaBlockedMessage } from "./billing"; + +const GUEST = "guest_AAAAAAAAAAAAAAAAAAAAAA"; +beforeEach(() => vi.clearAllMocks()); + +describe("guest quota", () => { + it("is 10,000 rows all-time, hard-capped, ignoring runs refused for size", async () => { + h.queryOne.mockResolvedValue({ total: "9999" }); + expect(await checkQuota(GUEST, null, GUEST)).toMatchObject({ plan: "guest", includedRows: 10_000, used: 9999, remaining: 1, blocked: false, hardCap: true }); + const [sql, params] = h.queryOne.mock.calls[0]; + expect(sql).not.toContain("created_at >="); + expect(sql).toContain("row_count_raw <= $2"); + expect(params).toEqual([GUEST, 5000]); + h.queryOne.mockResolvedValue({ total: "10000" }); + expect((await checkQuota(GUEST, null, GUEST)).blocked).toBe(true); + }); + + it("never treats a guest as admin or reads subscriptions", async () => { + expect(isAdmin("someone@example.com", GUEST)).toBe(false); + h.queryOne.mockResolvedValue({ total: "0" }); + await checkQuota(GUEST, null, GUEST); + expect(h.queryOne.mock.calls.some(([sql]) => String(sql).includes("subscriptions"))).toBe(false); + }); + + it("leaves signed-in users on their plan", async () => { + h.queryOne.mockImplementation(async (sql: string) => (sql.includes("subscriptions") ? null : { total: "10" })); + expect(await checkQuota("user_1", null, "user_1")).toMatchObject({ plan: "free", includedRows: 50_000, used: 10 }); + }); + + it("guest and user messages", () => { + const base = { used: 10000, includedRows: 10000, remaining: 0, hardCap: true, blocked: true, isAdmin: false }; + expect(quotaBlockedMessage({ ...base, plan: "guest" })).toBe("Guest row limit reached (10,000 / 10,000 rows). Sign up to keep going."); + expect(quotaBlockedMessage({ ...base, used: 50000, includedRows: 50000, plan: "free" })) + .toBe("Monthly row limit reached (50,000 / 50,000 rows on free plan). Upgrade at /pricing to continue."); + }); +}); diff --git a/src/lib/billing.ts b/src/lib/billing.ts index 2256efd..617f96e 100644 --- a/src/lib/billing.ts +++ b/src/lib/billing.ts @@ -1,4 +1,6 @@ import { queryOne } from "@/lib/db"; +import { isGuestId } from "@/lib/guest"; +import { GUEST_LIMITS } from "@/lib/guest-limits"; export const PLANS = { free: { name: "Free", includedRows: 50_000, overagePer100k: null, hardCap: true }, @@ -22,7 +24,7 @@ export interface Subscription { } export interface QuotaResult { - plan: PlanId | "admin"; + plan: PlanId | "admin" | "guest"; includedRows: number; used: number; remaining: number; @@ -46,6 +48,7 @@ const ADMIN_USER_IDS = new Set( ); export function isAdmin(email: string | null | undefined, userId?: string | null): boolean { + if (userId && isGuestId(userId)) return false; if (userId && ADMIN_USER_IDS.has(userId)) return true; return !!email && ADMIN_EMAILS.has(email.toLowerCase()); } @@ -75,6 +78,22 @@ export async function getMonthlyUsage(teamId: string): Promise { return Number(result?.total ?? 0); } +/** + * Rows a guest has used, all time (a guest lives 24 h), every pass. Runs over + * GUEST_LIMITS.rowsPerRun are refused before any AI call, so they do not count. + */ +export async function getGuestUsage(guestId: string): Promise { + const result = await queryOne<{ total: string }>( + `SELECT COALESCE(SUM(pr.row_count_raw), 0) AS total + FROM pipeline_runs pr + JOIN pipelines p ON pr.pipeline_id = p.id + WHERE p.team_id = $1 + AND pr.row_count_raw <= $2`, + [guestId, GUEST_LIMITS.rowsPerRun] + ); + return Number(result?.total ?? 0); +} + export async function checkQuota( teamId: string, email?: string | null, @@ -92,6 +111,20 @@ export async function checkQuota( }; } + if (isGuestId(teamId)) { + const used = await getGuestUsage(teamId); + const included = GUEST_LIMITS.rowsPerGuest; + return { + plan: "guest", + includedRows: included, + used, + remaining: Math.max(0, included - used), + hardCap: true, + blocked: used >= included, + isAdmin: false, + }; + } + const sub = await getSubscription(teamId); const planId: PlanId = (sub?.plan as PlanId) ?? "free"; const planConfig = PLANS[planId]; @@ -110,3 +143,11 @@ export async function checkQuota( isAdmin: false, }; } + +/** The message shown when quota.blocked. */ +export function quotaBlockedMessage(quota: QuotaResult): string { + const counts = `${quota.used.toLocaleString("en-US")} / ${quota.includedRows.toLocaleString("en-US")} rows`; + return quota.plan === "guest" + ? `Guest row limit reached (${counts}). Sign up to keep going.` + : `Monthly row limit reached (${counts} on ${quota.plan} plan). Upgrade at /pricing to continue.`; +} diff --git a/src/lib/guest-guard.ts b/src/lib/guest-guard.ts new file mode 100644 index 0000000..6be432a --- /dev/null +++ b/src/lib/guest-guard.ts @@ -0,0 +1,17 @@ +/** + * Route-level guest checks (server only). The middleware already blocks the + * paths in GUEST_BLOCKED_PATHS; routes call forbidGuest() too so a matcher + * mistake cannot expose a feature. + */ +import { NextResponse } from "next/server"; +import { isGuestId } from "@/lib/guest"; +import { GUEST_BLOCKED_MESSAGE } from "@/lib/guest-limits"; + +export function guestForbiddenResponse(): NextResponse { + return NextResponse.json({ error: GUEST_BLOCKED_MESSAGE, guest: true }, { status: 403 }); +} + +/** 403 for a guest id, null otherwise. */ +export function forbidGuest(userId: string | null | undefined): NextResponse | null { + return userId && isGuestId(userId) ? guestForbiddenResponse() : null; +} diff --git a/src/lib/guest-limits.test.ts b/src/lib/guest-limits.test.ts new file mode 100644 index 0000000..3e31474 --- /dev/null +++ b/src/lib/guest-limits.test.ts @@ -0,0 +1,21 @@ +import { describe, expect, it } from "vitest"; +import { GUEST_LIMITS, isGuestBlockedPath } from "./guest-limits"; + +describe("guest limits", () => { + it("match the approved numbers", () => { + expect(GUEST_LIMITS).toMatchObject({ + maxUploadBytes: 2 * 1024 * 1024, rowsPerRun: 5_000, rowsPerGuest: 10_000, uploadsPerGuest: 3, + uploadsPerIpPerDay: 10, aiCallsPerHour: 10, aiSpendPerGuestUsd: 0.25, aiSpendAllGuestsPerDayUsd: 5, + }); + }); + + it.each([ + ["/api/chat-builder", true], ["/api/chat-builder/generate-data", true], ["/api/runs/abc/auto-clean", true], + ["/api/export-training/abc", true], ["/api/alerts/configure", true], ["/api/account", true], + ["/api/templates", true], ["/api/templates/x/use", true], ["/api/admin/set-plan", true], ["/templates", true], + ["/api/runs/abc/iterate", false], ["/api/upload", false], ["/api/pipelines", false], ["/dashboard", false], + ["/api/chat-builderx", false], ["/templatesx", false], ["/api/runs/abc/auto-clean/x", false], + ])("%s blocked=%s", (path, blocked) => { + expect(isGuestBlockedPath(path)).toBe(blocked); + }); +}); diff --git a/src/lib/guest-limits.ts b/src/lib/guest-limits.ts index ac8204c..1303696 100644 --- a/src/lib/guest-limits.ts +++ b/src/lib/guest-limits.ts @@ -31,3 +31,28 @@ export const GUEST_LIMITS = { /** New guest sessions per 24 h, all IPs. */ guestsPerDay: 200, } as const; + +/** + * Features a guest cannot use: they cost money with no row/AI accounting + * (chat builder, synthetic data, auto-clean's unattended passes), send data + * elsewhere (Slack alerts, training export), are irreversible account actions + * (deletion; a guest session simply expires), or are shared/admin surfaces + * (templates, /api/admin). The middleware answers 403 (API) or redirects + * (pages) for these, and each route also refuses guest ids (forbidGuest). + */ +export const GUEST_BLOCKED_PATHS: readonly RegExp[] = [ + /^\/api\/chat-builder(\/|$)/, + /^\/api\/runs\/[^/]+\/auto-clean\/?$/, + /^\/api\/export-training(\/|$)/, + /^\/api\/alerts(\/|$)/, + /^\/api\/account(\/|$)/, + /^\/api\/templates(\/|$)/, + /^\/api\/admin(\/|$)/, + /^\/templates(\/|$)/, +]; + +export function isGuestBlockedPath(pathname: string): boolean { + return GUEST_BLOCKED_PATHS.some((re) => re.test(pathname)); +} + +export const GUEST_BLOCKED_MESSAGE = "This feature is not available in guest mode. Sign up to use it."; diff --git a/src/lib/guest-quota.ts b/src/lib/guest-quota.ts new file mode 100644 index 0000000..7950c15 --- /dev/null +++ b/src/lib/guest-quota.ts @@ -0,0 +1,66 @@ +/** + * Guest upload and pipeline caps, counted in Postgres (server only). See + * GUEST_LIMITS in src/lib/guest-limits.ts. + */ +import { queryOne } from "@/lib/db"; +import { GUEST_LIMITS } from "@/lib/guest-limits"; +import type { PipelineRun } from "@/lib/types"; + +/** + * Creates a guest's pending first-pass run only while (a) the guest session + * exists and has not expired, (b) the guest has fewer than uploadsPerGuest + * first-pass runs in 24 h and (c) all guests created from the same IP hash + * have fewer than uploadsPerIpPerDay. One statement, so a purged or expired + * session cannot upload. Two concurrent requests can both pass (READ + * COMMITTED); the per-user Upstash upload limit and the row cap bound that. + */ +export const GUEST_RUN_INSERT_SQL = ` +INSERT INTO pipeline_runs (id, pipeline_id, status, file_format, raw_s3_key, started_at) +SELECT $1, $2, 'pending', $3, $4, now() +FROM guest_sessions g +WHERE g.id = $5 + AND g.expires_at > now() + AND (SELECT count(*) FROM pipeline_runs pr JOIN pipelines p ON p.id = pr.pipeline_id + WHERE p.team_id = $5 AND pr.iteration = 1 + AND pr.created_at > now() - interval '24 hours') < $6 + AND (SELECT count(*) FROM pipeline_runs pr JOIN pipelines p ON p.id = pr.pipeline_id + JOIN guest_sessions gs ON gs.id = p.team_id + WHERE gs.ip_hash = g.ip_hash AND pr.iteration = 1 + AND pr.created_at > now() - interval '24 hours') < $7 +RETURNING *`; + +export async function insertGuestRun(opts: { + runId: string; pipelineId: string; fileFormat: string; s3Key: string; guestId: string; +}): Promise { + return queryOne(GUEST_RUN_INSERT_SQL, [ + opts.runId, opts.pipelineId, opts.fileFormat, opts.s3Key, opts.guestId, + GUEST_LIMITS.uploadsPerGuest, GUEST_LIMITS.uploadsPerIpPerDay, + ]); +} + +/** Why insertGuestRun refused (for the error message). */ +export async function guestUploadRefusal(guestId: string): Promise { + const row = await queryOne<{ live: boolean | null; mine: string; ip: string }>( + `SELECT + (SELECT expires_at > now() FROM guest_sessions WHERE id = $1) AS live, + (SELECT count(*) FROM pipeline_runs pr JOIN pipelines p ON p.id = pr.pipeline_id + WHERE p.team_id = $1 AND pr.iteration = 1 AND pr.created_at > now() - interval '24 hours') AS mine, + (SELECT count(*) FROM pipeline_runs pr JOIN pipelines p ON p.id = pr.pipeline_id + JOIN guest_sessions gs ON gs.id = p.team_id + WHERE gs.ip_hash = (SELECT ip_hash FROM guest_sessions WHERE id = $1) + AND pr.iteration = 1 AND pr.created_at > now() - interval '24 hours') AS ip`, + [guestId] + ); + if (!row?.live) return "Your guest session has expired. Start a new one or sign up."; + if (Number(row.mine) >= GUEST_LIMITS.uploadsPerGuest) { + return `Guests can upload ${GUEST_LIMITS.uploadsPerGuest} files. Sign up to keep going.`; + } + return `Guest upload limit for your network reached (${GUEST_LIMITS.uploadsPerIpPerDay} per day). Sign up to keep going.`; +} + +/** Inserts a guest pipeline only while the guest has fewer than pipelinesPerGuest. */ +export const GUEST_PIPELINE_INSERT_SQL = ` +INSERT INTO pipelines (name, description, owner_id, team_id) +SELECT $1, $2, $3, $3 +WHERE (SELECT count(*) FROM pipelines WHERE team_id = $3) < $4 +RETURNING *`; diff --git a/src/lib/migrations/003_guest_sessions.sql b/src/lib/migrations/003_guest_sessions.sql index 833d90a..2f3e52b 100644 --- a/src/lib/migrations/003_guest_sessions.sql +++ b/src/lib/migrations/003_guest_sessions.sql @@ -20,3 +20,8 @@ CREATE INDEX IF NOT EXISTS idx_guest_sessions_created CREATE INDEX IF NOT EXISTS idx_guest_sessions_expires ON guest_sessions (expires_at); + +-- checkAiBudget (src/lib/bedrock-meter.ts) sums today's spend of all guests +-- (team_id LIKE 'guest\_%' AND created_at >= today) before every guest AI call. +CREATE INDEX IF NOT EXISTS idx_bedrock_usage_created + ON bedrock_usage (created_at); diff --git a/src/lib/schema.sql b/src/lib/schema.sql index 17e991d..06b2a05 100644 --- a/src/lib/schema.sql +++ b/src/lib/schema.sql @@ -141,6 +141,8 @@ CREATE TABLE IF NOT EXISTS bedrock_usage ( CREATE INDEX IF NOT EXISTS idx_bedrock_usage_team_created ON bedrock_usage (team_id, created_at); +CREATE INDEX IF NOT EXISTS idx_bedrock_usage_created + ON bedrock_usage (created_at); CREATE TABLE IF NOT EXISTS ai_spend_limits ( team_id TEXT PRIMARY KEY, diff --git a/src/middleware.guest.test.ts b/src/middleware.guest.test.ts index f444e3e..2044bdf 100644 --- a/src/middleware.guest.test.ts +++ b/src/middleware.guest.test.ts @@ -39,4 +39,27 @@ describe("middleware: guests", () => { delete process.env.GUEST_COOKIE_SECRET; expect((await run("/api/pipelines", valid)).status).toBe(401); }); + + it("guests are blocked from chat, auto-clean, training export, alerts, account, templates and admin", async () => { + const cookie = `${GUEST_COOKIE}=${await guestToken()}`; + for (const path of ["/api/chat-builder", "/api/chat-builder/generate-data", "/api/runs/r1/auto-clean", "/api/export-training/r1", + "/api/alerts/configure", "/api/account", "/api/templates", "/api/templates/t1/use", "/api/admin/ai-spend"]) { + const res = await run(path, cookie); + expect(res.status, path).toBe(403); + expect((await res.json()).guest, path).toBe(true); + } + const page = await run("/templates", cookie); + expect(page.headers.get("location")).toBe("http://localhost:3000/dashboard?guest_blocked=1"); + // Not blocked: the normal flow, including later passes. + for (const path of ["/api/runs/r1/iterate", "/api/run-status/r1", "/api/download/r1"]) { + expect((await run(path, cookie)).headers.get("x-middleware-next"), path).toBe("1"); + } + }); + + it("a signed-in user with a stale guest cookie is not blocked", async () => { + const cookie = `${GUEST_COOKIE}=${await guestToken()}; cs_test_user=user_test_alice`; + for (const path of ["/api/chat-builder", "/templates", "/api/account"]) { + expect((await run(path, cookie)).headers.get("x-middleware-next"), path).toBe("1"); + } + }); }); diff --git a/src/middleware.ts b/src/middleware.ts index 8dadaf3..f4044be 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -2,6 +2,7 @@ import { clerkMiddleware, createRouteMatcher } from "@clerk/nextjs/server"; import { NextResponse, type NextFetchEvent, type NextRequest } from "next/server"; import { isTestMode } from "@/lib/test-mode"; import { GUEST_COOKIE, guestFromCookie } from "@/lib/guest"; +import { GUEST_BLOCKED_MESSAGE, isGuestBlockedPath } from "@/lib/guest-limits"; const isProtectedRoute = createRouteMatcher([ "/dashboard(.*)", @@ -27,13 +28,25 @@ async function hasGuestSession(req: NextRequest): Promise { return (await guestFromCookie(req.cookies.get(GUEST_COOKIE)?.value)) !== null; } +/** A guest on a feature guests cannot use: 403 JSON for the API, back to the dashboard for pages. */ +function guestBlocked(req: NextRequest): NextResponse { + if (req.nextUrl.pathname.startsWith("/api/")) { + return NextResponse.json({ error: GUEST_BLOCKED_MESSAGE, guest: true }, { status: 403 }); + } + return NextResponse.redirect(new URL("/dashboard?guest_blocked=1", req.url)); +} + const clerkHandler = clerkMiddleware(async (auth, req) => { - if (isProtectedRoute(req)) { - // Clerk users first; otherwise a guest may pass, otherwise Clerk's normal redirect / 401. + if (isProtectedRoute(req) || isGuestBlockedPath(req.nextUrl.pathname)) { + // Clerk users first; otherwise a guest may pass (except blocked features), + // otherwise Clerk's normal redirect / 401. const { userId } = await auth(); if (userId) return; - if (await hasGuestSession(req)) return; - await auth.protect(); + if (await hasGuestSession(req)) { + if (isGuestBlockedPath(req.nextUrl.pathname)) return guestBlocked(req); + return; + } + if (isProtectedRoute(req)) await auth.protect(); return; } @@ -57,7 +70,10 @@ const TEST_USER_ID = /^user_test_[A-Za-z0-9_-]{1,64}$/; async function testModeMiddleware(req: NextRequest) { const userId = req.cookies.get(TEST_USER_COOKIE)?.value; - const signedIn = (userId !== undefined && TEST_USER_ID.test(userId)) || (await hasGuestSession(req)); + const testUser = userId !== undefined && TEST_USER_ID.test(userId); + const guest = !testUser && (await hasGuestSession(req)); + const signedIn = testUser || guest; + if (guest && isGuestBlockedPath(req.nextUrl.pathname)) return guestBlocked(req); if (isProtectedRoute(req) && !signedIn) { if (req.nextUrl.pathname.startsWith("/api/")) { return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); diff --git a/tests/e2e/guest.spec.ts b/tests/e2e/guest.spec.ts index c1910d2..9e54b33 100644 --- a/tests/e2e/guest.spec.ts +++ b/tests/e2e/guest.spec.ts @@ -2,7 +2,8 @@ * Guest access in a real browser (production build, test mode): POST /api/guest * sets the signed cs_guest cookie, the guest can open protected pages, sees a * 2 MB per-file limit, uploads through the presigned POST and reaches review. - * Each test uses its own X-Forwarded-For so the per-IP cap never interferes. + * Guests are kept out of blocked features and capped at 3 uploads. Each test + * uses its own X-Forwarded-For so the per-IP caps never interfere. */ import { expect, test, type Page } from "@playwright/test"; import { join } from "node:path"; @@ -11,7 +12,7 @@ import { randomUUID } from "node:crypto"; const FIXTURE = join(__dirname, "fixtures", "orders.csv"); async function startGuest(page: Page) { - const res = await page.request.post("/api/guest", { headers: { "x-forwarded-for": `192.0.2.${Math.floor(Math.random() * 250) + 1}` } }); + const res = await page.request.post("/api/guest", { headers: { "x-forwarded-for": `e2e-${randomUUID()}` } }); expect(res.status()).toBe(201); const cookie = (await page.context().cookies()).find((c) => c.name === "cs_guest"); expect(cookie?.httpOnly).toBe(true); @@ -55,3 +56,22 @@ test("a forged or tampered guest cookie is not a session", async ({ request }) = expect(res.status()).toBe(401); } }); + +test("guests are kept out of blocked features and capped at 3 uploads", async ({ page }) => { + await startGuest(page); + await page.goto("/templates"); + await expect(page).toHaveURL(/\/dashboard\?guest_blocked=1$/); + for (const [method, path] of [["POST", "/api/chat-builder"], ["POST", "/api/chat-builder/generate-data"], ["GET", "/api/templates"], + ["POST", "/api/alerts/configure"], ["DELETE", "/api/account?confirm=true"], ["GET", "/api/export-training/00000000-0000-0000-0000-000000000000"]] as const) { + const res = await page.request.fetch(path, { method, data: method === "GET" ? undefined : {} }); + expect(res.status(), `${method} ${path}`).toBe(403); + expect((await res.json()).guest).toBe(true); + } + + const { pipeline } = await (await page.request.post("/api/pipelines", { data: { name: "guest caps" } })).json(); + const start = () => page.request.post("/api/upload", { data: { pipeline_id: pipeline.id, filename: "a.csv", size: 100 } }); + for (let i = 0; i < 3; i++) expect((await start()).status()).toBe(200); + const fourth = await start(); + expect(fourth.status()).toBe(429); + expect((await fourth.json()).error).toBe("Guests can upload 3 files. Sign up to keep going."); +}); diff --git a/tests/integration/guest-limits.test.ts b/tests/integration/guest-limits.test.ts new file mode 100644 index 0000000..4d956de --- /dev/null +++ b/tests/integration/guest-limits.test.ts @@ -0,0 +1,200 @@ +/** + * Guest quotas against real Postgres: the capped inserts for uploads and + * pipelines, the rows-per-run check, and checkAiBudget before suggest-transforms + * (fake model) and auto-validate. Every number comes from GUEST_LIMITS. + */ +import { afterAll, describe, expect, it, vi } from "vitest"; +import { randomUUID } from "node:crypto"; + +vi.mock("@/lib/auth", async () => (await import("./helpers")).authMock); + +import { authState, db, env, request } from "./helpers"; +import { POST as startGuest } from "@/app/api/guest/route"; +import { POST as createPipeline } from "@/app/api/pipelines/route"; +import { POST as upload } from "@/app/api/upload/route"; +import { POST as suggestTransforms } from "@/app/api/suggest-transforms/route"; +import { POST as autoValidate } from "@/app/api/auto-validate/[runId]/route"; +import { GUEST_COOKIE, guestFromCookie } from "@/lib/guest"; +import { GUEST_LIMITS } from "@/lib/guest-limits"; +import { closePgPool } from "@/lib/db-pg"; + +const guests: string[] = []; +const fakeGuests: string[] = []; +afterAll(async () => { + const all = [...guests, ...fakeGuests]; + await db.query("DELETE FROM bedrock_usage WHERE team_id = ANY($1)", [all]); + await db.query("DELETE FROM pipelines WHERE team_id = ANY($1)", [all]); + await db.query("DELETE FROM guest_sessions WHERE id = ANY($1)", [all]); + await closePgPool(); + await db.end(); +}); + +const uniqueIp = () => `198.18.${Math.floor(Math.random() * 250)}.${Math.floor(Math.random() * 250)}-${randomUUID().slice(0, 6)}`; +async function newGuest(ip = uniqueIp()): Promise { + const res = await startGuest(request("/api/guest", { method: "POST", body: {}, headers: { "x-forwarded-for": ip } })); + expect(res.status).toBe(201); + const id = (await guestFromCookie(res.cookies.get(GUEST_COOKIE)!.value))!.guestId; + guests.push(id); + return id; +} +async function pipelineFor(guest: string): Promise { + authState.userId = guest; + const res = await createPipeline(request("/api/pipelines", { method: "POST", body: { name: "g" } })); + expect(res.status).toBe(201); + return (await res.json()).pipeline.id; +} +const startUpload = (pipelineId: string) => + upload(request("/api/upload", { method: "POST", body: { pipeline_id: pipelineId, filename: "a.csv", size: 100 } })); + +/** A guest run at awaiting_ai with a raw profile, as the profiler leaves it. */ +async function awaitingAiRun(guest: string, rows: number, opts: { auto?: boolean } = {}): Promise { + const pipelineId = await pipelineFor(guest); + const runId = randomUUID(); + await db.query( + `INSERT INTO pipeline_runs (id, pipeline_id, status, file_format, raw_s3_key, row_count_raw, auto_mode) + VALUES ($1, $2, 'awaiting_ai', 'csv', $3, $4, $5)`, + [runId, pipelineId, `${guest}/${pipelineId}/${runId}/raw.csv`, rows, opts.auto ?? false] + ); + await db.query( + `INSERT INTO data_profiles (run_id, stage, quality_score, total_rows, null_percentage, duplicate_percentage, column_stats) + VALUES ($1, 'raw', 80, $2, 0, 0, $3)`, + [runId, rows, JSON.stringify({ amount: { dtype: "object", null_count: 0, unique_count: 3, sample_values: ["$1", "$2"] } })] + ); + return runId; +} +const suggest = (runId: string) => suggestTransforms(request("/api/suggest-transforms", { + method: "POST", body: { run_id: runId }, headers: { "x-webhook-secret": env("WEBHOOK_SECRET") }, +})); +const usage = async (guest: string) => + Number((await db.query("SELECT count(*) AS n FROM bedrock_usage WHERE team_id = $1", [guest])).rows[0].n); +const run = async (id: string) => (await db.query("SELECT status, error_message FROM pipeline_runs WHERE id = $1", [id])).rows[0]; +async function spend(teamId: string, calls: number, usd: number) { + for (let i = 0; i < calls; i++) { + await db.query("INSERT INTO bedrock_usage (team_id, model, call_type, estimated_cost_usd) VALUES ($1, 'm', 'seed', $2)", [teamId, usd / calls]); + } +} + +describe("guest upload and pipeline caps", () => { + it(`allows ${GUEST_LIMITS.uploadsPerGuest} uploads per guest, then 429`, async () => { + const guest = await newGuest(); + const p = await pipelineFor(guest); + for (let i = 0; i < GUEST_LIMITS.uploadsPerGuest; i++) expect((await startUpload(p)).status).toBe(200); + const res = await startUpload(p); + expect(res.status).toBe(429); + expect((await res.json()).error).toBe("Guests can upload 3 files. Sign up to keep going."); + }); + + it(`allows ${GUEST_LIMITS.uploadsPerIpPerDay} uploads per IP across guests, then 429`, async () => { + const ip = uniqueIp(); + let ok = 0; + for (let g = 0; g < 4; g++) { + const p = await pipelineFor(await newGuest(ip)); + for (let i = 0; i < GUEST_LIMITS.uploadsPerGuest && ok < GUEST_LIMITS.uploadsPerIpPerDay; i++) { + expect((await startUpload(p)).status).toBe(200); + ok++; + } + if (ok === GUEST_LIMITS.uploadsPerIpPerDay) { + const res = await startUpload(p); + expect(res.status).toBe(429); + expect((await res.json()).error).toMatch(/for your network reached \(10 per day\)/); + } + } + expect(ok).toBe(10); + }); + + it("an expired guest session cannot upload even with a valid cookie", async () => { + const guest = await newGuest(); + const p = await pipelineFor(guest); + await db.query("UPDATE guest_sessions SET expires_at = now() - interval '1 second' WHERE id = $1", [guest]); + const res = await startUpload(p); + expect(res.status).toBe(429); + expect((await res.json()).error).toMatch(/expired/); + }); + + it(`caps pipelines at ${GUEST_LIMITS.pipelinesPerGuest}`, async () => { + const guest = await newGuest(); + for (let i = 0; i < GUEST_LIMITS.pipelinesPerGuest; i++) await pipelineFor(guest); + const res = await createPipeline(request("/api/pipelines", { method: "POST", body: { name: "one too many" } })); + expect(res.status).toBe(429); + }); +}); + +describe("guest rows and AI budget before suggest-transforms", () => { + it("a guest run under the limits gets AI rules and is metered", async () => { + const guest = await newGuest(); + const runId = await awaitingAiRun(guest, 4); + const res = await suggest(runId); + expect(res.status).toBe(200); + expect((await run(runId)).status).toBe("awaiting_approval"); + expect(await usage(guest)).toBe(1); // meterBedrockCall is awaited, so the next check sees it + }); + + it(`fails a run over ${GUEST_LIMITS.rowsPerRun} rows before any AI call`, async () => { + const guest = await newGuest(); + const runId = await awaitingAiRun(guest, GUEST_LIMITS.rowsPerRun + 1); + expect((await suggest(runId)).status).toBe(402); + expect(await run(runId)).toEqual({ status: "failed", error_message: "Guest runs are limited to 5,000 rows; this file has 5,001. Sign up to clean larger files." }); + expect(await usage(guest)).toBe(0); + }); + + it(`stops at ${GUEST_LIMITS.rowsPerGuest} rows per guest, and oversized runs do not count`, async () => { + const guest = await newGuest(); + await awaitingAiRun(guest, 50_000); // refused for size: not counted + const first = await awaitingAiRun(guest, 5000); + expect((await suggest(first)).status).toBe(200); + const second = await awaitingAiRun(guest, 5000); + // used = 10,000 including this run's rows (the profiler already recorded them) + expect((await suggest(second)).status).toBe(402); + expect((await run(second)).error_message).toBe("Guest row limit reached (10,000 / 10,000 rows). Sign up to keep going."); + }); + + it(`refuses the call after ${GUEST_LIMITS.aiCallsPerHour} AI calls in an hour`, async () => { + const guest = await newGuest(); + await spend(guest, GUEST_LIMITS.aiCallsPerHour, 0.01); + const runId = await awaitingAiRun(guest, 4); + expect((await suggest(runId)).status).toBe(429); + expect((await run(runId)).error_message).toBe("Guests can make 10 AI calls per hour. Try again later or sign up."); + expect(await usage(guest)).toBe(GUEST_LIMITS.aiCallsPerHour); + }); + + it(`refuses at $${GUEST_LIMITS.aiSpendPerGuestUsd} spent by the guest`, async () => { + const guest = await newGuest(); + await spend(guest, 1, GUEST_LIMITS.aiSpendPerGuestUsd); + const runId = await awaitingAiRun(guest, 4); + expect((await suggest(runId)).status).toBe(402); + expect((await run(runId)).error_message).toBe("This guest session has used its AI allowance. Sign up to continue."); + }); + + it(`refuses every guest once all guests spent $${GUEST_LIMITS.aiSpendAllGuestsPerDayUsd} today, but not users`, async () => { + const other = `guest_${randomUUID().replace(/-/g, "").slice(0, 22)}`; + fakeGuests.push(other); + await spend(other, 1, GUEST_LIMITS.aiSpendAllGuestsPerDayUsd); + try { + const guest = await newGuest(); + const runId = await awaitingAiRun(guest, 4); + expect((await suggest(runId)).status).toBe(429); + expect((await run(runId)).error_message).toMatch(/Guest AI capacity is used up for today/); + } finally { + await db.query("DELETE FROM bedrock_usage WHERE team_id = $1", [other]); + } + }); + + it("auto-validate checks the budget for its 3 calls and falls back to manual review", async () => { + const guest = await newGuest(); + const runId = await awaitingAiRun(guest, 4, { auto: true }); + await db.query("UPDATE pipeline_runs SET status = 'queued' WHERE id = $1", [runId]); + await db.query( + "INSERT INTO transform_rules (pipeline_id, run_id, rule_type, column_name, parameters, status, order_index) SELECT pipeline_id, id, 'trim_whitespace', 'amount', '{}', 'pending', 0 FROM pipeline_runs WHERE id = $1", + [runId] + ); + await spend(guest, GUEST_LIMITS.aiCallsPerHour - 2, 0.001); // 8 used: 3 more would exceed 10 + const res = await autoValidate( + request(`/api/auto-validate/${runId}`, { method: "POST", headers: { "x-webhook-secret": env("WEBHOOK_SECRET") } }), + { params: Promise.resolve({ runId }) } + ); + expect(res.status).toBe(200); + expect(await res.json()).toMatchObject({ ok: true, fallback: "manual_review", scope: "guest_hourly" }); + expect((await run(runId)).status).toBe("awaiting_approval"); + expect(await usage(guest)).toBe(GUEST_LIMITS.aiCallsPerHour - 2); + }); +});