From c7f43833b82045d3569367f06197fe10eda397f4 Mon Sep 17 00:00:00 2001 From: Devarshi Shimpi Date: Wed, 11 Nov 2026 16:51:13 +0530 Subject: [PATCH] feat: add quota controls and harden review processing --- .github/workflows/trigger-docs-deploy.yml | 2 +- CONTRIBUTING.md | 15 +- README.md | 2 +- apps/node/src/env.ts | 4 + apps/node/src/index.ts | 25 ++- apps/node/src/runtime.ts | 24 ++- apps/worker/src/adapters/index.ts | 32 ++- apps/worker/src/env.ts | 2 + docs/contributing.mdx | 6 +- package.json | 2 + packages/api/src/create-deps.ts | 18 +- packages/api/src/index.ts | 1 + packages/api/src/platform/index.ts | 1 + packages/api/src/platform/quota.ts | 65 ++++++ packages/core/src/constants.ts | 3 + packages/core/src/index.ts | 12 ++ packages/core/src/token-tracker.ts | 24 ++- packages/db/src/client.ts | 36 +++- .../node-adapters/src/node-orchestrator.ts | 2 +- packages/provider-github/src/index.ts | 1 + packages/provider-github/src/review-post.ts | 192 ++++++++++++------ scripts/test.mjs | 73 +++++-- test/api/default-quota.spec.ts | 161 +++++++++++++++ test/db/connection-pool-reuse.spec.ts | 78 +++++++ test/helpers.ts | 6 +- test/review/comments.spec.ts | 65 +++++- test/review/subrequest-limits.spec.ts | 75 +++++++ test/setup.ts | 13 +- test/token-tracker.spec.ts | 26 +++ vitest.config.ts | 37 +++- 30 files changed, 897 insertions(+), 106 deletions(-) create mode 100644 packages/api/src/platform/quota.ts create mode 100644 test/api/default-quota.spec.ts create mode 100644 test/db/connection-pool-reuse.spec.ts create mode 100644 test/review/subrequest-limits.spec.ts diff --git a/.github/workflows/trigger-docs-deploy.yml b/.github/workflows/trigger-docs-deploy.yml index 47870817..93ed4f25 100644 --- a/.github/workflows/trigger-docs-deploy.yml +++ b/.github/workflows/trigger-docs-deploy.yml @@ -3,7 +3,7 @@ name: Trigger Docs Deploy on: push: branches: - - dev + - main paths: - "docs/**" diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e246d0a5..8045e3c4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -104,7 +104,7 @@ npm run typecheck ## 📝 Contributing to Documentation -All documentation content for [codra.run/docs](https://codra.run/docs) lives in `docs/*.mdx` in this repository. Even though the website repo is private, **all documentation is 100% open source** and deploys automatically when PRs are merged into `dev`. +All documentation content for [codra.run/docs](https://codra.run/docs) lives in `docs/*.mdx` in this repository. Even though the website repo is private, **all documentation is 100% open source** and deploys automatically when PRs are merged into `main`. ### MDX Guidelines - **Files & Routes**: Each file in `docs/*.mdx` maps directly to a route (e.g. `docs/docker.mdx` -> `codra.run/docs/docker`). @@ -118,16 +118,15 @@ All documentation content for [codra.run/docs](https://codra.run/docs) lives in ## 🚀 Pull Request Process -Contributions are merged into `dev` first and reach `main` when a release is cut, so `main` always -reflects what is deployed and published. Pull requests opened against `main` will be asked to -retarget to `dev`. +Pull requests are reviewed and merged directly into `main`. There is no separate `dev` integration +branch or release-branch promotion step. -1. **Fork & Branch**: Create a feature branch from `dev`. +1. **Fork & Branch**: Create a feature branch from `main`. 2. **Atomic Commits**: Keep your commits focused and descriptive. -3. **Sync**: Ensure your branch is up to date with `dev`. -4. **Target Branch**: Open pull requests against `dev`. +3. **Sync**: Ensure your branch is up to date with `main`. +4. **Target Branch**: Open pull requests against `main`. 5. **PR Description**: Use the provided template (if available) or clearly explain the _what_ and _why_ of your changes. -6. **CLA Check**: Once you open the PR, an automated check will verify your CLA status. If you haven't signed yet, follow the link in the check output. +6. **Checks & CLA**: Ensure the required CI checks pass. An automated check verifies your CLA status; if you haven't signed yet, follow the link in the check output. --- diff --git a/README.md b/README.md index bbe3e047..b06d4fdc 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ The full setup and operations guides live at [codra.run/docs](https://codra.run/ ## Contributing -Contributions are welcome. Please read [CONTRIBUTING.md](CONTRIBUTING.md) before opening a pull request against `dev`. Codra uses a Contributor License Agreement for contributions. +Contributions are welcome. Please read [CONTRIBUTING.md](CONTRIBUTING.md) before opening a pull request against `main`. Codra uses a Contributor License Agreement for contributions. ## License diff --git a/apps/node/src/env.ts b/apps/node/src/env.ts index f056f906..b95080a5 100644 --- a/apps/node/src/env.ts +++ b/apps/node/src/env.ts @@ -28,6 +28,8 @@ export interface NodeAppBindings { LLM_CONFIG_ENCRYPTION_KEY: string; BOT_USERNAME: string; ENVIRONMENT: string; + SUBREQUEST_LIMIT?: string | number; + CF_SUBREQUEST_LIMIT?: string | number; } export function createNodeEnv(stubs: { @@ -59,5 +61,7 @@ export function createNodeEnv(stubs: { LLM_CONFIG_ENCRYPTION_KEY: requireEnv("LLM_CONFIG_ENCRYPTION_KEY"), BOT_USERNAME: process.env.BOT_USERNAME || "codra-app", ENVIRONMENT: process.env.ENVIRONMENT || "development", + SUBREQUEST_LIMIT: process.env.SUBREQUEST_LIMIT, + CF_SUBREQUEST_LIMIT: process.env.CF_SUBREQUEST_LIMIT, }; } diff --git a/apps/node/src/index.ts b/apps/node/src/index.ts index ac156f62..534df9ee 100644 --- a/apps/node/src/index.ts +++ b/apps/node/src/index.ts @@ -10,7 +10,8 @@ dotenv.config({ path: path.resolve(process.cwd(), "../../.env.local") }); import { serve } from "@hono/node-server"; import { createApiRouter } from "@codraoss/api"; -import { runWithDb } from "@codraoss/db/client"; +import { runWithDb, getDb, type DbClient } from "@codraoss/db/client"; +import type { DbEnv } from "@codraoss/db/env"; import { NodeOrchestrator, @@ -60,6 +61,13 @@ env.REVIEW_ORCHESTRATOR = new NodeOrchestrator( stubs.REVIEW_QUEUE, ); +const dbEnv: DbEnv = { + ...env, + HYPERDRIVE: env.DATABASE_CONFIG, + workerMode: false, +}; +const dbClient: DbClient = getDb(dbEnv); + const distClientPath = path.resolve(__dirname, "../../../dist/client"); const app = createApiRouter({ serveIndex: async () => @@ -79,8 +87,6 @@ app.use("/*.svg", serveStatic({ root: distClientPath })); app.use("/*.ico", serveStatic({ root: distClientPath })); app.get("*", serveStatic({ root: distClientPath, path: "index.html" })); -app.get("/health", (c) => c.text("OK")); - const port = parseInt(process.env.PORT || "3000", 10); let worker: ReturnType | undefined; @@ -105,8 +111,9 @@ if (process.env.START_API !== "false") { try { return await runWithDb( - { ...env, HYPERDRIVE: env.DATABASE_CONFIG }, + dbEnv, () => app.fetch(request, apiEnv as any), + dbClient, ); } catch (e) { console.error("SERVE ERROR:", e); @@ -136,7 +143,8 @@ if (process.env.START_API !== "false") { import("@codraoss/models").then(({ ModelRegistryService }) => { const service = ModelRegistryService.getInstance(); service.refreshFromRemote().catch(() => {}); - setInterval(() => service.refreshFromRemote().catch(() => {}), 60 * 60 * 1000).unref(); + const timer = setInterval(() => service.refreshFromRemote().catch(() => {}), 60 * 60 * 1000); + (timer as unknown as { unref?: () => void }).unref?.(); }).catch(() => {}); const shutdown = async () => { @@ -156,7 +164,12 @@ const shutdown = async () => { logger.info("Redis client closed."); } - // Postgres client manages its own pool. + // Close Postgres client pool + if (dbClient?.close) { + logger.info("Closing Postgres client..."); + await dbClient.close(); + logger.info("Postgres client closed."); + } logger.info("All services stopped. Exiting."); }; diff --git a/apps/node/src/runtime.ts b/apps/node/src/runtime.ts index 842e9ae8..3dbc791f 100644 --- a/apps/node/src/runtime.ts +++ b/apps/node/src/runtime.ts @@ -1,8 +1,28 @@ import type { ReviewRuntime } from "@codraoss/core/ports"; import type { DbEnv } from "@codraoss/db/env"; -import { TokenTracker } from "@codraoss/core/token-tracker"; +import { + TokenTracker, + type TokenTrackerOptions, +} from "@codraoss/core"; import type { NodeAppBindings } from "./env"; +export function resolveNodeSubrequestLimits( + env?: Partial, +): TokenTrackerOptions { + const envLimit = + env?.SUBREQUEST_LIMIT ?? + env?.CF_SUBREQUEST_LIMIT ?? + process.env.SUBREQUEST_LIMIT ?? + process.env.CF_SUBREQUEST_LIMIT; + if (envLimit !== undefined) { + const limit = Number(envLimit); + if (!Number.isNaN(limit) && limit > 0) { + return { maxSubrequests: limit }; + } + } + return { maxSubrequests: Number.POSITIVE_INFINITY }; +} + import { makeJobStore, makeFileReviewStore, @@ -52,7 +72,7 @@ export function createReviewRuntime( }, telemetry: makeTelemetrySink(env), - createTokenTracker: () => new TokenTracker(), + createTokenTracker: () => new TokenTracker(resolveNodeSubrequestLimits(env)), createGitHub: makeGitHubFactory(env), createModel: makeModelFactory(env), createFormatter: makeFormatterFactory(env), diff --git a/apps/worker/src/adapters/index.ts b/apps/worker/src/adapters/index.ts index 89f87431..df58c07c 100644 --- a/apps/worker/src/adapters/index.ts +++ b/apps/worker/src/adapters/index.ts @@ -1,6 +1,34 @@ import type { ReviewRuntime } from "@codraoss/core/ports"; -import { TokenTracker } from "@codraoss/core/token-tracker"; +import { + TokenTracker, + type TokenTrackerOptions, + DEFAULT_MAX_SUBREQUESTS, + DEFAULT_SUBREQUEST_SAFE_MARGIN, + PAID_WORKER_MAX_SUBREQUESTS, +} from "@codraoss/core"; import type { AppBindings } from "../env"; + +export function resolveWorkerSubrequestLimits( + env: Partial, +): TokenTrackerOptions { + if (env.CF_SUBREQUEST_LIMIT !== undefined) { + const limit = Number(env.CF_SUBREQUEST_LIMIT); + if (!Number.isNaN(limit) && limit > 0) { + return { maxSubrequests: limit }; + } + } + const isPaid = + env.CF_PAID_WORKER === true || + env.CF_PAID_WORKER === "true" || + env.CF_PAID_WORKER === "1"; + if (isPaid) { + return { maxSubrequests: PAID_WORKER_MAX_SUBREQUESTS }; + } + return { + maxSubrequests: DEFAULT_MAX_SUBREQUESTS, + safeMargin: DEFAULT_SUBREQUEST_SAFE_MARGIN, + }; +} // Imported for its side effect as well as never: this installs the AsyncLocalStorage-backed logger as // the sink @codraoss/core's logger facade delegates to. Explicit here so engine log lines carry request // context by construction, rather than because some other module happened to be loaded first. @@ -50,7 +78,7 @@ export function createReviewRuntime(env: AppBindings): ReviewRuntime { repoConfig: makeRepoConfigLoader(env), telemetry: makeTelemetrySink(env), - createTokenTracker: () => new TokenTracker(), + createTokenTracker: () => new TokenTracker(resolveWorkerSubrequestLimits(env)), createGitHub: makeGitHubFactory(env), createModel: makeModelFactory(env), createFormatter: makeFormatterFactory(env), diff --git a/apps/worker/src/env.ts b/apps/worker/src/env.ts index f59aaf29..92424be0 100644 --- a/apps/worker/src/env.ts +++ b/apps/worker/src/env.ts @@ -50,6 +50,8 @@ export interface AppBindings { ENVIRONMENT: string; CF_API_TOKEN: string; CF_ACCOUNT_ID: string; + CF_PAID_WORKER?: string | boolean; + CF_SUBREQUEST_LIMIT?: string | number; } export interface AppVariables { diff --git a/docs/contributing.mdx b/docs/contributing.mdx index b0962c1d..97ac08ba 100644 --- a/docs/contributing.mdx +++ b/docs/contributing.mdx @@ -28,10 +28,10 @@ Contributions are welcome. The recommended workflow keeps changes focused and ea 3. **Install dependencies** — `npm install` and copy `.dev.vars.example` to `.dev.vars` with your secrets filled in. 4. **Make your change** — Keep the scope tight. Prefer existing patterns over new abstractions. 5. **Run typecheck and tests** — `npm run typecheck` and `npm test` must both pass before opening a PR. -6. **Open a pull request** — Open your PR against the `dev` branch. Write a clear description of what changed and why, and reference any related issues. +6. **Open a pull request** — Open your PR against `main`. Write a clear description of what changed and why, and reference any related issues. -Before opening a PR against dev, run npm run typecheck and npm test. Both must pass. PRs that break type coverage or existing tests will not be merged. +Pull requests are reviewed and merged directly into main; there is no separate dev integration branch. Run npm run typecheck and npm test before opening a PR, and ensure the required CI checks pass.

Engineering Conventions

@@ -51,7 +51,7 @@ Before opening a PR against dev, run npm run typecheck All documentation content for [codra.run/docs](https://codra.run/docs) is stored in the public `docs/` folder in the Codra repository. -Even though the frontend hosting is managed separately in a private repo, documentation content is 100% open source. When changes are merged into the `dev` branch, an automated GitHub Actions workflow copies the updated MDX files and deploys them to the website. +Even though the frontend hosting is managed separately in a private repo, documentation content is 100% open source. When changes are merged into `main`, an automated GitHub Actions workflow copies the updated MDX files and deploys them to the website. ### Guidelines for Docs * **Files & Routing**: Stored in `docs/.mdx` mapping directly to `https://codra.run/docs/`. diff --git a/package.json b/package.json index 039e90b5..e0cf0996 100644 --- a/package.json +++ b/package.json @@ -36,6 +36,8 @@ "setup:cloudflare": "node scripts/setup-cloudflare.js", "migrate": "node packages/db/scripts/migrate.mjs", "test": "node scripts/test.mjs", + "test:unit": "node scripts/test.mjs unit", + "test:integration": "node scripts/test.mjs integration", "test:all": "npm run test --workspaces --if-present", "test:watch": "vitest", "typecheck": "tsc --noEmit", diff --git a/packages/api/src/create-deps.ts b/packages/api/src/create-deps.ts index 70c220de..5444956d 100644 --- a/packages/api/src/create-deps.ts +++ b/packages/api/src/create-deps.ts @@ -36,9 +36,17 @@ export interface PlatformDeps { llmConfigEncryptionKey: string; cfApiToken?: string; cfAccountId?: string; + checkQuota?: ((input: QuotaCheckInput) => Promise) | null; + authz?: AuthzPort; } -import type { ApiRouterDeps, PlatformPort } from "./ports"; +import type { + ApiRouterDeps, + PlatformPort, + QuotaCheckInput, + QuotaResult, + AuthzPort, +} from "./ports"; import type { SessionStore, ReviewRuntime, @@ -81,6 +89,7 @@ import { createOAuthState, consumeOAuthState } from "./platform/oauth"; import { verifyGitHubWebhookSignature } from "@codraoss/core/verify"; import { createModelRunnerDeps } from "./model-runner-deps"; +import { createDefaultQuotaChecker } from "./platform/quota"; function toAppBindingsConfig(p: PlatformDeps) { return { @@ -184,5 +193,12 @@ export function createSharedApiDeps(p: PlatformDeps): ApiRouterDeps { normalizeGitHubWebhook(eventName, payload), extractReviewRequest: (input: any) => extractReviewRequest(input), }, + checkQuota: + p.checkQuota !== undefined + ? (p.checkQuota ?? undefined) + : p.kv + ? createDefaultQuotaChecker(p.kv) + : undefined, + authz: p.authz, }; } diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index d28962cb..2755c61c 100644 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -16,5 +16,6 @@ export type { export { requirePermission, requireQuota } from "./middleware/authorize"; export { requireSession } from "./middleware/auth"; export { requireCsrfHeader } from "./middleware/csrf"; +export { createDefaultQuotaChecker } from "./platform/quota"; export * from "./create-deps"; diff --git a/packages/api/src/platform/index.ts b/packages/api/src/platform/index.ts index 241be71e..94817d7b 100644 --- a/packages/api/src/platform/index.ts +++ b/packages/api/src/platform/index.ts @@ -2,3 +2,4 @@ export * from "./kv"; export * from "./config"; export * from "./oauth"; export * from "./updates-email"; +export * from "./quota"; diff --git a/packages/api/src/platform/quota.ts b/packages/api/src/platform/quota.ts new file mode 100644 index 00000000..8bfeeae5 --- /dev/null +++ b/packages/api/src/platform/quota.ts @@ -0,0 +1,65 @@ +import type { QuotaCheckInput, QuotaResult } from "../ports"; +import type { KvCompat } from "./kv"; + +export interface DefaultQuotaOptions { + windowSeconds?: number; + maxReviewsPerWindow?: number; + maxRetriesPerWindow?: number; +} + +export const DEFAULT_QUOTA_WINDOW_SECONDS = 600; // 10 minutes +export const DEFAULT_MAX_REVIEWS_PER_WINDOW = 30; // 30 review requests per 10 minutes per repo/installation +export const DEFAULT_MAX_RETRIES_PER_WINDOW = 20; // 20 retries per 10 minutes per user/action + +export function createDefaultQuotaChecker( + kv: KvCompat, + options: DefaultQuotaOptions = {}, +): (input: QuotaCheckInput) => Promise { + const windowSeconds = options.windowSeconds ?? DEFAULT_QUOTA_WINDOW_SECONDS; + const maxReviews = options.maxReviewsPerWindow ?? DEFAULT_MAX_REVIEWS_PER_WINDOW; + const maxRetries = options.maxRetriesPerWindow ?? DEFAULT_MAX_RETRIES_PER_WINDOW; + + return async (input: QuotaCheckInput): Promise => { + let limit = maxReviews; + if (input.action === "jobs.retry" || input.action === "jobs.rerun") { + limit = maxRetries; + } + + const subjectKey = + input.subject?.installationId + ? `install:${input.subject.installationId}` + : input.subject?.owner && input.subject?.repo + ? `repo:${input.subject.owner}/${input.subject.repo}` + : input.user?.login + ? `user:${input.user.login}` + : "global"; + + const currentTimestamp = Math.floor(Date.now() / 1000); + const windowSlot = Math.floor(currentTimestamp / windowSeconds); + const key = `quota:${input.action}:${subjectKey}:${windowSlot}`; + + try { + const raw = await kv.get(key, "text"); + const count = raw ? parseInt(raw, 10) : 0; + if (Number.isFinite(count) && count >= limit) { + const retryAfterSeconds = Math.max( + 1, + (windowSlot + 1) * windowSeconds - currentTimestamp, + ); + return { + allowed: false, + retryAfterSeconds, + reason: `Quota exceeded for ${input.action}. Limit is ${limit} per ${windowSeconds}s.`, + }; + } + + await kv.put(key, String(count + 1), { + expirationTtl: windowSeconds * 2, + }); + return { allowed: true }; + } catch { + // In case of KV lookup error, fail open to avoid blocking operations + return { allowed: true }; + } + }; +} diff --git a/packages/core/src/constants.ts b/packages/core/src/constants.ts index 32ac1385..fefcded9 100644 --- a/packages/core/src/constants.ts +++ b/packages/core/src/constants.ts @@ -31,6 +31,9 @@ export const MAX_FINALIZE_CONTINUATIONS = 3; export const MAX_ATTEMPTS_PER_WORKFLOW_INSTANCE = 6; // Budget & Attempts +export const DEFAULT_MAX_SUBREQUESTS = 50; +export const DEFAULT_SUBREQUEST_SAFE_MARGIN = 25; +export const PAID_WORKER_MAX_SUBREQUESTS = 1000; export const FILE_FIXED_SUBREQUESTS = 2; export const MAX_MODEL_ATTEMPTS_ESTIMATE = 4; export const MISSING_FILE_ERROR = diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 93686371..d54bf242 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -26,5 +26,17 @@ export { type VerifyOutcome, } from "./review"; +export { + TokenTracker, + type TokenTrackerOptions, +} from "./token-tracker"; + +export { + DEFAULT_MAX_SUBREQUESTS, + DEFAULT_SUBREQUEST_SAFE_MARGIN, + PAID_WORKER_MAX_SUBREQUESTS, +} from "./constants"; + export * from "./ports"; export * from "./review/model-gates"; + diff --git a/packages/core/src/token-tracker.ts b/packages/core/src/token-tracker.ts index 717165bd..58504002 100644 --- a/packages/core/src/token-tracker.ts +++ b/packages/core/src/token-tracker.ts @@ -1,4 +1,8 @@ import { logger } from "./logger"; +import { + DEFAULT_MAX_SUBREQUESTS, + DEFAULT_SUBREQUEST_SAFE_MARGIN, +} from "./constants"; export interface TokenUsage { input: number; @@ -19,13 +23,23 @@ export interface WastedUsage { byReason: Record; } +export interface TokenTrackerOptions { + maxSubrequests?: number; + safeMargin?: number; +} + export class TokenTracker { private usage: Map = new Map(); private wasted = { attempts: 0, estimatedInput: 0, skips: 0 }; private wastedByReason: Map = new Map(); private subrequests = 0; - private readonly MAX_SUBREQUESTS = 50; - private readonly SAFE_MARGIN = 25; + private readonly MAX_SUBREQUESTS: number; + private readonly SAFE_MARGIN: number; + + constructor(options?: TokenTrackerOptions) { + this.MAX_SUBREQUESTS = options?.maxSubrequests ?? DEFAULT_MAX_SUBREQUESTS; + this.SAFE_MARGIN = options?.safeMargin ?? DEFAULT_SUBREQUEST_SAFE_MARGIN; + } incrementSubrequests(count = 1) { this.subrequests += count; @@ -40,10 +54,16 @@ export class TokenTracker { } isNearLimit() { + if (!Number.isFinite(this.MAX_SUBREQUESTS)) { + return false; + } return this.subrequests >= this.MAX_SUBREQUESTS - this.SAFE_MARGIN; } remainingSafeBudget() { + if (!Number.isFinite(this.MAX_SUBREQUESTS)) { + return Number.POSITIVE_INFINITY; + } return Math.max( 0, this.MAX_SUBREQUESTS - this.SAFE_MARGIN - this.subrequests, diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index c3d106b2..c87bd001 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -2,14 +2,15 @@ import type { DbEnv } from "./env"; import { AsyncLocalStorage } from "node:async_hooks"; import postgres from "postgres"; -type DbClient = { +export type DbClient = { query(sqlText: string, params?: unknown[]): Promise; transaction(fn: (tx: DbClient) => Promise): Promise; + close?: () => Promise; }; const dbStorage = new AsyncLocalStorage(); -function createDbClient(env: DbEnv): DbClient { +export function createDbClient(env: DbEnv): DbClient { const sql = postgres(env.HYPERDRIVE.connectionString, { max: process.env.DB_POOL_SIZE ? parseInt(process.env.DB_POOL_SIZE, 10) : 5, fetch_types: false, @@ -49,6 +50,14 @@ function createDbClient(env: DbEnv): DbClient { return await fn(txClient); })) as T; }, + async close() { + fallbackClients.delete(env.HYPERDRIVE.connectionString); + try { + await sql.end({ timeout: 5 }); + } catch { + // ignore errors on close + } + }, }; } @@ -71,8 +80,10 @@ function toPostgresArrayElement(value: unknown) { return `"${text}"`; } -export function runWithDb(env: DbEnv, fn: () => T): T { - return dbStorage.run(createDbClient(env), fn); +export function runWithDb(env: DbEnv, fn: () => T, client?: DbClient): T { + const resolvedClient = + client ?? (env.workerMode === false ? getDb(env) : createDbClient(env)); + return dbStorage.run(resolvedClient, fn); } // Module-scoped Map pools connections outside runWithDb, but must self-heal when request context changes. @@ -100,6 +111,21 @@ export function getDb(env: DbEnv) { return client; } +export async function closeDb(env?: DbEnv): Promise { + if (env) { + const connectionString = env.HYPERDRIVE.connectionString; + const client = fallbackClients.get(connectionString); + if (client) { + fallbackClients.delete(connectionString); + await client.close?.(); + } + } else { + const clients = Array.from(fallbackClients.values()); + fallbackClients.clear(); + await Promise.all(clients.map((c) => c.close?.())); + } +} + // Retries op once on a fresh client if the module-cached socket has gone stale. async function withStaleConnectionRecovery( env: DbEnv, @@ -113,7 +139,9 @@ async function withStaleConnectionRecovery( throw error; const connectionString = env.HYPERDRIVE.connectionString; + const stale = fallbackClients.get(connectionString); fallbackClients.delete(connectionString); + stale?.close?.().catch?.(() => {}); const fresh = createDbClient(env); fallbackClients.set(connectionString, fresh); diff --git a/packages/node-adapters/src/node-orchestrator.ts b/packages/node-adapters/src/node-orchestrator.ts index 3581837d..7288842a 100644 --- a/packages/node-adapters/src/node-orchestrator.ts +++ b/packages/node-adapters/src/node-orchestrator.ts @@ -30,7 +30,7 @@ export class NodeOrchestrator implements JobOrchestrator { nextParams.jobId = result.jobId; } await this.queue.send(nextParams, { - delaySeconds: result.delaySeconds, + delaySeconds: 0, }); } else if (result.action === "retry") { await this.queue.send(currentParams, { diff --git a/packages/provider-github/src/index.ts b/packages/provider-github/src/index.ts index 8aa20e77..93a2b5cf 100644 --- a/packages/provider-github/src/index.ts +++ b/packages/provider-github/src/index.ts @@ -15,3 +15,4 @@ export { type GitHubOAuthProfile, } from "./oauth"; export { normalizeGitHubWebhook } from "./webhook"; +export { findBotReviewForCommit, postReview } from "./review-post"; diff --git a/packages/provider-github/src/review-post.ts b/packages/provider-github/src/review-post.ts index 426ef5c5..e4fc4be1 100644 --- a/packages/provider-github/src/review-post.ts +++ b/packages/provider-github/src/review-post.ts @@ -7,6 +7,53 @@ import { } from "./http"; import type { ReviewComment } from "./types"; +export function appendDroppedCommentsToSummary( + summaryBody: string, + droppedComments: ReviewComment[], +): string { + if (droppedComments.length === 0) return summaryBody; + const rows = droppedComments.map((comment) => { + const file = `\`${comment.path}\``; + const line = + typeof comment.line === "number" && comment.line > 0 + ? String(comment.line) + : typeof comment.position === "number" && comment.position > 0 + ? `pos ${comment.position}` + : "-"; + const body = comment.body + .replace(/\|/g, "\\|") + .replace(/\r?\n/g, "
"); + return `| ${file} | ${line} | ${body} |`; + }); + const appendix = [ + "### ⚠️ Comments on Out-of-Diff Lines", + "", + "The following comments could not be posted inline (lines outside the pull request diff):", + "", + "| File | Line | Finding |", + "| :--- | :--- | :--- |", + ...rows, + ].join("\n"); + const trimmed = summaryBody.trimEnd(); + return trimmed ? `${trimmed}\n\n${appendix}` : appendix; +} + +export function parseRejectedCommentIndices( + errorText: string, + commentCount: number, +): number[] { + try { + const rejected = new Set(); + for (const match of errorText.matchAll(/(?:comments?[.[]|"index":\s*)(\d+)/gi)) { + const idx = Number.parseInt(match[1], 10); + if (idx >= 0 && idx < commentCount) rejected.add(idx); + } + return Array.from(rejected); + } catch { + return []; + } +} + export async function postReview( ctx: GitHubRequestContext, owner: string, @@ -21,89 +68,119 @@ export async function postReview( ) { return withRetry(`createReview ${owner}/${repo}#${pullNumber}`, async () => { // Address by `line` + `side`, or legacy `position`. - const mapped = input.comments.map((comment) => { + const mapped = input.comments.map((comment, index) => { if (typeof comment.line === "number" && comment.line > 0) { return { - path: comment.path, - line: comment.line, - side: comment.side ?? "RIGHT", - body: comment.body, + originalIndex: index, + comment: { + path: comment.path, + line: comment.line, + side: comment.side ?? "RIGHT", + body: comment.body, + }, }; } if (typeof comment.position === "number" && comment.position > 0) { return { - path: comment.path, - position: comment.position, - body: comment.body, + originalIndex: index, + comment: { + path: comment.path, + position: comment.position, + body: comment.body, + }, }; } return null; }); - // Which of the caller's comments survived mapping; marking a dropped one "posted" would suppress it on every future commit unseen. - let postedIndices = mapped.flatMap((c, index) => - c === null ? [] : [index], + const addressable = mapped.filter( + (m): m is NonNullable => m !== null, ); + let postedIndices = addressable.map((a) => a.originalIndex); + const comments = addressable.map((a) => a.comment); - const comments = mapped.filter( - (c): c is NonNullable => c !== null, - ); - const unaddressable = mapped.length - comments.length; + const unaddressable = input.comments.length - addressable.length; if (unaddressable > 0) { logger.warn("Dropping review comments with no usable line/position", { owner, repo, pullNumber, unaddressable, - total: mapped.length, + total: input.comments.length, }); } - const body = { - commit_id: input.commitSha, - event: input.event, - body: input.body, - comments, - }; - const reviewPath = `${repoApiPath(owner, repo)}/pulls/${pullNumber}/reviews`; let response = await ctx.request(reviewPath, { method: "POST", - headers: { - "content-type": "application/json", - }, - body: JSON.stringify(body), + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + commit_id: input.commitSha, + event: input.event, + body: input.body, + comments, + }), }); - if (response.status === 422 && body.comments.length > 0) { - // A 422 here almost always means a comment pointed at a line outside the diff; log the response text or this failure is invisible. + if (response.status === 422 && comments.length > 0) { const reason = await response .clone() .text() .catch(() => ""); - logger.warn( - `GitHub review creation failed with 422, retrying without inline comments`, - { - owner, - repo, - pullNumber, - droppedComments: body.comments.length, - reason: reason.slice(0, 500), - }, - ); - response = await ctx.request(reviewPath, { - method: "POST", - headers: { - "content-type": "application/json", - }, - body: JSON.stringify({ - commit_id: input.commitSha, - event: input.event, - body: input.body, - comments: [], - }), - }); - postedIndices = []; // Summary posts, but no inline comments. + const rejected = new Set(parseRejectedCommentIndices(reason, comments.length)); + + // Attempt to salvage remaining valid comments if specific invalid ones were identified + if (rejected.size > 0 && rejected.size < addressable.length) { + const remaining = addressable.filter((_, i) => !rejected.has(i)); + const salvagedIndices = remaining.map((a) => a.originalIndex); + const dropped = input.comments.filter((_, i) => !salvagedIndices.includes(i)); + const body = appendDroppedCommentsToSummary(input.body, dropped); + + logger.info( + `GitHub review rejected ${rejected.size} comment(s); retrying with ${remaining.length} salvaged comment(s)`, + { owner, repo, pullNumber, rejectedCount: rejected.size, salvagedCount: remaining.length }, + ); + + response = await ctx.request(reviewPath, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + commit_id: input.commitSha, + event: input.event, + body, + comments: remaining.map((a) => a.comment), + }), + }); + + if (response.status !== 422) { + postedIndices = salvagedIndices; + } + } + + // If salvage was skipped or still returned 422, drop all inline comments and append to summary + if (response.status === 422) { + logger.warn( + `GitHub review creation failed with 422, retrying without inline comments`, + { + owner, + repo, + pullNumber, + droppedComments: input.comments.length, + reason: reason.slice(0, 500), + }, + ); + response = await ctx.request(reviewPath, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + commit_id: input.commitSha, + event: input.event, + body: appendDroppedCommentsToSummary(input.body, input.comments), + comments: [], + }), + }); + postedIndices = []; + } } await assertResponseOk(response, reviewPath, "GitHub review creation"); @@ -134,11 +211,12 @@ export async function findBotReviewForCommit( user?: { login?: string | null } | null; }>; const login = botLogin.toLowerCase(); - const match = reviews.find( - (review) => - review.commit_id === commitSha && - (review.user?.login ?? "").toLowerCase().startsWith(login), - ); + const botSuffix = `${login}[bot]`; + const match = reviews.find((review) => { + if (review.commit_id !== commitSha) return false; + const userLogin = (review.user?.login ?? "").toLowerCase(); + return userLogin === login || userLogin === botSuffix; + }); return match ? { id: match.id } : null; }, ); diff --git a/scripts/test.mjs b/scripts/test.mjs index 6a813132..c28c6009 100644 --- a/scripts/test.mjs +++ b/scripts/test.mjs @@ -72,17 +72,66 @@ function run(command, args) { loadEnvFiles(); -if (!usableEnvValue(process.env.TEST_DATABASE_URL)) { - console.error( - [ - "TEST_DATABASE_URL is required to run the full test suite.", - "Copy .env.test.example to .env.test and point TEST_DATABASE_URL at a disposable Postgres database.", - ].join("\n"), - ); - process.exit(1); -} +const rawArgs = process.argv.slice(2); +const isUnit = rawArgs.includes("unit") || rawArgs.includes("--unit"); +const isIntegration = + rawArgs.includes("integration") || rawArgs.includes("--integration"); +const passthroughArgs = rawArgs.filter( + (arg) => + arg !== "unit" && + arg !== "--unit" && + arg !== "integration" && + arg !== "--integration", +); + +if (isUnit) { + process.env.CODRA_TEST_SUITE = "unit"; + delete process.env.TEST_DATABASE_URL; + delete process.env.DATABASE_URL; + run(process.execPath, [ + "node_modules/vitest/vitest.mjs", + "run", + ...passthroughArgs, + ]); +} else if (isIntegration) { + process.env.CODRA_TEST_SUITE = "integration"; + if (!usableEnvValue(process.env.TEST_DATABASE_URL)) { + console.error( + [ + "TEST_DATABASE_URL is required to run the integration test suite.", + "Copy .env.test.example to .env.test and point TEST_DATABASE_URL at a disposable Postgres database.", + ].join("\n"), + ); + process.exit(1); + } -process.env.DATABASE_URL = process.env.TEST_DATABASE_URL; + process.env.DATABASE_URL = process.env.TEST_DATABASE_URL; -run(process.execPath, ["packages/db/scripts/migrate.mjs"]); -run(process.execPath, ["node_modules/vitest/vitest.mjs", "run"]); + run(process.execPath, ["packages/db/scripts/migrate.mjs"]); + run(process.execPath, [ + "node_modules/vitest/vitest.mjs", + "run", + ...passthroughArgs, + ]); +} else { + // Default: full test suite (backwards compatible) + if (!usableEnvValue(process.env.TEST_DATABASE_URL)) { + console.error( + [ + "TEST_DATABASE_URL is required to run the full test suite.", + "Copy .env.test.example to .env.test and point TEST_DATABASE_URL at a disposable Postgres database.", + "To run tests without a database, use: npm run test:unit", + ].join("\n"), + ); + process.exit(1); + } + + process.env.DATABASE_URL = process.env.TEST_DATABASE_URL; + + run(process.execPath, ["packages/db/scripts/migrate.mjs"]); + run(process.execPath, [ + "node_modules/vitest/vitest.mjs", + "run", + ...passthroughArgs, + ]); +} diff --git a/test/api/default-quota.spec.ts b/test/api/default-quota.spec.ts new file mode 100644 index 00000000..95108956 --- /dev/null +++ b/test/api/default-quota.spec.ts @@ -0,0 +1,161 @@ +import { describe, it, expect } from "vitest"; +import { + createDefaultQuotaChecker, + createSharedApiDeps, + createApiRouter, +} from "@codraoss/api"; +import { findBotReviewForCommit } from "@codraoss/provider-github"; +import type { KvCompat } from "@codraoss/api/platform"; + +class MockMemoryKV implements KvCompat { + private store = new Map(); + + async put( + key: string, + value: string, + _options?: { expirationTtl?: number }, + ): Promise { + this.store.set(key, value); + } + + async get(key: string, type?: "json" | "text"): Promise { + const val = this.store.get(key); + if (val === undefined) return null; + if (type === "json") return JSON.parse(val); + return val; + } + + async delete(key: string): Promise { + this.store.delete(key); + } +} + +describe("Default Quota Checker & Rate Limiting", () => { + it("allows requests under the review limit and enforces quota when limit is reached", async () => { + const kv = new MockMemoryKV(); + const checker = createDefaultQuotaChecker(kv, { + maxReviewsPerWindow: 3, + windowSeconds: 60, + }); + + const input = { + action: "reviews.enqueue" as const, + subject: { installationId: "inst-1", owner: "org-1", repo: "repo-1" }, + }; + + // First 3 requests should be allowed + expect((await checker(input)).allowed).toBe(true); + expect((await checker(input)).allowed).toBe(true); + expect((await checker(input)).allowed).toBe(true); + + // 4th request must be rejected + const blocked = await checker(input); + expect(blocked.allowed).toBe(false); + if (!blocked.allowed) { + expect(blocked.reason).toContain("Quota exceeded"); + expect(blocked.retryAfterSeconds).toBeGreaterThan(0); + } + }); + + it("isolates counters by subject (different installations/repos have independent quotas)", async () => { + const kv = new MockMemoryKV(); + const checker = createDefaultQuotaChecker(kv, { + maxReviewsPerWindow: 2, + windowSeconds: 60, + }); + + const repo1 = { + action: "reviews.enqueue" as const, + subject: { installationId: "inst-1", owner: "org-1", repo: "repo-1" }, + }; + const repo2 = { + action: "reviews.enqueue" as const, + subject: { installationId: "inst-2", owner: "org-2", repo: "repo-2" }, + }; + + expect((await checker(repo1)).allowed).toBe(true); + expect((await checker(repo1)).allowed).toBe(true); + expect((await checker(repo1)).allowed).toBe(false); + + // repo2 should still be allowed + expect((await checker(repo2)).allowed).toBe(true); + expect((await checker(repo2)).allowed).toBe(true); + expect((await checker(repo2)).allowed).toBe(false); + }); + + it("applies maxRetriesPerWindow for dashboard retry/rerun actions", async () => { + const kv = new MockMemoryKV(); + const checker = createDefaultQuotaChecker(kv, { + maxRetriesPerWindow: 2, + windowSeconds: 60, + }); + + const retryInput = { + action: "jobs.retry" as const, + user: { login: "developer" } as any, + }; + + expect((await checker(retryInput)).allowed).toBe(true); + expect((await checker(retryInput)).allowed).toBe(true); + const blocked = await checker(retryInput); + expect(blocked.allowed).toBe(false); + }); + + it("fails open if KV access throws", async () => { + const failingKv: KvCompat = { + async put() { + throw new Error("KV offline"); + }, + async get() { + throw new Error("KV offline"); + }, + async delete() {}, + }; + + const checker = createDefaultQuotaChecker(failingKv); + const res = await checker({ + action: "reviews.enqueue" as const, + subject: { installationId: "123" }, + }); + expect(res.allowed).toBe(true); + }); +}); + +describe("Bot Review Exact Matching", () => { + function makeMockCtx(reviews: Array<{ commit_id: string; user: { login: string } }>): any { + return { + async request() { + return new Response(JSON.stringify(reviews), { status: 200 }); + }, + async requestAndCheck() { + return new Response(JSON.stringify(reviews), { status: 200 }); + }, + }; + } + + it("matches exact bot username", async () => { + const ctx = makeMockCtx([ + { commit_id: "sha1", user: { login: "codra" } }, + ]); + const found = await findBotReviewForCommit(ctx, "owner", "repo", 1, "sha1", "codra"); + expect(found).not.toBeNull(); + }); + + it("matches exact bot username with [bot] suffix", async () => { + const ctx = makeMockCtx([ + { commit_id: "sha1", user: { login: "codra[bot]" } }, + ]); + const found = await findBotReviewForCommit(ctx, "owner", "repo", 1, "sha1", "codra"); + expect(found).not.toBeNull(); + }); + + it("rejects spoofed username that merely starts with the bot username", async () => { + const ctx = makeMockCtx([ + { commit_id: "sha1", user: { login: "codra-attacker" } }, + { commit_id: "sha1", user: { login: "codrabot" } }, + { commit_id: "sha1", user: { login: "codra123" } }, + ]); + const found = await findBotReviewForCommit(ctx, "owner", "repo", 1, "sha1", "codra"); + expect(found).toBeNull(); + }); +}); diff --git a/test/db/connection-pool-reuse.spec.ts b/test/db/connection-pool-reuse.spec.ts new file mode 100644 index 00000000..4c69802d --- /dev/null +++ b/test/db/connection-pool-reuse.spec.ts @@ -0,0 +1,78 @@ +import { describe, it, expect } from "vitest"; +import { runWithDb, getDb, type DbClient } from "@codraoss/db/client"; +import type { DbEnv } from "@codraoss/db/env"; + +describe("Node.js connection pool reuse and client lifecycle", () => { + const dummyEnv: DbEnv = { + HYPERDRIVE: { connectionString: "postgres://postgres:password@localhost:5432/codra_test" }, + APP_KV: { + put: async () => {}, + get: async () => null, + delete: async () => {}, + }, + workerMode: false, + }; + + it("reuses process-level singleton DbClient across runWithDb calls when workerMode is false", () => { + let client1: DbClient | undefined; + let client2: DbClient | undefined; + + runWithDb(dummyEnv, () => { + client1 = getDb(dummyEnv); + }); + + runWithDb(dummyEnv, () => { + client2 = getDb(dummyEnv); + }); + + expect(client1).toBeDefined(); + expect(client2).toBeDefined(); + expect(client1).toBe(client2); + }); + + it("reuses explicit singleton DbClient passed into runWithDb", () => { + const explicitClient = getDb(dummyEnv); + let capturedClient: DbClient | undefined; + + runWithDb(dummyEnv, () => { + capturedClient = getDb(dummyEnv); + }, explicitClient); + + expect(capturedClient).toBe(explicitClient); + }); + + it("preserves Cloudflare Workers semantics by creating separate instances when workerMode is not false", () => { + const workerEnv: DbEnv = { + HYPERDRIVE: { connectionString: "postgres://postgres:password@localhost:5432/codra_worker_test" }, + APP_KV: dummyEnv.APP_KV, + workerMode: true, + }; + + let clientA: DbClient | undefined; + let clientB: DbClient | undefined; + + runWithDb(workerEnv, () => { + clientA = getDb(workerEnv); + }); + + runWithDb(workerEnv, () => { + clientB = getDb(workerEnv); + }); + + expect(clientA).toBeDefined(); + expect(clientB).toBeDefined(); + expect(clientA).not.toBe(clientB); + }); + + it("allows graceful close of DbClient", async () => { + const envToClose: DbEnv = { + HYPERDRIVE: { connectionString: "postgres://postgres:password@localhost:5432/codra_close_test" }, + APP_KV: dummyEnv.APP_KV, + workerMode: false, + }; + + const client = getDb(envToClose); + expect(typeof client.close).toBe("function"); + await expect(client.close?.()).resolves.toBeUndefined(); + }); +}); diff --git a/test/helpers.ts b/test/helpers.ts index 08ce8a8e..c17c7c7a 100644 --- a/test/helpers.ts +++ b/test/helpers.ts @@ -1,4 +1,4 @@ -import { describe } from "vitest"; +import { describe } from "vitest"; import type { AppBindings } from "@server/env"; import { InMemorySessionStore } from "@codraoss/core"; import { encryptLlmApiKey, ModelRunner } from "@codraoss/models"; @@ -144,7 +144,9 @@ export function createTestEnv( REVIEW_WORKFLOW: new MockWorkflow() as any, ASSETS: new MockAssets() as any, HYPERDRIVE: { - connectionString: getTestDatabaseUrl(), + connectionString: hasConfiguredTestDatabaseUrl() + ? getTestDatabaseUrl() + : "postgresql://mock:mock@127.0.0.1:5432/codra_mock", }, APP_PRIVATE_KEY: "private_key", GITHUB_APP_ID: "app_id", diff --git a/test/review/comments.spec.ts b/test/review/comments.spec.ts index 8a6d7e2a..747c7d96 100644 --- a/test/review/comments.spec.ts +++ b/test/review/comments.spec.ts @@ -137,12 +137,75 @@ describe("createReview inline comment payload", () => { expect(review.id).toBe(777); expect(sent[0].comments).toHaveLength(1); expect(sent[1].comments).toHaveLength(0); - expect(sent[1].body).toBe("summary"); + expect(sent[1].body).toContain("summary"); + expect(sent[1].body).toContain("### ⚠️ Comments on Out-of-Diff Lines"); + expect(sent[1].body).toContain("| `src/app.ts` | 12 | Something is wrong here. |"); // Nothing was actually shown, so nothing may be recorded as posted -- otherwise the finding // would be suppressed on every later commit without a human ever having seen it. expect(review.postedIndices).toEqual([]); }); + it("salvages valid inline comments when GitHub 422 identifies the rejected comment", async () => { + const client = new GitHubClient({} as never, "123"); + const sent: any[] = []; + let call = 0; + vi.spyOn(client as any, "request").mockImplementation( + async (...args: unknown[]) => { + const [, init] = args as [string, RequestInit]; + sent.push(JSON.parse(String(init.body))); + call += 1; + // First attempt (with 2 comments) rejects the second comment as out-of-diff + if (call === 1) { + return new Response( + JSON.stringify({ + message: "Validation Failed", + errors: [ + { + resource: "PullRequestReviewComment", + field: "comments[1].line", + code: "custom", + message: "line must be part of the diff", + }, + ], + }), + { status: 422 }, + ); + } + return new Response(JSON.stringify({ id: 888 }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }, + ); + + const review = await client.createReview("o", "r", 7, { + commitSha: "abc123", + event: "COMMENT", + body: "summary", + comments: [ + comment({ path: "src/good.ts", line: 10, body: "Valid comment" }), + comment({ path: "src/bad.ts", line: 99, body: "Out of diff comment" }), + ], + }); + + expect(review.id).toBe(888); + expect(sent).toHaveLength(2); + // First attempt sent both comments + expect(sent[0].comments).toHaveLength(2); + // Salvage retry sent only the valid comment inline + expect(sent[1].comments).toHaveLength(1); + expect(sent[1].comments[0]).toMatchObject({ + path: "src/good.ts", + line: 10, + }); + // Dropped comment finding was preserved in the summary body + expect(sent[1].body).toContain("summary"); + expect(sent[1].body).toContain("### ⚠️ Comments on Out-of-Diff Lines"); + expect(sent[1].body).toContain("| `src/bad.ts` | 99 | Out of diff comment |"); + // Only the salvaged comment was posted inline + expect(review.postedIndices).toEqual([0]); + }); + it("reports which comments GitHub accepted, by caller index", async () => { const client = new GitHubClient({} as never, "123"); vi.spyOn(client as any, "request").mockResolvedValue( diff --git a/test/review/subrequest-limits.spec.ts b/test/review/subrequest-limits.spec.ts new file mode 100644 index 00000000..cc6652a2 --- /dev/null +++ b/test/review/subrequest-limits.spec.ts @@ -0,0 +1,75 @@ +import { describe, it, expect } from "vitest"; +import { + budgetAwareFileLimit, + DEFAULT_MAX_SUBREQUESTS, + DEFAULT_SUBREQUEST_SAFE_MARGIN, + PAID_WORKER_MAX_SUBREQUESTS, + TokenTracker, +} from "@codraoss/core"; +import { resolveWorkerSubrequestLimits } from "../../apps/worker/src/adapters"; +import { resolveNodeSubrequestLimits } from "../../apps/node/src/runtime"; + +describe("Subrequest Limits & Decoupling", () => { + describe("resolveWorkerSubrequestLimits", () => { + it("preserves Cloudflare free tier default (50 max, 25 safe margin) when no override provided", () => { + const limits = resolveWorkerSubrequestLimits({}); + expect(limits).toEqual({ + maxSubrequests: DEFAULT_MAX_SUBREQUESTS, + safeMargin: DEFAULT_SUBREQUEST_SAFE_MARGIN, + }); + expect(limits.maxSubrequests).toBe(50); + expect(limits.safeMargin).toBe(25); + }); + + it("uses paid worker limit (1000) when CF_PAID_WORKER is set to true / 'true' / '1'", () => { + expect(resolveWorkerSubrequestLimits({ CF_PAID_WORKER: true })).toEqual({ + maxSubrequests: PAID_WORKER_MAX_SUBREQUESTS, + }); + expect(resolveWorkerSubrequestLimits({ CF_PAID_WORKER: "true" })).toEqual({ + maxSubrequests: 1000, + }); + expect(resolveWorkerSubrequestLimits({ CF_PAID_WORKER: "1" })).toEqual({ + maxSubrequests: 1000, + }); + }); + + it("respects explicit CF_SUBREQUEST_LIMIT override", () => { + expect(resolveWorkerSubrequestLimits({ CF_SUBREQUEST_LIMIT: 250 })).toEqual({ + maxSubrequests: 250, + }); + expect(resolveWorkerSubrequestLimits({ CF_SUBREQUEST_LIMIT: "500" })).toEqual({ + maxSubrequests: 500, + }); + }); + }); + + describe("resolveNodeSubrequestLimits", () => { + it("defaults to unconstrained (Infinity) subrequests in Node.js / Docker", () => { + const limits = resolveNodeSubrequestLimits({}); + expect(limits.maxSubrequests).toBe(Number.POSITIVE_INFINITY); + }); + + it("allows optional SUBREQUEST_LIMIT override in Node.js", () => { + const limits = resolveNodeSubrequestLimits({ SUBREQUEST_LIMIT: 300 }); + expect(limits.maxSubrequests).toBe(300); + }); + + it("allows optional CF_SUBREQUEST_LIMIT override in Node.js", () => { + const limits = resolveNodeSubrequestLimits({ CF_SUBREQUEST_LIMIT: "400" }); + expect(limits.maxSubrequests).toBe(400); + }); + }); + + + describe("budgetAwareFileLimit with unconstrained budget", () => { + it("does not constrain file limits when budget is Infinity", () => { + const tracker = new TokenTracker({ maxSubrequests: Number.POSITIVE_INFINITY }); + const safeBudget = tracker.remainingSafeBudget(); + expect(safeBudget).toBe(Number.POSITIVE_INFINITY); + + // Concurrency limits are fully respected without being squeezed down by subrequests + expect(budgetAwareFileLimit(safeBudget, 4, 3, true, true)).toBe(4); + expect(budgetAwareFileLimit(safeBudget, 10, 1, false, false)).toBe(10); + }); + }); +}); diff --git a/test/setup.ts b/test/setup.ts index 107f4999..edaff04a 100644 --- a/test/setup.ts +++ b/test/setup.ts @@ -86,7 +86,12 @@ function loadTestEnvFromFiles() { } function assertRequiredTestEnv() { - const missing = REQUIRED_TEST_ENV_KEYS.filter( + const isUnit = process.env.CODRA_TEST_SUITE === "unit"; + const requiredKeys = isUnit + ? REQUIRED_TEST_ENV_KEYS.filter((key) => key !== "TEST_DATABASE_URL") + : REQUIRED_TEST_ENV_KEYS; + + const missing = requiredKeys.filter( (key) => !usableEnvValue(process.env[key]), ); if (missing.length === 0) return; @@ -95,7 +100,11 @@ function assertRequiredTestEnv() { [ `Missing required test environment variables: ${missing.join(", ")}.`, "Set these values in .env.test, .env.local, .env, .dev.vars, .env.test.example, or CI.", - "TEST_DATABASE_URL must point to a disposable Postgres database so the full test suite can run.", + ...(isUnit + ? [] + : [ + "TEST_DATABASE_URL must point to a disposable Postgres database so the full test suite can run.", + ]), ].join("\n"), ); } diff --git a/test/token-tracker.spec.ts b/test/token-tracker.spec.ts index 4893273f..71211186 100644 --- a/test/token-tracker.spec.ts +++ b/test/token-tracker.spec.ts @@ -37,6 +37,32 @@ describe("TokenTracker.remainingSafeBudget", () => { expect(tracker.isNearLimit()).toBe(true); expect(tracker.remainingSafeBudget()).toBe(0); }); + + it("handles unconstrained subrequests without artificial limits", () => { + const tracker = new TokenTracker({ maxSubrequests: Number.POSITIVE_INFINITY }); + expect(tracker.remainingSafeBudget()).toBe(Number.POSITIVE_INFINITY); + expect(tracker.isNearLimit()).toBe(false); + expect(tracker.hasRemainingSubrequests(100_000)).toBe(true); + + tracker.incrementSubrequests(500); + expect(tracker.remainingSafeBudget()).toBe(Number.POSITIVE_INFINITY); + expect(tracker.isNearLimit()).toBe(false); + expect(tracker.hasRemainingSubrequests(100_000)).toBe(true); + }); + + it("supports configurable paid worker limits", () => { + const tracker = new TokenTracker({ maxSubrequests: 1000, safeMargin: 25 }); + expect(tracker.remainingSafeBudget()).toBe(975); // 1000 - 25 + expect(tracker.isNearLimit()).toBe(false); + + tracker.incrementSubrequests(974); + expect(tracker.isNearLimit()).toBe(false); + expect(tracker.remainingSafeBudget()).toBe(1); + + tracker.incrementSubrequests(1); + expect(tracker.isNearLimit()).toBe(true); + expect(tracker.remainingSafeBudget()).toBe(0); + }); }); // record() only ran on success, so retried/429'd sends were invisible to accounting. diff --git a/vitest.config.ts b/vitest.config.ts index b048c9db..3ecf62e5 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -2,6 +2,40 @@ import { defineConfig } from "vitest/config"; import react from "@vitejs/plugin-react"; import { resolve } from "path"; +export const INTEGRATION_TEST_PATTERNS = [ + "test/db/**/*.spec.ts", + "test/jsonb-encoding.spec.ts", + "test/findings/suppression.spec.ts", + "test/api/auth.spec.ts", + "test/api/authorize.spec.ts", + "test/api/jobs.spec.ts", + "test/api/quota.spec.ts", + "test/api/repos.spec.ts", + "test/api/router-options.spec.ts", + "test/api/webhook-handling.spec.ts", + "test/review/async-batch.spec.ts", + "test/review/batch-flow.spec.ts", + "test/review/flow-chunking.spec.ts", + "test/review/flow-lifecycle.spec.ts", + "test/review/flow-retry.spec.ts", + "test/review/quota-deferral.spec.ts", + "test/review/resumable-queue.spec.ts", + "test/comment-feedback.spec.ts", + "test/**/*.integration.spec.ts", +]; + +const testSuite = process.env.CODRA_TEST_SUITE; + +const testInclude = + testSuite === "integration" + ? INTEGRATION_TEST_PATTERNS + : ["test/**/*.spec.ts", "test/**/*.spec.tsx"]; + +const testExclude = + testSuite === "unit" + ? ["**/node_modules/**", "**/dist/**", ...INTEGRATION_TEST_PATTERNS] + : ["**/node_modules/**", "**/dist/**"]; + export default defineConfig({ plugins: [react()], resolve: { @@ -18,7 +52,8 @@ export default defineConfig({ test: { globals: true, environment: "node", - include: ["test/**/*.spec.ts", "test/**/*.spec.tsx"], + include: testInclude, + exclude: testExclude, passWithNoTests: true, setupFiles: [resolve(__dirname, "./test/setup.ts")], // The suite is dominated by round trips to a remote Postgres, so wall clock is latency-bound,