diff --git a/packages/corsair/core/endpoints/bind.ts b/packages/corsair/core/endpoints/bind.ts index ca104dfad..605076ce3 100644 --- a/packages/corsair/core/endpoints/bind.ts +++ b/packages/corsair/core/endpoints/bind.ts @@ -81,6 +81,7 @@ export function bindEndpointsRecursively({ permissionsConfig?: { mode: PermissionMode; overrides?: Record; + limits?: import('../plugins').UsageLimit[]; }; /** Risk level metadata per dot-notation endpoint path. Defaults riskLevel to 'write' when missing. */ endpointMeta?: Record; @@ -98,6 +99,15 @@ export function bindEndpointsRecursively({ allPlugins?: readonly CorsairPlugin[]; multiTenancy?: boolean; }): void { + if ( + (permissionsOptions?.limits?.length || permissionsConfig?.limits?.length) && + !database + ) { + throw new Error( + 'Cannot enable usage limits without a database connection. Please configure a database or remove the limits configuration.', + ); + } + for (const [key, value] of Object.entries(endpoints)) { // we have to retype this now because it's nested webhooks const nodeHooks = hooks?.[key] as Record | undefined; @@ -122,7 +132,7 @@ export function bindEndpointsRecursively({ // ── Permission guard ──────────────────────────────────────────────────────────────── let onPermissionComplete: (() => Promise) | undefined; - if (permissionsConfig) { + if (permissionsConfig || permissionsOptions?.limits?.length) { const meta = endpointMetaEntry; const { result: permResult, @@ -135,8 +145,8 @@ export function bindEndpointsRecursively({ pluginId, endpointPath: operationPath, args, - mode: permissionsConfig.mode, - override: permissionsConfig.overrides?.[operationPath], + mode: permissionsConfig?.mode ?? 'open', + override: permissionsConfig?.overrides?.[operationPath], // Default to 'write' when no meta declared — conservative fallback riskLevel: meta?.riskLevel ?? 'write', meta, @@ -146,6 +156,8 @@ export function bindEndpointsRecursively({ : undefined, tenantId, approvalMode: permissionsOptions?.mode, + globalLimits: permissionsOptions?.limits, + pluginLimits: permissionsConfig?.limits, }); if (permResult === 'blocked') { let msg: string; @@ -153,6 +165,10 @@ export function bindEndpointsRecursively({ msg = `Action '${operationPath}' was denied by the user. Await further instructions before proceeding.`; } else if (permReason === 'policy') { msg = `Action '${operationPath}' is blocked by the permission policy. Update the corsair config to allow it.`; + } else if (permReason === 'rate_limit_exceeded') { + msg = `Action '${operationPath}' is currently rate limited. Please wait before trying again.`; + } else if (permReason === 'budget_exhausted') { + msg = `Action '${operationPath}' has exhausted its configured budget quota.`; } else if (permReason === 'timeout') { msg = `Action '${operationPath}' timed out waiting for approval.`; } else if (permToken && permId) { diff --git a/packages/corsair/core/permissions/index.ts b/packages/corsair/core/permissions/index.ts index faffa2a4a..17c1c1b60 100644 --- a/packages/corsair/core/permissions/index.ts +++ b/packages/corsair/core/permissions/index.ts @@ -8,6 +8,7 @@ import type { EndpointRiskLevel, PermissionMode, PermissionPolicy, + UsageLimit, } from '../plugins'; // ───────────────────────────────────────────────────────────────────────────── @@ -236,12 +237,20 @@ export type EnforcePermissionOptions = { | 'synchronous' | 'asynchronous' | (() => 'synchronous' | 'asynchronous'); + globalLimits?: (UsageLimit & { scope?: 'global' | 'tenant' })[]; + pluginLimits?: UsageLimit[]; }; export type EnforcePermissionResult = { result: 'allow' | 'blocked'; /** Why the call was blocked. Only present when result === 'blocked'. */ - reason?: 'denied' | 'policy' | 'timeout' | 'pending'; + reason?: + | 'denied' + | 'policy' + | 'timeout' + | 'pending' + | 'rate_limit_exceeded' + | 'budget_exhausted'; /** Permission record ID. Present when a pending approval record exists. */ id?: string; /** Permission token (the value embedded in review URLs). Present when a pending approval record exists. */ @@ -322,7 +331,6 @@ export async function enforcePermission( opts: EnforcePermissionOptions, ): Promise { const policy = evaluatePermission(opts.riskLevel, opts.mode, opts.override); - if (policy === 'allow') return { result: 'allow' }; const irreversibleNote = opts.meta?.irreversible ? ' (irreversible)' : ''; const description = opts.meta?.description @@ -338,9 +346,10 @@ export async function enforcePermission( return { result: 'blocked', reason: 'policy' }; } + const tenantId = opts.tenantId ?? 'default'; + const argsJson = JSON.stringify(opts.args); const now = new Date().toISOString(); - const tenantId = opts.tenantId ?? 'default'; // Check for an existing, non-expired permission record for this plugin + endpoint + args + tenant const existing = await opts.db.db @@ -406,6 +415,61 @@ export async function enforcePermission( }; } + // Evaluate limits + const applicableLimits = [ + ...(opts.globalLimits || []).map((l) => ({ + ...l, + // global configs default to 'global' scope + resolvedScope: l.scope === 'tenant' ? `tenant:${tenantId}` : `global`, + })), + ...(opts.pluginLimits || []).map((l) => ({ + ...l, + // plugin configs natively apply to the plugin + resolvedScope: `plugin:${opts.pluginId}`, + })), + ].filter((l) => !l.riskLevel || l.riskLevel === opts.riskLevel); + + if (applicableLimits.length > 0) { + const { sql } = await import('kysely'); + const nowTs = Date.now(); + for (const limit of applicableLimits) { + const windowMs = parseDurationMs(limit.window); + const epoch = Math.floor(nowTs / windowMs); + // Hash properties to create a stable limit bucket + const limitFingerprint = `${limit.type}_${limit.max}_${limit.window}`; + const key = `usage:${limit.resolvedScope}:${limitFingerprint}:${epoch}`; + const expiresAt = new Date(nowTs + windowMs).toISOString(); + + const res = await opts.db.db + .insertInto('corsair_usage_counters') + .values({ key, count: 1, expires_at: expiresAt }) + .onConflict((oc) => + oc + .column('key') + .doUpdateSet({ count: sql`corsair_usage_counters.count + 1` }), + ) + .returning('count') + .executeTakeFirst(); + + if (res && res.count > limit.max) { + console.log( + `[corsair/${opts.pluginId}] '${opts.endpointPath}' blocked — ${limit.type} exceeded.`, + `\n Action: ${description}`, + `\n Limit: ${limit.max} per ${limit.window}`, + ); + return { + result: 'blocked', + reason: + limit.type === 'budget' + ? 'budget_exhausted' + : 'rate_limit_exceeded', + }; + } + } + } + + if (policy === 'allow') return { result: 'allow' }; + // No existing actionable record — create a new pending approval request const id = uuidv4(); const token = randomBytes(32).toString('hex'); diff --git a/packages/corsair/core/plugins/index.ts b/packages/corsair/core/plugins/index.ts index 1160a1b95..169956893 100644 --- a/packages/corsair/core/plugins/index.ts +++ b/packages/corsair/core/plugins/index.ts @@ -37,6 +37,20 @@ import type { */ export type EndpointRiskLevel = 'read' | 'write' | 'destructive'; +/** + * Usage limit configuration for rate limiting and budget quotas. + */ +export type UsageLimit = { + /** The maximum number of tool calls allowed within the time window. */ + max: number; + /** The time window for the limit, e.g., '1m' for a minute, '1d' for a day. */ + window: string; + /** Dictates the 'blocked' reason returned when the limit is exceeded. */ + type: 'rate_limit' | 'budget'; + /** Optional filter to restrict the limit to a specific risk level. */ + riskLevel?: EndpointRiskLevel; +}; + /** * Permission mode controlling what the AI agent is allowed to do by default. * @@ -158,6 +172,10 @@ export type PluginPermissionsConfig = { * Only valid paths for this specific plugin compile — typos are type errors. */ overrides?: Partial, PermissionPolicy>>; + /** + * Limits applied to calls made to this plugin. + */ + limits?: UsageLimit[]; }; // ───────────────────────────────────────────────────────────────────────────── @@ -737,6 +755,10 @@ export type CorsairPermissionsOptions = { | 'synchronous' | 'asynchronous' | (() => 'synchronous' | 'asynchronous'); + /** + * Global usage limits applied to all calls across all plugins. + */ + limits?: (UsageLimit & { scope?: 'global' | 'tenant' })[]; /** * @deprecated Use `manual.onApprovalRequired` instead. TODO: delete ~April 2026. */ diff --git a/packages/corsair/db/index.ts b/packages/corsair/db/index.ts index 711981bd0..a51561f4f 100644 --- a/packages/corsair/db/index.ts +++ b/packages/corsair/db/index.ts @@ -157,6 +157,27 @@ export type CorsairPermissionInsert = { error?: string | null; }; +// ───────────────────────────────────────────────────────────────────────────── +// Corsair Usage Counters +// ───────────────────────────────────────────────────────────────────────────── + +export const CorsairUsageSchema = z.object({ + /** Composite key identifying the limit scope and time window epoch */ + key: z.string(), + /** Number of logical tool attempts within this window */ + count: z.number().int().default(1), + /** ISO8601 timestamp — when this counter can be safely pruned */ + expires_at: z.string(), +}); + +export type CorsairUsage = z.infer; + +export type CorsairUsageInsert = { + key: string; + count?: number; + expires_at: string; +}; + // ───────────────────────────────────────────────────────────────────────────── // Table Names // ───────────────────────────────────────────────────────────────────────────── @@ -167,6 +188,7 @@ export type CorsairTableName = | 'corsair_entities' | 'corsair_events' | 'corsair_permissions' + | 'corsair_usage_counters' | (string & {}); // ───────────────────────────────────────────────────────────────────────────── @@ -178,6 +200,7 @@ export type CorsairTableRow = { corsair_accounts: CorsairAccount; corsair_entities: CorsairEntity; corsair_events: CorsairEvent; + corsair_usage_counters: CorsairUsage; }; export type TableRowType = @@ -234,6 +257,7 @@ export type CorsairTableInsert = { corsair_accounts: CorsairAccountInsert; corsair_entities: CorsairEntityInsert; corsair_events: CorsairEventInsert; + corsair_usage_counters: CorsairUsageInsert; }; export type TableInsertType = @@ -266,6 +290,7 @@ export type CorsairTableUpdate = { corsair_accounts: CorsairAccountUpdate; corsair_entities: CorsairEntityUpdate; corsair_events: CorsairEventUpdate; + corsair_usage_counters: Partial; }; export type TableUpdateType = diff --git a/packages/corsair/db/kysely/database.ts b/packages/corsair/db/kysely/database.ts index 813ad8a29..bcde36c26 100644 --- a/packages/corsair/db/kysely/database.ts +++ b/packages/corsair/db/kysely/database.ts @@ -9,6 +9,7 @@ import type { CorsairEvent, CorsairIntegration, CorsairPermission, + CorsairUsage, } from '../index'; import { SqliteDatePlugin } from './sqlite-date-plugin.js'; @@ -18,6 +19,7 @@ export type CorsairKyselyDatabase = { corsair_entities: CorsairEntity; corsair_events: CorsairEvent; corsair_permissions: CorsairPermission; + corsair_usage_counters: CorsairUsage; }; export type CorsairDatabase = { diff --git a/packages/corsair/tests/bind-limits.test.ts b/packages/corsair/tests/bind-limits.test.ts new file mode 100644 index 000000000..292089a4a --- /dev/null +++ b/packages/corsair/tests/bind-limits.test.ts @@ -0,0 +1,55 @@ +import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; +import { bindEndpointsRecursively } from '../core/endpoints/bind'; +import { createTestDatabase } from './setup-db'; + +describe('bindEndpointsRecursively - Global Limits Bypass', () => { + let testDb: ReturnType; + + beforeEach(() => { + testDb = createTestDatabase(); + }); + + afterEach(() => { + testDb.cleanup(); + }); + + it('enforces global limits on a plugin that does NOT define options.permissions', async () => { + const endpoints = { + testMethod: () => 'success', + }; + + const tree: Record = {}; + + bindEndpointsRecursively({ + endpoints, + hooks: undefined, + ctx: {}, + tree, + pluginId: 'test-plugin', + errorHandlers: { + handle: (e: any) => { + throw e; + }, + wrap: (e: any) => { + throw e; + }, + }, + currentPath: [], + // NOTE: permissionsConfig is intentionally undefined! + permissionsConfig: undefined, + permissionsOptions: { + limits: [{ max: 1, window: '1m', type: 'rate_limit' }], + }, + database: testDb.database, + }); + + const boundMethod = tree.testMethod as () => Promise; + + // First call should succeed and consume the 1 quota + const result1 = await boundMethod(); + expect(result1).toBe('success'); + + // Second call should fail with rate limit error + await expect(boundMethod()).rejects.toThrow('rate limited'); + }); +}); diff --git a/packages/corsair/tests/permissions-limits.test.ts b/packages/corsair/tests/permissions-limits.test.ts new file mode 100644 index 000000000..c07aa2670 --- /dev/null +++ b/packages/corsair/tests/permissions-limits.test.ts @@ -0,0 +1,241 @@ +import { enforcePermission } from '../core/permissions'; +import { createTestDatabase } from './setup-db'; + +describe('enforcePermission - Rate Limits and Budget', () => { + let testDb: ReturnType; + + beforeEach(() => { + testDb = createTestDatabase(); + }); + + afterEach(() => { + testDb.cleanup(); + jest.useRealTimers(); + }); + + it('allows calls below the limit and blocks when exceeded', async () => { + const limits = [{ max: 2, window: '1m', type: 'rate_limit' as const }]; + + const call = () => + enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args: {}, + mode: 'open', + riskLevel: 'read', + db: testDb.database, + pluginLimits: limits, + }); + + // 1st call + expect((await call()).result).toBe('allow'); + // 2nd call + expect((await call()).result).toBe('allow'); + + // 3rd call - blocked + const blockedRes = await call(); + expect(blockedRes.result).toBe('blocked'); + expect(blockedRes.reason).toBe('rate_limit_exceeded'); + }); + + it('returns budget_exhausted for budget limits', async () => { + const limits = [{ max: 1, window: '1d', type: 'budget' as const }]; + + const call = () => + enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args: {}, + mode: 'open', + riskLevel: 'write', + db: testDb.database, + pluginLimits: limits, + }); + + expect((await call()).result).toBe('allow'); + + const blockedRes = await call(); + expect(blockedRes.result).toBe('blocked'); + expect(blockedRes.reason).toBe('budget_exhausted'); + }); + + it('resets the counter when the time window passes', async () => { + jest.useFakeTimers(); + + const limits = [{ max: 1, window: '10s', type: 'rate_limit' as const }]; + + const call = () => + enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args: {}, + mode: 'open', + riskLevel: 'read', + db: testDb.database, + globalLimits: limits, + }); + + expect((await call()).result).toBe('allow'); + expect((await call()).result).toBe('blocked'); + + // Advance time by 11 seconds to enter the next window epoch + jest.advanceTimersByTime(11000); + + // Now it should be allowed again + expect((await call()).result).toBe('allow'); + }); + + it('isolates counters by tenant and plugin scope', async () => { + const call = (tenantId: string, pluginId: string, limits: any) => + enforcePermission({ + pluginId, + endpointPath: 'test.endpoint', + args: {}, + mode: 'open', + riskLevel: 'read', + db: testDb.database, + tenantId, + ...limits, + }); + + // Global limit (applies to all plugins and tenants) + await call('tenant1', 'pluginA', { + globalLimits: [{ max: 10, window: '1m', type: 'rate_limit' }], + }); + + // Tenant limit (applies to tenant1 across all plugins) + await call('tenant1', 'pluginA', { + globalLimits: [ + { max: 1, window: '1m', type: 'rate_limit', scope: 'tenant' }, + ], + }); + + // Tenant limit exhausted for tenant1, should block + const blocked = await call('tenant1', 'pluginB', { + globalLimits: [ + { max: 1, window: '1m', type: 'rate_limit', scope: 'tenant' }, + ], + }); + expect(blocked.result).toBe('blocked'); + + // Tenant2 should still be allowed + const allowed = await call('tenant2', 'pluginB', { + globalLimits: [ + { max: 1, window: '1m', type: 'rate_limit', scope: 'tenant' }, + ], + }); + expect(allowed.result).toBe('allow'); + }); + + it('isolates plugin limits by pluginId', async () => { + const call = (pluginId: string, pluginLimits: any) => + enforcePermission({ + pluginId, + endpointPath: 'test.endpoint', + args: {}, + mode: 'open', + riskLevel: 'read', + db: testDb.database, + pluginLimits, + }); + + const limits = [{ max: 1, window: '1m', type: 'rate_limit' as const }]; + + // First call through pluginA is allowed + const firstA = await call('pluginA', limits); + expect(firstA.result).toBe('allow'); + + // Second call through pluginA is blocked + const secondA = await call('pluginA', limits); + expect(secondA.result).toBe('blocked'); + expect(secondA.reason).toBe('rate_limit_exceeded'); + + // pluginB with the same plugin limit remains allowed + const firstB = await call('pluginB', limits); + expect(firstB.result).toBe('allow'); + }); + + it('does not increment counter for policy-denied calls', async () => { + const limits = [{ max: 1, window: '1m', type: 'rate_limit' as const }]; + + // This call is denied by policy + const deniedRes = await enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args: {}, + mode: 'readonly', // Policy deny! + riskLevel: 'write', + db: testDb.database, + pluginLimits: limits, + }); + + expect(deniedRes.result).toBe('blocked'); + expect(deniedRes.reason).toBe('policy'); + + // Now an allowed call should pass because the quota was not consumed + const allowedRes = await enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args: {}, + mode: 'open', + riskLevel: 'write', + db: testDb.database, + pluginLimits: limits, + }); + + expect(allowedRes.result).toBe('allow'); + }); + it('does not double-charge quota when an approved request is replayed', async () => { + const limits = [{ max: 1, window: '1m', type: 'rate_limit' as const }]; + const args = { data: 'test' }; + + // 1. Initial request: should hit limit check, consume 1 quota, and return blocked (pending) + const initialRes = await enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args, + mode: 'open', + override: 'require_approval', // Triggers pending record creation + riskLevel: 'write', + db: testDb.database, + pluginLimits: limits, + approvalMode: 'asynchronous', // Return immediately + }); + + expect(initialRes.result).toBe('blocked'); + expect(initialRes.reason).toBe('pending'); + const token = initialRes.token!; + const id = initialRes.id!; + + // 2. Simulate human approval + await testDb.database.db + .updateTable('corsair_permissions') + .set({ status: 'approved', updated_at: new Date() }) + .where('id', '=', id) + .execute(); + + // 3. Replay the request (simulate executePermission or client retry) + const replayRes = await enforcePermission({ + pluginId: 'test-plugin', + endpointPath: 'test.endpoint', + args, + mode: 'open', + override: 'require_approval', + riskLevel: 'write', + db: testDb.database, + pluginLimits: limits, + }); + + // It should be allowed, because the existing approved record bypasses the limits check! + expect(replayRes.result).toBe('allow'); + + // 4. Verify that the quota was NOT double-charged + // If it were double-charged, the counter would be 2. Let's check the database. + const usage = await testDb.database.db + .selectFrom('corsair_usage_counters') + .selectAll() + .executeTakeFirst(); + + expect(usage?.count).toBe(1); // Still 1! + }); +}); diff --git a/packages/corsair/tests/setup-db.ts b/packages/corsair/tests/setup-db.ts index ed6159d4c..1427f4c5c 100644 --- a/packages/corsair/tests/setup-db.ts +++ b/packages/corsair/tests/setup-db.ts @@ -55,6 +55,25 @@ export function createTestDatabase(): { payload TEXT NOT NULL, status TEXT ); + + CREATE TABLE IF NOT EXISTS corsair_usage_counters ( + key TEXT PRIMARY KEY, + count INTEGER NOT NULL DEFAULT 1, + expires_at TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS corsair_permissions ( + id TEXT PRIMARY KEY, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + token TEXT NOT NULL, + plugin TEXT NOT NULL, + endpoint TEXT NOT NULL, + args TEXT NOT NULL, + tenant_id TEXT NOT NULL, + status TEXT NOT NULL, + expires_at TEXT NOT NULL + ); `); const db = new Kysely({