diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e786ac5..9a227672 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,16 @@ This project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). Packages without a separate changelog are covered by the cross-package notes below. -## [Unreleased] +## [Unreleased - Major] + +### Added + +- Registration caps for devices (nodes), agents and human members: a new registration past the deployment's limit is refused with HTTP 429 `plan_limit_exceeded` carrying `limit`, `max`, `current` and the host's context (`plan`, `upgrade_url` on the hosted gateway), on HTTP and on the node `agent.register` frame. Existing nodes and agents re-register freely. +- Self-hosted operators can cap their own instance with `config.limits`; absent keys stay unlimited. + +### Breaking changes + +- The engine no longer ships plan tiers: the `PLAN_LIMITS` table and every plan name/number are gone, `PlanLimits` is replaced by the metric-keyed `Limits` (`null` = unlimited), and all limits come from the deployment's entitlements provider. Self-hosted deployments are unlimited by default. ## [11.0.0] - 2026-10-22 diff --git a/README.md b/README.md index aee41832..a5e55a17 100644 --- a/README.md +++ b/README.md @@ -328,31 +328,45 @@ deployments — where a control plane owns the canonical key — leave it unset and rotate through that plane, so the route answers `workspace_key_rotation_disabled`. -### Throttling +### Throttling and limits -Two different conditions return HTTP 429, and the error code tells you which: +Three conditions return HTTP 429, and the error code tells you which: | Code | What it is | Retry | | --- | --- | --- | | `rate_limit_exceeded` | Per-minute request ceiling | Always clears inside the window; `Retry-After` is at most 60 | -| `plan_limit_exceeded` | Plan API-call quota for the current billing period | Clears only at the period boundary; `Retry-After` reports it | +| `plan_limit_exceeded` | API-call quota for the current usage period | Clears only at the period boundary; `Retry-After` reports it | +| `plan_limit_exceeded` with `limit` | Registration cap (`devices`, `agents` or `humans`) hit by a NEW node/agent registration; body carries `max`, `current` and the host's context (the hosted gateway adds `plan` and `upgrade_url`) | Clears when a registration is removed or the host raises the limit; no `Retry-After` | -Both carry `Retry-After` and `X-RateLimit-Reset`; successful responses carry -`X-RateLimit-Limit` and `X-RateLimit-Remaining`. `Retry-After` is authoritative +The first two carry `Retry-After` and `X-RateLimit-Reset`; when a finite per-minute +limit applies, successful responses carry `X-RateLimit-Limit` and +`X-RateLimit-Remaining` (an unlimited workspace — the self-hosted default — gets +neither). `Retry-After` is authoritative — a value beyond your retry budget means retrying cannot help, so fail fast rather than spending the budget on it. -A workspace's `plan` (returned by `GET /v1/workspace`) sets both limits: - -| Plan | Agents | API calls / month | Requests / minute | -| --- | --- | --- | --- | -| `free` | 5 | 100,000 | 300 | -| `connect` | 10 | 10,000,000 | 1,200 | -| `pro` | 100 | 10,000,000 | 6,000 | -| `enterprise` | unlimited | unlimited | 30,000 | - -`connect` is sized for short-lived collaboration rooms. Self-hosted deployments -default to `selfhost`, which is unlimited. +The engine carries no tiers or numbers of its own. Every limit comes from the +deployment's `EntitlementsProvider`, which maps a workspace to a `Limits` +object keyed by metric (`messages`, `agents`, `file_bytes`, `api_calls`, +`rate_per_min`, `devices`, `humans`, `message_retention_days`; `null` is +unlimited). A device is any node a host enrolls through `POST /v1/nodes`; the +implicit direct node every agent gets for its own WebSocket is not one. + +- Self-hosted deployments are unlimited by default. To cap your own instance, + set static limits in the engine config; absent keys stay unlimited: + + ```ts + createNodeRuntime({ + dbPath: './relaycast.db', + baseUrl: 'https://relay.example.com', + config: { limits: { devices: 5, agents: 15, humans: 2, rate_per_min: 600 } }, + }); + ``` + +- A hosting product injects its own provider (`getLimits`, `getUsage`, optional + `getUsageResetAt` and `getLimitContext`). `getLimitContext` returns fields the + engine spreads verbatim onto the top level of a registration-cap 429 — on + `cast.agentrelay.com` that is `plan` and `upgrade_url`. `GET /v1/workspace` is limited in its own bucket, so the identity read used for credential validation and launch preflight stays reachable while the workspace's diff --git a/openapi.yaml b/openapi.yaml index 79f3809a..3cdacc15 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -19,6 +19,20 @@ info: credential validation and launch preflight stays reachable when the workspace's data-plane traffic is at its ceiling. + The deployment's entitlements provider may also cap how many devices + (nodes), agents and human members a workspace can register. Enrolling a NEW + node (`POST /nodes`) or registering a NEW agent (`POST /agents`, + `POST /a2a/register`) beyond the cap returns HTTP 429 `plan_limit_exceeded` + with the top-level fields `error_code`, `limit` (`devices`|`agents`|`humans`), + `max` and `current`, plus whatever context the provider supplies (the hosted + gateway adds `plan` and `upgrade_url`) so a client can explain the refusal + and offer an upgrade. An `agent.register` node-control frame at the cap is + answered on the WebSocket with an `error` frame carrying `code: + plan_limit_exceeded` and the same fields, not with an HTTP response. + Re-enrolling or rotating an existing node or agent is never capped. No + `Retry-After` is sent: the cap clears only when a registration is removed + or the provider raises the limit. + Workspace delivery capacity applies to all local message producers. A write that would exceed it returns HTTP 429 with error code `workspace_delivery_depth_exceeded` and `Retry-After: 30`. A required @@ -182,8 +196,60 @@ components: application/json: schema: $ref: '#/components/schemas/ErrorResponse' + PlanLimitExceeded: + description: >- + The workspace is at its registration cap for the requested kind. + Only a NEW registration is refused; re-enrolling or rotating an existing + node or agent never is. The cap clears when a registration is removed or + the provider raises the limit, so no `Retry-After` is sent. + content: + application/json: + schema: + $ref: '#/components/schemas/PlanLimitExceededResponse' schemas: + PlanLimitExceededResponse: + type: object + required: [ok, error, error_code, limit, max, current] + description: >- + Beyond the required fields, the deployment's entitlements provider may + add context at the top level (for example `plan` and `upgrade_url` on + the hosted gateway); the engine passes it through verbatim. + additionalProperties: true + properties: + ok: + type: boolean + enum: [false] + error: + type: object + required: [code, message] + properties: + code: + type: string + enum: [plan_limit_exceeded] + message: + type: string + error_code: + type: string + enum: [plan_limit_exceeded] + limit: + type: string + enum: [devices, agents, humans] + description: Which registration cap refused the request. + max: + type: integer + description: The cap for `limit`. + current: + type: integer + description: How many are registered now. + plan: + type: string + description: Provider-supplied, optional. The hosted gateway reports the workspace's billing tier here. + upgrade_url: + type: string + format: uri + description: Provider-supplied, optional. Where the workspace can raise its limits. + WorkspaceProvenanceInput: type: object description: >- @@ -261,8 +327,7 @@ components: description: Workspace API key (only returned on creation) plan: type: string - enum: [free, connect, pro, enterprise, selfhost] - description: Stored plan that sets the workspace's usage limits. `connect` is sized for short-lived collaboration rooms. + description: Host-assigned plan key stored on the workspace; opaque to the engine, which takes every limit from its entitlements provider. created_at: type: string format: date-time @@ -2495,6 +2560,8 @@ paths: description: Malformed registration identity_key (invalid_agent_identity_key) '409': description: Agent name already exists in this workspace + '429': + $ref: '#/components/responses/PlanLimitExceeded' get: summary: List agents @@ -5455,6 +5522,8 @@ paths: application/json: schema: $ref: '#/components/schemas/ErrorResponse' + '429': + $ref: '#/components/responses/PlanLimitExceeded' get: summary: List nodes description: > @@ -6774,6 +6843,8 @@ paths: application/json: schema: type: object + '429': + $ref: '#/components/responses/PlanLimitExceeded' /a2a/agents: get: diff --git a/packages/engine/CHANGELOG.md b/packages/engine/CHANGELOG.md index 644e050d..ddd3191d 100644 --- a/packages/engine/CHANGELOG.md +++ b/packages/engine/CHANGELOG.md @@ -8,7 +8,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [Unreleased - Major] + +### Added + +- Registration caps are enforced at every path that creates a node or agent: `POST /v1/nodes`, `POST /v1/agents`, `POST /v1/a2a/register` and the `agent.register` node frame. A refusal is HTTP 429 `plan_limit_exceeded` with top-level `error_code`, `limit`, `max`, `current` and the provider's `getLimitContext` fields spread verbatim (the same fields on the WebSocket `error` frame). Re-enrolling or rotating an existing node/agent is never capped; an agent's implicit direct node is not a device, and the `node_direct_` id prefix is reserved for it (`reserved_node_id`). +- `Limits` (metric-keyed, `null` = unlimited) with `UNLIMITED_LIMITS` and `resolveLimits`; `UsageMetric` gains `devices`/`humans` and `EntitlementsProvider.getUsage` reports live registration counts for them (`countRegistrations` is exported for providers). Optional `EntitlementsProvider.getLimitContext` supplies the refusal's host context. +- `EngineConfig.limits` sets static limits for a self-hosted instance; `StaticEntitlementsProvider(kv, limits, db)` applies them and counts registrations from `db`. +- `workspaces.limit_overrides` (migration `0076`) is a host-owned JSON column the engine never interprets. +- `createRegistrationCapacityGuard` and `HandleNodeControlMessageArgs.capacity` let a host's node-control transport enforce the agent cap on `agent.register`. + +### Breaking changes + +- `PLAN_LIMITS` and every plan name, number and rank are removed from the engine; `PlanLimits` is replaced by `Limits`, whose metered fields are `number | null` (`null` = unlimited) instead of `number`/`Infinity`. `StaticEntitlementsProvider` no longer reads `workspace.plan`: it is unlimited unless `config.limits` is set. Hosts that relied on the built-in tiers must supply them from their own provider. +- `WorkspaceDeliveryPolicyConfig.resolve` receives `plan: string | null` (the row's host-assigned key, no `'free'` default). +- The `rate_limit_exceeded` and `plan_limit_exceeded` quota messages no longer name a plan. ## [11.0.0] - 2026-10-22 diff --git a/packages/engine/scripts/public-consumer-compile.mjs b/packages/engine/scripts/public-consumer-compile.mjs index 15f9947e..a19659af 100644 --- a/packages/engine/scripts/public-consumer-compile.mjs +++ b/packages/engine/scripts/public-consumer-compile.mjs @@ -22,7 +22,7 @@ const config = { workspaceDelivery: { resolve: async (workspace) => { const id: string = workspace.id; - const plan: string = workspace.plan; + const plan: string | null = workspace.plan; await Promise.resolve(); return id && plan === 'enterprise' ? { cap: 5000, reserve: 0 } : undefined; }, diff --git a/packages/engine/src/__tests__/conformance/connectObserver.test.ts b/packages/engine/src/__tests__/conformance/connectObserver.test.ts index 71cb417a..9c3d3fd0 100644 --- a/packages/engine/src/__tests__/conformance/connectObserver.test.ts +++ b/packages/engine/src/__tests__/conformance/connectObserver.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { eq } from 'drizzle-orm'; -import { messages, observerTokens, workspaces } from '../../db/schema.js'; +import { messages, observerTokens } from '../../db/schema.js'; import { dmFilteredCountConversationChunkSize, dmMessagePageLimit } from '../../engine/dmAll.js'; import { createWorkspace, makeNodeStack, registerAgent, type TestStack } from './harness.js'; @@ -55,10 +55,6 @@ describe('Relay Connect observer capability', () => { it('lists more DM conversations than one D1 statement can bind, for the workspace and an observer', async () => { const room = await createWorkspace(stack.app, 'connect-many-dms'); - await stack.runtime.deps.db - .update(workspaces) - .set({ plan: 'enterprise' }) - .where(eq(workspaces.id, room.workspaceId)); const alice = await registerAgent(stack.app, room.workspaceKey, 'alice'); const peers = Array.from({ length: 505 }, (_, index) => `peer-${String(index).padStart(3, '0')}`); const conversationIds: string[] = []; diff --git a/packages/engine/src/__tests__/conformance/harness.ts b/packages/engine/src/__tests__/conformance/harness.ts index 99983d84..0f608260 100644 --- a/packages/engine/src/__tests__/conformance/harness.ts +++ b/packages/engine/src/__tests__/conformance/harness.ts @@ -53,6 +53,8 @@ export function makeNodeStack(options?: { workspaceBootstrapProofRequired?: boolean; allowWorkspaceKeyRotation?: boolean; entitlements?: EntitlementsProvider; + /** Static limits for the default provider (`EngineConfig.limits`). */ + limits?: EngineConfig['limits']; }): TestStack { const tasks = new BackgroundTasks(); observePresence(); @@ -79,6 +81,7 @@ export function makeNodeStack(options?: { workspaceBootstrapSecret: options?.workspaceBootstrapSecret ?? 'test-bootstrap-secret', workspaceBootstrapProofRequired: options?.workspaceBootstrapProofRequired, allowWorkspaceKeyRotation: options?.allowWorkspaceKeyRotation, + limits: options?.limits, }, entitlements: options?.entitlements, // Disable the auto-sweep timer; tests drive presence.sweep() explicitly. diff --git a/packages/engine/src/__tests__/conformance/inboxArchivedMentions.test.ts b/packages/engine/src/__tests__/conformance/inboxArchivedMentions.test.ts index d0116b7d..d9c2f7cc 100644 --- a/packages/engine/src/__tests__/conformance/inboxArchivedMentions.test.ts +++ b/packages/engine/src/__tests__/conformance/inboxArchivedMentions.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { makeNodeStack, createWorkspace, registerAgent, type TestStack } from './harness.js'; +import { resolveLimits } from '../../ports/entitlements.js'; /** * Regression for the captured cross-channel leak: a native MCP post result @@ -13,9 +14,9 @@ import { makeNodeStack, createWorkspace, registerAgent, type TestStack } from '. * non-member never seeing the mention. */ // Test-only entitlements: a high per-minute ceiling so the bounded >200-candidate -// fixture is not rate-limited by the free plan. No production limit changes. +// fixture exercises the limiter without being throttled. const testEntitlements = { - getLimits: async () => ({ messages: 1_000_000, agents: 1_000_000, file_bytes: 1_000_000_000, api_calls: 1_000_000, rate_per_min: 100_000 }), + getLimits: async () => resolveLimits({ messages: 1_000_000, agents: 1_000_000, file_bytes: 1_000_000_000, api_calls: 1_000_000, rate_per_min: 100_000 }), getUsage: async () => 0, }; diff --git a/packages/engine/src/__tests__/conformance/planCapacity.test.ts b/packages/engine/src/__tests__/conformance/planCapacity.test.ts new file mode 100644 index 00000000..62981903 --- /dev/null +++ b/packages/engine/src/__tests__/conformance/planCapacity.test.ts @@ -0,0 +1,289 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import type { EntitlementsProvider, Limits, Workspace } from '../../ports/index.js'; +import { resolveLimits } from '../../ports/entitlements.js'; +import { countRegistrations } from '../../engine/planCapacity.js'; +import { createWorkspace, FakeSocket, makeNodeStack, registerAgent, type TestStack } from './harness.js'; + +type Created = { workspaceKey: string; workspaceId: string }; + +/** The caps a hosted tier might carry; the engine only knows the numbers. */ +const CAPPED = { devices: 5, agents: 15, humans: 2 } as const; + +describe('registration caps', () => { + let stack: TestStack; + + afterEach(async () => { + vi.restoreAllMocks(); + await stack?.close(); + }); + + function enrollNode(ws: Created, nodeId: string, name: string, currentToken?: string) { + return stack.app.request('/v1/nodes', { + method: 'POST', + headers: { + 'content-type': 'application/json', + authorization: `Bearer ${ws.workspaceKey}`, + ...(currentToken ? { 'X-Relaycast-Node-Token': currentToken } : {}), + }, + body: JSON.stringify({ node_id: nodeId, name, role: 'broker', capabilities: [], max_agents: 4, tags: ['test'], version: 'v0' }), + }); + } + + function registerAgentRaw(ws: Created, body: Record) { + return stack.app.request('/v1/agents', { + method: 'POST', + headers: { 'content-type': 'application/json', authorization: `Bearer ${ws.workspaceKey}` }, + body: JSON.stringify(body), + }); + } + + const refusal = (limit: string, max: number, current: number) => ({ + ok: false, + error: { code: 'plan_limit_exceeded' }, + error_code: 'plan_limit_exceeded', + limit, + max, + current, + }); + + it('allows the 5th device and refuses the 6th with the machine-readable body', async () => { + stack = makeNodeStack({ limits: CAPPED }); + const ws = await createWorkspace(stack.app, 'cap-devices'); + + const tokens: string[] = []; + for (let i = 1; i <= 5; i += 1) { + const res = await enrollNode(ws, `node_${i}`, `device-${i}`); + expect(res.status, `device ${i}`).toBe(201); + tokens.push((await res.json() as { data: { token: string } }).data.token); + } + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'devices')).resolves.toBe(5); + + const sixth = await enrollNode(ws, 'node_6', 'device-6'); + expect(sixth.status).toBe(429); + const body = await sixth.json() as Record; + expect(body).toMatchObject(refusal('devices', 5, 5)); + // The engine adds nothing it does not know: no tier name without a provider context. + expect(body).not.toHaveProperty('plan'); + + // Re-enrolling (rotating) an established device is not a new registration. + const rotate = await enrollNode(ws, 'node_1', 'device-1', tokens[0]); + expect(rotate.status).toBe(201); + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'devices')).resolves.toBe(5); + }); + + it('does not count an agent\'s implicit direct node as a device', async () => { + stack = makeNodeStack({ limits: { devices: 1 } }); + const ws = await createWorkspace(stack.app, 'cap-direct'); + expect((await enrollNode(ws, 'node_1', 'only-device')).status).toBe(201); + + const agent = await registerAgent(stack.app, ws.workspaceKey, 'worker'); + const nodeToken = await stack.app.request('/v1/agent/node-token', { + method: 'POST', + headers: { authorization: `Bearer ${agent.token}` }, + }); + expect(nodeToken.status).toBe(200); + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'devices')).resolves.toBe(1); + expect((await enrollNode(ws, 'node_2', 'second-device')).status).toBe(429); + }); + + it('refuses host enrollment into the reserved implicit-direct id namespace', async () => { + stack = makeNodeStack({ limits: { devices: 1 } }); + const ws = await createWorkspace(stack.app, 'cap-reserved-id'); + expect((await enrollNode(ws, 'node_1', 'only-device')).status).toBe(201); + + const spoofed = await stack.app.request('/v1/nodes', { + method: 'POST', + headers: { 'content-type': 'application/json', authorization: `Bearer ${ws.workspaceKey}` }, + body: JSON.stringify({ node_id: 'node_direct_custom', name: 'spoof', role: 'direct', max_agents: 1 }), + }); + expect(spoofed.status).toBe(400); + await expect(spoofed.json()).resolves.toMatchObject({ error: { code: 'reserved_node_id' } }); + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'devices')).resolves.toBe(1); + + // A near miss on the prefix (`_` is a LIKE wildcard, not a GLOB one) is an + // ordinary host-enrolled direct node: not reserved, and counted. + const nearMiss = await stack.app.request('/v1/nodes', { + method: 'POST', + headers: { 'content-type': 'application/json', authorization: `Bearer ${ws.workspaceKey}` }, + body: JSON.stringify({ node_id: 'nodeXdirectXcustom', name: 'near-miss', role: 'direct', max_agents: 1 }), + }); + expect(nearMiss.status).toBe(429); + await expect(nearMiss.json()).resolves.toMatchObject(refusal('devices', 1, 1)); + }); + + it('counts an explicitly enrolled direct node as a device, so a direct-then-broker re-enroll cannot bypass the cap', async () => { + stack = makeNodeStack({ limits: { devices: 1 } }); + const ws = await createWorkspace(stack.app, 'cap-explicit-direct'); + const direct = await stack.app.request('/v1/nodes', { + method: 'POST', + headers: { 'content-type': 'application/json', authorization: `Bearer ${ws.workspaceKey}` }, + body: JSON.stringify({ node_id: 'node_d', name: 'webhook', kind: 'http_push', role: 'direct', max_agents: 1, delivery: { url: 'https://hooks.example/relay' } }), + }); + expect(direct.status).toBe(201); + const token = (await direct.json() as { data: { token: string } }).data.token; + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'devices')).resolves.toBe(1); + + const broker = await enrollNode(ws, 'node_b', 'broker'); + expect(broker.status).toBe(429); + await expect(broker.json()).resolves.toMatchObject(refusal('devices', 1, 1)); + + // Re-enrolling the counted node as a broker changes its shape, not the count. + const reshaped = await enrollNode(ws, 'node_d', 'webhook', token); + expect(reshaped.status).toBe(201); + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'devices')).resolves.toBe(1); + expect((await enrollNode(ws, 'node_b', 'broker')).status).toBe(429); + }); + + it('allows the 15th agent and refuses the 16th; humans are counted separately', async () => { + stack = makeNodeStack({ limits: CAPPED }); + const ws = await createWorkspace(stack.app, 'cap-agents'); + + for (let i = 1; i <= 15; i += 1) { + expect((await registerAgentRaw(ws, { name: `agent-${i}` })).status, `agent ${i}`).toBe(201); + } + const sixteenth = await registerAgentRaw(ws, { name: 'agent-16' }); + expect(sixteenth.status).toBe(429); + await expect(sixteenth.json()).resolves.toMatchObject(refusal('agents', 15, 15)); + + // Humans have their own cap and do not consume agent slots. + expect((await registerAgentRaw(ws, { name: 'will', type: 'human' })).status).toBe(201); + expect((await registerAgentRaw(ws, { name: 'khaliq', type: 'human' })).status).toBe(201); + const third = await registerAgentRaw(ws, { name: 'third-human', type: 'human' }); + expect(third.status).toBe(429); + await expect(third.json()).resolves.toMatchObject(refusal('humans', 2, 2)); + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'agents')).resolves.toBe(15); + await expect(countRegistrations(stack.runtime.handle.db, ws.workspaceId, 'humans')).resolves.toBe(2); + }); + + it('answers an established name with the conflict, never the cap', async () => { + stack = makeNodeStack({ limits: { agents: 1 } }); + const ws = await createWorkspace(stack.app, 'cap-conflict'); + expect((await registerAgentRaw(ws, { name: 'only' })).status).toBe(201); + const again = await registerAgentRaw(ws, { name: 'only' }); + expect(again.status).toBe(409); + await expect(again.json()).resolves.toMatchObject({ error: { code: 'agent_already_exists' } }); + }); + + it('refuses a WebSocket agent.register at the cap with the same detail on the error frame', async () => { + stack = makeNodeStack({ limits: { agents: 1 } }); + const ws = await createWorkspace(stack.app, 'cap-ws'); + expect((await enrollNode(ws, 'node_a', 'alpha')).status).toBe(201); + + const sock = new FakeSocket(); + const handle = stack.runtime.realtime.attachNodeSocket(ws.workspaceId, 'node_a', sock); + await handle.handleMessage(JSON.stringify({ + v: 1, id: 'reg', type: 'node.register', name: 'alpha', node_id: 'node_a', + capabilities: [{ name: 'spawn:claude', kind: 'capacity' }], + max_agents: 4, tags: [], version: 'test-node', resume_cursor: null, + })); + await handle.handleMessage(JSON.stringify({ v: 1, id: 'a1', type: 'agent.register', name: 'worker-1' })); + expect(sock.ofType('reply').at(-1)).toMatchObject({ id: 'a1', ok: true }); + + await handle.handleMessage(JSON.stringify({ v: 1, id: 'a2', type: 'agent.register', name: 'worker-2' })); + expect(sock.ofType('error').at(-1)).toMatchObject({ + id: 'a2', + ok: false, + code: 'plan_limit_exceeded', + error_code: 'plan_limit_exceeded', + limit: 'agents', + max: 1, + current: 1, + }); + + // Re-registering the established name is a conflict, not a cap refusal. + await handle.handleMessage(JSON.stringify({ v: 1, id: 'a3', type: 'agent.register', name: 'worker-1' })); + expect(sock.ofType('error').at(-1)).toMatchObject({ id: 'a3', code: 'agent_already_exists' }); + }); + + it('counts an A2A proxy identity as an agent', async () => { + stack = makeNodeStack({ limits: { agents: 1 } }); + vi.spyOn(globalThis, 'fetch').mockImplementation(async () => new Response('{}', { status: 200 })); + const ws = await createWorkspace(stack.app, 'cap-a2a'); + expect((await registerAgentRaw(ws, { name: 'local' })).status).toBe(201); + + const refused = await stack.app.request('/v1/a2a/register', { + method: 'POST', + headers: { 'content-type': 'application/json', authorization: `Bearer ${ws.workspaceKey}` }, + body: JSON.stringify({ + agent_card: { + name: 'northwind', + url: 'https://northwind.example/a2a/rpc', + version: '1.0.0', + skills: [{ id: 'lead', name: 'lead' }], + }, + }), + }); + expect(refused.status).toBe(429); + await expect(refused.json()).resolves.toMatchObject(refusal('agents', 1, 1)); + }); + + it('caps nothing by default', async () => { + stack = makeNodeStack(); + const ws = await createWorkspace(stack.app, 'cap-default'); + for (let i = 1; i <= 6; i += 1) { + expect((await enrollNode(ws, `node_${i}`, `device-${i}`)).status).toBe(201); + } + for (let i = 1; i <= 3; i += 1) { + expect((await registerAgentRaw(ws, { name: `human-${i}`, type: 'human' })).status).toBe(201); + } + }); + + it('fails open when the entitlements provider cannot answer', async () => { + const broken: EntitlementsProvider = { + async getLimits(): Promise { + throw new Error('billing unavailable'); + }, + async getUsage() { + return 0; + }, + }; + stack = makeNodeStack({ entitlements: broken }); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const ws = await createWorkspace(stack.app, 'cap-fail-open'); + expect((await enrollNode(ws, 'node_1', 'device')).status).toBe(201); + expect(warn).toHaveBeenCalledWith('[plan-capacity] entitlements lookup failed; allowing registration', expect.objectContaining({ metric: 'devices' })); + }); + + it('spreads the provider\'s limit context onto the refusal, on HTTP and on the error frame', async () => { + const hosted: EntitlementsProvider = { + async getLimits(_workspace: Workspace): Promise { + return resolveLimits({ devices: 1, agents: 1 }); + }, + getUsage: (workspaceId, metric) => + countRegistrations(stack.runtime.handle.db, workspaceId, metric as 'devices' | 'agents' | 'humans'), + async getLimitContext() { + return { plan: 'tier-small', upgrade_url: 'https://host.example/upgrade', ok: 'cannot clobber' }; + }, + }; + stack = makeNodeStack({ entitlements: hosted }); + const ws = await createWorkspace(stack.app, 'cap-context'); + expect((await enrollNode(ws, 'node_1', 'first')).status).toBe(201); + const second = await enrollNode(ws, 'node_2', 'second'); + expect(second.status).toBe(429); + await expect(second.json()).resolves.toEqual({ + ...refusal('devices', 1, 1), + error: { code: 'plan_limit_exceeded', message: expect.any(String) }, + plan: 'tier-small', + upgrade_url: 'https://host.example/upgrade', + }); + + const sock = new FakeSocket(); + const handle = stack.runtime.realtime.attachNodeSocket(ws.workspaceId, 'node_1', sock); + await handle.handleMessage(JSON.stringify({ + v: 1, id: 'reg', type: 'node.register', name: 'first', node_id: 'node_1', + capabilities: [], max_agents: 4, tags: [], version: 'test-node', resume_cursor: null, + })); + expect((await registerAgentRaw(ws, { name: 'only-agent' })).status).toBe(201); + await handle.handleMessage(JSON.stringify({ v: 1, id: 'a2', type: 'agent.register', name: 'second-agent' })); + expect(sock.ofType('error').at(-1)).toMatchObject({ + id: 'a2', + ok: false, + code: 'plan_limit_exceeded', + limit: 'agents', + max: 1, + current: 1, + plan: 'tier-small', + upgrade_url: 'https://host.example/upgrade', + }); + }); +}); diff --git a/packages/engine/src/__tests__/conformance/rateLimitContract.test.ts b/packages/engine/src/__tests__/conformance/rateLimitContract.test.ts index 6ba66ad6..4c460cb3 100644 --- a/packages/engine/src/__tests__/conformance/rateLimitContract.test.ts +++ b/packages/engine/src/__tests__/conformance/rateLimitContract.test.ts @@ -1,18 +1,12 @@ import { afterEach, describe, expect, it } from 'vitest'; -import type { EntitlementsProvider, PlanLimits, Workspace } from '../../ports/index.js'; +import { resolveLimits, type EntitlementsProvider, type Limits, type Workspace } from '../../ports/index.js'; import { InProcessRateLimiter } from '../../adapters/node/rate-limit.js'; import { createWorkspace, makeNodeStack, type TestStack } from './harness.js'; function entitlementsWithRate(ratePerMin: number): EntitlementsProvider { return { - async getLimits(_workspace: Workspace): Promise { - return { - messages: Infinity, - agents: Infinity, - file_bytes: Infinity, - api_calls: Infinity, - rate_per_min: ratePerMin, - }; + async getLimits(_workspace: Workspace): Promise { + return resolveLimits({ rate_per_min: ratePerMin }); }, async getUsage() { return 0; }, }; diff --git a/packages/engine/src/__tests__/conformance/usageQuota.test.ts b/packages/engine/src/__tests__/conformance/usageQuota.test.ts index 5fc2046d..68eb464c 100644 --- a/packages/engine/src/__tests__/conformance/usageQuota.test.ts +++ b/packages/engine/src/__tests__/conformance/usageQuota.test.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; -import type { EntitlementsProvider, PlanLimits, UsageMetric, Workspace } from '../../ports/index.js'; +import { resolveLimits, type EntitlementsProvider, type Limits, type UsageMetric, type Workspace } from '../../ports/index.js'; import { usageCounterKey, usagePeriodResetAt } from '../../engine/usage.js'; import { createWorkspace, makeNodeStack, registerAgent, type TestStack } from './harness.js'; @@ -9,14 +9,8 @@ class ApiCallQuotaEntitlements implements EntitlementsProvider { private readonly apiCallLimit: number, ) {} - async getLimits(_workspace: Workspace): Promise { - return { - messages: Infinity, - agents: Infinity, - file_bytes: Infinity, - api_calls: this.apiCallLimit, - rate_per_min: 300, - }; + async getLimits(_workspace: Workspace): Promise { + return resolveLimits({ api_calls: this.apiCallLimit, rate_per_min: 300 }); } // Reads the same period-scoped counter the usage tracker writes, so the @@ -39,7 +33,7 @@ describe('usage quotas', () => { await stack?.close(); }); - it('records authenticated API calls and enforces the api_calls plan limit', async () => { + it('records authenticated API calls and enforces the api_calls quota', async () => { stack = makeNodeStack({ entitlements: new ApiCallQuotaEntitlements(() => stack, 2), }); diff --git a/packages/engine/src/adapters/node/index.ts b/packages/engine/src/adapters/node/index.ts index 068ad8bb..b9a09738 100644 --- a/packages/engine/src/adapters/node/index.ts +++ b/packages/engine/src/adapters/node/index.ts @@ -161,7 +161,7 @@ export function createNodeRuntime(options: NodeRuntimeOptions): NodeRuntime { webhookQueue.start(); const auth = options.auth ?? new SqliteApiKeyAuthProvider(); - const entitlements = options.entitlements ?? new StaticEntitlementsProvider(kv); + const entitlements = options.entitlements ?? new StaticEntitlementsProvider(kv, options.config?.limits, db); // An explicit host-provided boundary is authoritative (the host may prune // outside this process). Otherwise report the same fallback the local Node @@ -193,6 +193,7 @@ export function createNodeRuntime(options: NodeRuntimeOptions): NodeRuntime { }; realtime.setNodeCompletionDeps(deps); + realtime.setEntitlements(entitlements); const runDeliveryMaintenance = createDeliveryMaintenanceRunner(deps); diff --git a/packages/engine/src/adapters/node/realtime.ts b/packages/engine/src/adapters/node/realtime.ts index 89e2e0da..92254941 100644 --- a/packages/engine/src/adapters/node/realtime.ts +++ b/packages/engine/src/adapters/node/realtime.ts @@ -26,6 +26,8 @@ import { providerAttachDecision } from '../../engine/placement.js'; import { drainNodeInvocations } from '../../engine/action.js'; import { observerAllowsEvent } from '../../engine/observerToken.js'; import type { InvocationCompletionDeps } from '../../engine/invocationCompletion.js'; +import type { EntitlementsProvider } from '../../ports/entitlements.js'; +import { createRegistrationCapacityGuard } from '../../engine/planCapacity.js'; import type { FleetRelaycastToBrokerMessage } from '@relaycast/types'; /** @@ -108,6 +110,7 @@ export class InProcessRealtime implements RealtimeBus, ConnectionRegistry, NodeC /** Serializes drains per node so concurrent triggers never overlap. */ private readonly nodeDrainChains = new Map>(); private nodeCompletionDeps: InvocationCompletionDeps | undefined; + private entitlements: EntitlementsProvider | undefined; private connectionSeq = 0; constructor(private readonly db: EngineDb) {} @@ -123,6 +126,11 @@ export class InProcessRealtime implements RealtimeBus, ConnectionRegistry, NodeC this.nodeCompletionDeps = deps; } + /** Enforce registration caps on `agent.register` frames. */ + setEntitlements(entitlements: EntitlementsProvider): void { + this.entitlements = entitlements; + } + private nodeKey(workspaceId: string, nodeId: string): string { return `${workspaceId}:${nodeId}`; } @@ -844,6 +852,9 @@ export class InProcessRealtime implements RealtimeBus, ConnectionRegistry, NodeC nodeId, socket, connectionId, + capacity: this.entitlements + ? createRegistrationCapacityGuard({ db: this.db, entitlements: this.entitlements, workspace: workspaceId }) + : undefined, raw, }); }, diff --git a/packages/engine/src/db/__tests__/compactMigrations.test.ts b/packages/engine/src/db/__tests__/compactMigrations.test.ts index 1ea2d008..182d4191 100644 --- a/packages/engine/src/db/__tests__/compactMigrations.test.ts +++ b/packages/engine/src/db/__tests__/compactMigrations.test.ts @@ -53,13 +53,14 @@ function snapshot(handle: SqliteDbHandle) { // so this regression continues to compare all pre-existing user data. // 0070 adds a nullable liveness-recovery marker without backfilling it. if (name === 'workspaces') { - // 0073 adds the rotation grace slot and the generation fence as - // additive columns; exclude them so this comparison stays - // byte-for-byte over every pre-existing value. + // 0073 adds the rotation grace slot and the generation fence, and + // 0076 the host limit overrides, as additive columns; exclude them so + // this comparison stays byte-for-byte over every pre-existing value. const { previous_api_key_hash: _previousApiKeyHash, previous_api_key_expires_at: _previousApiKeyExpiresAt, api_key_generation: _apiKeyGeneration, + limit_overrides: _limitOverrides, ...existing } = row as Record; return JSON.stringify(existing); @@ -158,16 +159,17 @@ function expectConstraintsPreserved(before: ReturnType, afte if (table.name === 'agents' && original) { return { ...table, sql: original.sql }; } - // 0073 appends the rotation grace slot and the generation fence. - // Remove only their exact additive DDL so every pre-existing workspace - // constraint stays under comparison. + // 0073 appends the rotation grace slot and the generation fence, and + // 0076 the host limit overrides. Remove only their exact additive DDL so + // every pre-existing workspace constraint stays under comparison. if (table.name === 'workspaces') { table = { ...table, sql: table.sql .replace(', previous_api_key_hash TEXT', '') .replace(', previous_api_key_expires_at INTEGER', '') - .replace(', api_key_generation INTEGER NOT NULL DEFAULT 0', ''), + .replace(', api_key_generation INTEGER NOT NULL DEFAULT 0', '') + .replace(', limit_overrides TEXT', ''), }; } // 0063 appends four nullable recovery fields. Remove only their exact diff --git a/packages/engine/src/db/__tests__/taskMigration.test.ts b/packages/engine/src/db/__tests__/taskMigration.test.ts index b5cfd36c..b0c03e80 100644 --- a/packages/engine/src/db/__tests__/taskMigration.test.ts +++ b/packages/engine/src/db/__tests__/taskMigration.test.ts @@ -39,6 +39,7 @@ it('upgrades existing actions and results additively, preserving rows and every '0073_workspace_key_rotation.sql', '0074_delivery_active_location_node_index.sql', '0075_message_mention_fts.sql', + '0076_workspace_limit_overrides.sql', ]); for (const table of before) { const columns = db.pragma(`table_info(${table.name})`) as { name: string }[]; diff --git a/packages/engine/src/db/migrations/0076_workspace_limit_overrides.sql b/packages/engine/src/db/migrations/0076_workspace_limit_overrides.sql new file mode 100644 index 00000000..83b0d43a --- /dev/null +++ b/packages/engine/src/db/migrations/0076_workspace_limit_overrides.sql @@ -0,0 +1,5 @@ +-- Host-owned per-workspace limit overrides (JSON). The hosting control plane +-- stores whatever its entitlements provider needs to derive this workspace's +-- limits; the engine never reads the value itself. NULL means the host's +-- defaults alone decide. +ALTER TABLE workspaces ADD COLUMN limit_overrides TEXT; diff --git a/packages/engine/src/db/schema.ts b/packages/engine/src/db/schema.ts index 4ee26daa..df4756a5 100644 --- a/packages/engine/src/db/schema.ts +++ b/packages/engine/src/db/schema.ts @@ -67,6 +67,10 @@ export const workspaces = sqliteTable( createdAt: integer('created_at', { mode: 'timestamp' }).notNull().default(sql`(unixepoch())`), metadata: text('metadata', { mode: 'json' }).$type>().default({}), retention: text('retention', { mode: 'json' }).$type(), + // Host-owned per-workspace limit overrides (0076). The engine never + // interprets them: the host's entitlements provider maps them, with + // whatever else it knows about the workspace, to the engine's `Limits`. + limitOverrides: text('limit_overrides', { mode: 'json' }).$type>(), // A non-null deadline is an explicit, immutable opt-in to whole-workspace // deletion. Persistent workspaces always leave this null. expiresAt: integer('expires_at', { mode: 'timestamp' }), diff --git a/packages/engine/src/engine.ts b/packages/engine/src/engine.ts index 9b7e07a5..cef64eed 100644 --- a/packages/engine/src/engine.ts +++ b/packages/engine/src/engine.ts @@ -1,4 +1,5 @@ import { safeTelemetrySink } from './ports/telemetry.js'; +import { planLimitExceeded, planLimitExceededBody } from './engine/planCapacity.js'; import { NoopTelemetrySink } from './providers/noop-telemetry.js'; import { Hono } from 'hono'; import type { ContentfulStatusCode } from 'hono/utils/http-status'; @@ -255,6 +256,8 @@ export function createEngine(deps: EngineDeps): Hono { if (error.code === 'invalid_json' || err instanceof SyntaxError) { return jsonMalformedBody(c); } + const refusal = planLimitExceeded(error); + if (refusal) return c.json(planLimitExceededBody(error, refusal), 429); const requestId = c.get('requestId'); return c.json( { diff --git a/packages/engine/src/engine/__tests__/workspaceDeliveryDepth.test.ts b/packages/engine/src/engine/__tests__/workspaceDeliveryDepth.test.ts index 15103824..bea2b8d9 100644 --- a/packages/engine/src/engine/__tests__/workspaceDeliveryDepth.test.ts +++ b/packages/engine/src/engine/__tests__/workspaceDeliveryDepth.test.ts @@ -173,18 +173,18 @@ describe('workspace delivery growth guard (channel broadcast)', () => { it('resolves the dynamic host policy through the async resolver, clamped within cap', async () => { const config = { workspaceDelivery: { - resolve: async (w: { id: string; plan: string }) => - ({ cap: w.plan === 'enterprise' ? 5000 : 500, reserve: 16 }), + resolve: async (w: { id: string; plan: string | null }) => + ({ cap: w.plan === 'tier-large' ? 5000 : 500, reserve: 16 }), }, }; - expect(await resolveWorkspaceDeliveryPolicyFor(config, { id: 'ws', plan: 'enterprise' })) + expect(await resolveWorkspaceDeliveryPolicyFor(config, { id: 'ws', plan: 'tier-large' })) .toEqual({ cap: 5000, reserve: 16 }); - expect(await resolveWorkspaceDeliveryPolicyFor(config, { id: 'ws', plan: 'free' })) + expect(await resolveWorkspaceDeliveryPolicyFor(config, { id: 'ws', plan: 'tier-small' })) .toEqual({ cap: 500, reserve: 16 }); // No host policy => undefined (self-host has no workspace guard). - expect(await resolveWorkspaceDeliveryPolicyFor(undefined, { id: 'ws', plan: 'free' })).toBeUndefined(); + expect(await resolveWorkspaceDeliveryPolicyFor(undefined, { id: 'ws', plan: 'tier-small' })).toBeUndefined(); // A reserve >= cap is clamped into `0 <= reserve < cap`. - expect(await resolveWorkspaceDeliveryPolicyFor({ workspaceDelivery: { cap: 100, reserve: 200 } }, { id: 'ws', plan: 'free' })) + expect(await resolveWorkspaceDeliveryPolicyFor({ workspaceDelivery: { cap: 100, reserve: 200 } }, { id: 'ws', plan: 'tier-small' })) .toEqual({ cap: 100, reserve: 99 }); }); diff --git a/packages/engine/src/engine/a2a.ts b/packages/engine/src/engine/a2a.ts index bc96f4ca..150195d2 100644 --- a/packages/engine/src/engine/a2a.ts +++ b/packages/engine/src/engine/a2a.ts @@ -26,6 +26,7 @@ import type { FileAttachment } from '@relaycast/types'; import type { getDb } from '../db/index.js'; import { a2aAgents, agents } from '../db/schema.js'; import { registerAgent, getAgentByName, updateAgent, updateAgentById } from './agent.js'; +import type { RegistrationCapacityGuard } from './planCapacity.js'; import { rotateAgentToken } from './tokenRotate.js'; import { createAndRunCertification } from './certify.js'; import { isSafeExternalUrl } from '../lib/ssrf.js'; @@ -90,6 +91,8 @@ export interface RegisterA2aAgentInput { authCredential?: string; /** Remote skill/agent that DMs to this proxy should address. */ targetAgent?: string; + /** Registration-cap guard, consulted only when a NEW proxy identity is provisioned. */ + capacity?: RegistrationCapacityGuard; } export interface RegisterA2aAgentResult { @@ -342,6 +345,7 @@ export async function registerA2aAgent( } else { const provisioned = await registerAgent(db, workspaceId, { name: relayName, + capacity: input.capacity, type: 'external', persona: `External A2A agent proxy for ${agentCard.name}`, metadata: proxyMetadata, diff --git a/packages/engine/src/engine/agent.ts b/packages/engine/src/engine/agent.ts index 40413c10..656d8776 100644 --- a/packages/engine/src/engine/agent.ts +++ b/packages/engine/src/engine/agent.ts @@ -6,6 +6,7 @@ import { generateId } from './snowflake.js'; import { invalidateChannelCache } from './cache.js'; import { queryInChunks } from '../lib/queryChunks.js'; import { codedError } from '../lib/httpError.js'; +import type { RegistrationCapacityGuard } from './planCapacity.js'; import { directNodeIdForAgent } from './node.js'; import { addressNodeSelection, formatAgentAddress } from './address.js'; import { runAtomicWrites, type AtomicWrite } from '../ports/database.js'; @@ -216,6 +217,8 @@ export async function registerAgent( autoJoinGeneral?: boolean; recoveryProofHash?: string; workUnitId?: string; + /** Registration-cap guard, consulted only when `name` is NEW to the workspace. */ + capacity?: RegistrationCapacityGuard; }, /** * Internal. Set only by the verifier-reclaim retry below, so that retry can @@ -242,6 +245,16 @@ export async function registerAgent( 400, ); } + if (data.capacity && !afterVerifierReclaim) { + // Only a NEW identity counts against the cap. An established name is + // answered by the unique-constraint conflict below — never by the cap — + // so a register-or-rotate client at its cap still reaches its rotation. + const [established] = await db + .select({ id: agents.id }) + .from(agents) + .where(and(eq(agents.workspaceId, workspaceId), eq(agents.name, data.name))); + if (!established) await data.capacity((data.type ?? 'agent') === 'human' ? 'humans' : 'agents'); + } const agentId = generateId(); const directMachineId = data.directMachinePrefix ? `${data.directMachinePrefix}-${agentId}` : null; const token = `at_live_${randomHex(16)}`; diff --git a/packages/engine/src/engine/node.ts b/packages/engine/src/engine/node.ts index 31f48111..4748b813 100644 --- a/packages/engine/src/engine/node.ts +++ b/packages/engine/src/engine/node.ts @@ -24,6 +24,7 @@ import type { getDb } from '../db/index.js'; import { actionInvocations, agents, agentNodeBindings, channelMembers, channels, nodeProviders, nodes } from '../db/schema.js'; import { randomHex, sha256Hex } from '../lib/crypto.js'; import { codedError } from '../lib/httpError.js'; +import { planLimitExceeded, type RegistrationCapacityGuard } from './planCapacity.js'; import { acceptTaskInvocation, completeTaskInvocation } from './taskInvocation.js'; import { databaseConstraintKind, runAtomic, runAtomicWrites } from '../ports/database.js'; import { retryD1Write } from '../lib/d1Retry.js'; @@ -88,6 +89,11 @@ export interface HandleNodeControlMessageArgs { * caller drives node control without a registry-managed socket. */ connectionId?: string; + /** + * Registration-cap guard for `agent.register`. Absent when the host + * enforces no caps (or drives node control without entitlements). + */ + capacity?: RegistrationCapacityGuard; raw: string; } @@ -101,8 +107,11 @@ function resolveProviderIdentity( type CapabilityLike = string | FleetCapability; +/** Id prefix of the implicit direct node minted for every agent; reserved. */ +export const IMPLICIT_DIRECT_NODE_ID_PREFIX = 'node_direct_'; + export function directNodeIdForAgent(agentId: string): string { - return `node_direct_${agentId}`; + return `${IMPLICIT_DIRECT_NODE_ID_PREFIX}${agentId}`; } function directNodeNameForAgent(agentId: string): string { @@ -439,6 +448,8 @@ export async function createNodeToken( idempotencyKeyHash: string; requestDigest: string; }; + /** Registration-cap guard, consulted only when this enrolls a NEW device. */ + capacity?: RegistrationCapacityGuard; } = {}, ) { const token = authorization.rotationRecovery?.replacementToken ?? `nt_live_${randomHex(24)}`; @@ -450,6 +461,20 @@ export async function createNodeToken( if (!name || name.startsWith('#')) { throw codedError('Node name must be non-empty and cannot begin with "#"', 'invalid_node_name', 400); } + // `node_direct_` is the engine-minted implicit direct node, which + // the device count excludes. The namespace is reserved so a host cannot + // enroll a node that is miscounted as implicit. + if ( + data.node_id !== undefined + && authorization.directAgentId === undefined + && data.node_id.startsWith(IMPLICIT_DIRECT_NODE_ID_PREFIX) + ) { + throw codedError( + `Node ids beginning with "${IMPLICIT_DIRECT_NODE_ID_PREFIX}" are reserved for implicit direct nodes`, + 'reserved_node_id', + 400, + ); + } // Enrollment keys on node_id (the stable identity) when supplied; the name // lookup is only a fallback for callers without a stable id. A name held by // a *different* node id is a conflict — mirroring node.register — never a @@ -484,6 +509,13 @@ export async function createNodeToken( if (role === 'direct' && maxAgents !== 1) { throw codedError('Direct nodes can bind at most one agent', 'direct_node_capacity_exceeded', 400); } + // Only a NEW device counts against the cap. Rotating or re-enrolling an + // established row never does (its role may change, but it was counted when + // it was enrolled), and an agent's implicit direct node — minted on its + // behalf, never enrolled by a host — is its own WebSocket, not a device. + if (!existing && authorization.directAgentId === undefined && authorization.capacity) { + await authorization.capacity('devices'); + } if (existing) { const hasPresentedTokenProof = authorization.expectedTokenHash !== undefined; @@ -2116,10 +2148,19 @@ export async function registerAgentViaNode( nodeId: string, providerName: string, message: FleetAgentRegisterMessage, - options: { deliveryCursorSupported?: boolean } = {}, + options: { deliveryCursorSupported?: boolean; capacity?: RegistrationCapacityGuard } = {}, ): Promise { assertRegistrableAgentName(message.name); const wantsAutoJoin = message.auto_join_general !== false; + if (options.capacity) { + // Only a NEW identity counts against the cap: an established name is + // answered by the conflict below, never by the cap. + const [established] = await db + .select({ id: agents.id }) + .from(agents) + .where(and(eq(agents.workspaceId, workspaceId), eq(agents.name, message.name))); + if (!established) await options.capacity('agents'); + } const registered = await runAtomic(db, async (tx) => { // Every awaited statement is a separate D1 round-trip — and this path runs // under the node-control liveness deadline whose misses surface as @@ -3551,7 +3592,10 @@ export async function handleNodeControlMessage(args: HandleNodeControlMessageArg args.nodeId, frameProviderName, message, - { deliveryCursorSupported: supportsProviderDeliveryReadiness(args.registry) }, + { + deliveryCursorSupported: supportsProviderDeliveryReadiness(args.registry), + capacity: args.capacity, + }, ); // The relay broker's node_control client awaits a `reply` frame keyed by // the request id (it matches `pending_agent_registrations` by `reply.id` @@ -3767,13 +3811,19 @@ export async function handleNodeControlMessage(args: HandleNodeControlMessageArg provider: frameProviderName, message: error.message, }); + // A cap refusal carries the same machine-readable detail as the HTTP 429 + // (`error_code`, `limit`, `max`, `current`) plus the host's context. + const refusal = planLimitExceeded(err); sendControl(args.socket, { + // Host context first: the frame's own fields always win over it. + ...(refusal?.context ?? {}), v: 1, id: requestId(message), type: 'error', ok: false, code, message: error.message, + ...(refusal ? { error_code: code, ...refusal.detail } : {}), }); } }); diff --git a/packages/engine/src/engine/planCapacity.ts b/packages/engine/src/engine/planCapacity.ts new file mode 100644 index 00000000..eb7bed50 --- /dev/null +++ b/packages/engine/src/engine/planCapacity.ts @@ -0,0 +1,186 @@ +import { and, eq, ne, not, sql } from 'drizzle-orm'; +import type { EngineDb } from '../ports/database.js'; +import type { EntitlementsProvider, RegistrationMetric } from '../ports/entitlements.js'; +import type { Workspace } from '../ports/auth.js'; +import { agents, nodes, workspaces } from '../db/schema.js'; +import type { CodedError } from '../lib/httpError.js'; + +export const PLAN_LIMIT_EXCEEDED = 'plan_limit_exceeded'; + +/** The machine-readable detail a registration-cap refusal carries. */ +export interface PlanLimitExceededDetail { + limit: RegistrationMetric; + max: number; + current: number; +} + +export interface PlanLimitExceededError extends CodedError { + code: typeof PLAN_LIMIT_EXCEEDED; + status: 429; + planLimit: PlanLimitExceededDetail; + /** Host-owned context from `EntitlementsProvider.getLimitContext`, passed through verbatim. */ + limitContext: Record; +} + +const REGISTRATION_NOUN: Record = { + devices: 'device', + agents: 'agent', + humans: 'human member', +}; + +export function planLimitExceededError( + detail: PlanLimitExceededDetail, + limitContext: Record = {}, +): PlanLimitExceededError { + const message = + `Limit exceeded for ${detail.limit}: this workspace allows ${detail.max} ` + + `${REGISTRATION_NOUN[detail.limit]}${detail.max === 1 ? '' : 's'} and ${detail.current} ` + + `${detail.current === 1 ? 'is' : 'are'} registered.`; + return Object.assign(new Error(message), { + code: PLAN_LIMIT_EXCEEDED, + status: 429, + clientSafe: true as const, + planLimit: detail, + limitContext, + }) as PlanLimitExceededError; +} + +/** The cap detail and host context when `err` is a registration-cap refusal, else `null`. */ +export function planLimitExceeded( + err: unknown, +): { detail: PlanLimitExceededDetail; context: Record } | null { + if (!err || typeof err !== 'object') return null; + const candidate = err as Partial; + if (candidate.code !== PLAN_LIMIT_EXCEEDED) return null; + const detail = candidate.planLimit; + if (!detail || typeof detail !== 'object') return null; + if (typeof detail.limit !== 'string' || typeof detail.max !== 'number' || typeof detail.current !== 'number') { + return null; + } + const context = candidate.limitContext; + return { + detail: { limit: detail.limit, max: detail.max, current: detail.current }, + context: context && typeof context === 'object' && !Array.isArray(context) ? context : {}, + }; +} + +/** + * The HTTP 429 body for a registration-cap refusal: the standard error + * envelope plus top-level `error_code`/`limit`/`max`/`current` and the + * provider's context spread at top level (so a hosted gateway's `plan` and + * `upgrade_url` land where clients key on them). The fixed fields win over a + * context key of the same name. + */ +export function planLimitExceededBody( + error: CodedError, + refusal: { detail: PlanLimitExceededDetail; context: Record }, +) { + return { + ...refusal.context, + ok: false as const, + error: { code: PLAN_LIMIT_EXCEEDED, message: error.message }, + error_code: PLAN_LIMIT_EXCEEDED, + ...refusal.detail, + }; +} + +/** A finite non-negative cap, or `null` for unlimited / unset / malformed. */ +export function registrationCap(value: unknown): number | null { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 ? Math.floor(value) : null; +} + +/** + * How many of `metric` the workspace currently has registered — the number a + * new registration is compared against. + * + * - `devices`: every node a host enrolled, whatever its role or transport. + * Every agent also owns an implicit direct node (`node_direct_`) + * for its own WebSocket; those are not devices. The prefix is reserved at + * enrollment (`createNodeToken`), so a host cannot enroll into it. + * - `agents`: non-human identities that have not been released. + * - `humans`: `type = 'human'` identities that have not been released. + */ +export async function countRegistrations( + db: EngineDb, + workspaceId: string, + metric: RegistrationMetric, +): Promise { + const count = sql`count(*)`; + const rows = metric === 'devices' + ? await db + .select({ count }) + .from(nodes) + .where(and( + eq(nodes.workspaceId, workspaceId), + // GLOB, not LIKE: in LIKE `_` is a wildcard, so `nodeXdirectX…` would + // slip out of the count while passing the literal reserved-prefix check. + not(and(eq(nodes.role, 'direct'), sql`${nodes.id} GLOB 'node_direct_*'`)!), + )) + : await db + .select({ count }) + .from(agents) + .where(and( + eq(agents.workspaceId, workspaceId), + // Tombstoned identities (`RELEASED_AGENT_STATUS` in agent.ts) hold no slot. + ne(agents.status, 'released'), + metric === 'humans' ? eq(agents.type, 'human') : ne(agents.type, 'human'), + )); + return Number(rows[0]?.count ?? 0); +} + +/** + * Refuses a NEW registration of `metric` when the workspace is at its cap. + * Callers invoke it only on the create path — re-registering an existing + * node or agent never counts as new. + */ +export type RegistrationCapacityGuard = (metric: RegistrationMetric) => Promise; + +export interface RegistrationCapacityArgs { + db: EngineDb; + entitlements: EntitlementsProvider; + /** The authenticated workspace row, or its id when the caller only has that. */ + workspace: Workspace | string; +} + +async function resolveWorkspace(db: EngineDb, workspace: Workspace | string): Promise { + if (typeof workspace !== 'string') return workspace; + const [row] = await db.select().from(workspaces).where(eq(workspaces.id, workspace)); + return row ?? null; +} + +/** + * Build the guard the registration paths run before inserting a node or agent. + * The provider decides the cap (`getLimits`), the current count (`getUsage`) + * and the refusal's host context (`getLimitContext`); a provider failure fails + * open with a warning, matching the api_calls quota middleware — a billing + * outage must not block registration. + */ +export function createRegistrationCapacityGuard(args: RegistrationCapacityArgs): RegistrationCapacityGuard { + return async (metric) => { + let detail: PlanLimitExceededDetail | null = null; + let context: Record = {}; + try { + const workspace = await resolveWorkspace(args.db, args.workspace); + if (!workspace) return; + const limits = await args.entitlements.getLimits(workspace); + const max = registrationCap(limits[metric]); + if (max === null) return; + const current = await args.entitlements.getUsage(workspace.id, metric); + if (current < max) return; + detail = { limit: metric, max, current }; + try { + context = (await args.entitlements.getLimitContext?.(workspace)) ?? {}; + } catch { + // The refusal stands without the host's extra fields. + } + } catch (err) { + console.warn('[plan-capacity] entitlements lookup failed; allowing registration', { + metric, + workspaceId: typeof args.workspace === 'string' ? args.workspace : args.workspace.id, + error: err instanceof Error ? err.message : String(err), + }); + return; + } + throw planLimitExceededError(detail, context); + }; +} diff --git a/packages/engine/src/engine/workspaceDeliveryPolicy.ts b/packages/engine/src/engine/workspaceDeliveryPolicy.ts index a7792288..dd24a927 100644 --- a/packages/engine/src/engine/workspaceDeliveryPolicy.ts +++ b/packages/engine/src/engine/workspaceDeliveryPolicy.ts @@ -51,7 +51,7 @@ export async function resolveWorkspaceDeliveryPolicyFor( ): Promise { const dynamic = await config?.workspaceDelivery?.resolve?.({ id: workspace.id, - plan: workspace.plan ?? 'free', + plan: workspace.plan ?? null, }); if (dynamic !== undefined) return dynamicWorkspaceDeliveryPolicySchema.parse(dynamic); return resolveWorkspaceDeliveryPolicy(config, workspace.id); diff --git a/packages/engine/src/index.ts b/packages/engine/src/index.ts index a31718b5..911d641f 100644 --- a/packages/engine/src/index.ts +++ b/packages/engine/src/index.ts @@ -30,7 +30,7 @@ export type { Workspace, Agent, EntitlementsProvider, - PlanLimits, + Limits, UsageMetric, TelemetrySink, TelemetryEvent, @@ -43,7 +43,20 @@ export { observerAllowsEvent, type ObserverToken } from './engine/observerToken. // OSS default providers (self-host). Cloud injects its own. export { SqliteApiKeyAuthProvider, hashToken } from './auth/index.js'; -export { StaticEntitlementsProvider, PLAN_LIMITS } from './providers/static-entitlements.js'; +export { StaticEntitlementsProvider } from './providers/static-entitlements.js'; +export { UNLIMITED_LIMITS, resolveLimits } from './ports/entitlements.js'; +// Registration caps (devices/agents/humans): the live counts a provider's +// `getUsage` reports, the cap normalizer, and the guard a host wires into its +// node-control transport so `agent.register` frames are capped like HTTP. +export { + countRegistrations, + createRegistrationCapacityGuard, + planLimitExceeded, + registrationCap, + PLAN_LIMIT_EXCEEDED, + type PlanLimitExceededDetail, + type RegistrationCapacityGuard, +} from './engine/planCapacity.js'; // Usage-counter key contract. An `EntitlementsProvider` reads the counters the // engine writes, so it has to agree with the engine on the billing period the // key is scoped to — build the key with these rather than by hand. diff --git a/packages/engine/src/lib/httpError.ts b/packages/engine/src/lib/httpError.ts index a1384cf9..08725b78 100644 --- a/packages/engine/src/lib/httpError.ts +++ b/packages/engine/src/lib/httpError.ts @@ -1,6 +1,7 @@ import type { Context } from 'hono'; import type { ContentfulStatusCode } from 'hono/utils/http-status'; import { jsonError, jsonMalformedBody } from './httpResponse.js'; +import { planLimitExceeded, planLimitExceededBody } from '../engine/planCapacity.js'; /** * Errors thrown inside the engine optionally carry a stable error `code` and an @@ -91,6 +92,10 @@ export function errorResponse(c: Context, err: unknown) { if (err instanceof SyntaxError) { return jsonMalformedBody(c); } + // A registration-cap refusal carries machine-readable detail beyond the + // envelope (`error_code`, `limit`, `max`, `current`) plus the host's context. + const refusal = planLimitExceeded(error); + if (refusal) return c.json(planLimitExceededBody(error, refusal), 429); const status = effectiveHttpStatus(error); // Throw the normalized Error: Hono hands only `Error` instances to `onError`. if (!error.code && status >= 500) throw error; diff --git a/packages/engine/src/middleware/planLimits.ts b/packages/engine/src/middleware/planLimits.ts index 7fd81dd5..a6571d72 100644 --- a/packages/engine/src/middleware/planLimits.ts +++ b/packages/engine/src/middleware/planLimits.ts @@ -1,13 +1,12 @@ import { createMiddleware } from 'hono/factory'; +import type { Context } from 'hono'; import type { AppEnv } from '../env.js'; +import { createRegistrationCapacityGuard, type RegistrationCapacityGuard } from '../engine/planCapacity.js'; import type { EntitlementsProvider, UsageMetric } from '../ports/entitlements.js'; import { usagePeriodResetAt } from '../engine/usage.js'; import { jsonError } from '../lib/httpResponse.js'; import { setRetryContract } from '../lib/throttle.js'; -// Re-exported for back-compat with callers/tests that imported the table here. -export { PLAN_LIMITS } from '../providers/static-entitlements.js'; - /** * A provider billing on its own cycle reports that cycle's end. An absent or * failing hook falls back to the engine's UTC-month period, so how the reset is @@ -39,21 +38,33 @@ export function checkPlanLimit(metric: UsageMetric) { return; } const limit = limits[metric]; - if (typeof limit !== 'number' || limit === Infinity) { await next(); return; } + if (typeof limit !== 'number' || !Number.isFinite(limit)) { await next(); return; } try { const current = await entitlements.getUsage(workspace.id, metric); if (current >= limit) { - const plan = workspace.plan || 'free'; // Quota clears on the billing period, not on a backoff. Advertising the // real reset lets a caller with a bounded retry budget fail fast instead // of spending it on a condition retrying cannot fix. const resetAt = await resolveUsageResetAt(entitlements, workspace.id, metric); setRetryContract(c, resetAt); - return jsonError(c, 'plan_limit_exceeded', `Plan limit exceeded for ${metric}. Current plan: ${plan}. Limit: ${limit} per usage period. Current usage: ${current}. Quota resets at ${new Date(resetAt).toISOString()}. Upgrade your plan to increase limits.`, 429); + return jsonError(c, 'plan_limit_exceeded', `Limit exceeded for ${metric}: ${limit} per usage period allowed, ${current} used. Quota resets at ${new Date(resetAt).toISOString()}.`, 429); } } catch { /* fail open */ } await next(); }); } + +/** + * The registration-cap guard for the authenticated workspace on this + * request. Registration routes hand it to the engine, which runs it only when + * the request would create a NEW node or agent. + */ +export function registrationCapacityGuard(c: Context): RegistrationCapacityGuard { + return createRegistrationCapacityGuard({ + db: c.get('db'), + entitlements: c.get('engine').entitlements, + workspace: c.get('workspace'), + }); +} diff --git a/packages/engine/src/middleware/rateLimit.ts b/packages/engine/src/middleware/rateLimit.ts index d2214979..f2a279a0 100644 --- a/packages/engine/src/middleware/rateLimit.ts +++ b/packages/engine/src/middleware/rateLimit.ts @@ -109,12 +109,17 @@ export const rateLimit = createMiddleware(async (c, next) => { // Resolve the per-minute limit. If entitlements are unavailable, fall back to // a conservative default and still enforce it — an entitlements outage must // NOT translate into unlimited throughput. - let globalLimit: number; + let globalLimit: number | null; try { globalLimit = (await entitlements.getLimits(workspace)).rate_per_min; } catch { globalLimit = FALLBACK_RATE_PER_MIN; } + // A null ceiling is unlimited: the provider chose not to throttle this workspace. + if (globalLimit === null || !Number.isFinite(globalLimit)) { + await presenceRefresh(c, async () => {}); + return usageTracker(c, next); + } const limit = routeKey ? Math.ceil(globalLimit * ROUTE_MULTIPLIERS[routeKey]) : globalLimit; try { @@ -128,7 +133,7 @@ export const rateLimit = createMiddleware(async (c, next) => { // Transient by construction: the bucket is keyed to this minute, so the // advertised wait is the real one. setRetryContract(c, resetAt, now); - return jsonError(c, 'rate_limit_exceeded', `Rate limit exceeded. ${limit} requests per minute allowed for ${workspace.plan} plan.`, 429); + return jsonError(c, 'rate_limit_exceeded', `Rate limit exceeded. ${limit} requests per minute allowed.`, 429); } } catch { // Only the limiter backend failing is fail-open — a transient infra hiccup diff --git a/packages/engine/src/ports/entitlements.ts b/packages/engine/src/ports/entitlements.ts index 43050b37..b89f59ff 100644 --- a/packages/engine/src/ports/entitlements.ts +++ b/packages/engine/src/ports/entitlements.ts @@ -1,31 +1,76 @@ import type { Workspace } from './auth.js'; -export type UsageMetric = 'messages' | 'agents' | 'file_bytes' | 'api_calls'; - -export interface PlanLimits { - messages: number; - agents: number; - file_bytes: number; - /** Total authenticated API calls allowed for the current billing window. */ - api_calls: number; - /** Global requests-per-minute ceiling, fed into the rate limiter. */ - rate_per_min: number; +/** + * `messages`/`file_bytes`/`api_calls` are metered per billing period; + * `devices`/`agents`/`humans` are registration counts — how many of each a + * workspace currently has enrolled — compared against the hard cap when a new + * one is registered. + */ +export type UsageMetric = 'messages' | 'agents' | 'file_bytes' | 'api_calls' | 'devices' | 'humans'; + +/** The metrics enforced as hard registration caps. */ +export type RegistrationMetric = 'devices' | 'agents' | 'humans'; + +/** + * A workspace's effective limits, keyed by metric. `null` is unlimited. The + * engine carries no tiers or numbers of its own: a host's entitlements + * provider maps whatever it knows about a workspace (billing plan, stored + * overrides, operator config) to this object. + */ +export interface Limits { + /** Messages per usage period. */ + messages: number | null; + /** Registered non-human agents (hard cap). */ + agents: number | null; + /** Stored file bytes. */ + file_bytes: number | null; + /** Authenticated API calls per usage period. */ + api_calls: number | null; + /** Requests per minute, fed into the rate limiter. */ + rate_per_min: number | null; + /** Registered devices: nodes a host enrolled, excluding each agent's implicit direct node (hard cap). */ + devices: number | null; + /** Registered human members (`agents.type = 'human'`, hard cap). */ + humans: number | null; + /** Message retention in days the host applies for this workspace; `null` = never prune. */ + message_retention_days: number | null; +} + +/** Every metric unlimited — the self-hosted default. */ +export const UNLIMITED_LIMITS: Readonly = Object.freeze({ + messages: null, + agents: null, + file_bytes: null, + api_calls: null, + rate_per_min: null, + devices: null, + humans: null, + message_retention_days: null, +}); + +/** A full {@link Limits} object from a partial one; absent keys are unlimited. */ +export function resolveLimits(limits?: Partial | null): Limits { + return { ...UNLIMITED_LIMITS, ...(limits ?? {}) }; } /** - * Pluggable plan limits / usage — the billing seam. + * Pluggable limits / usage — the hosting seam. * - * Self-host uses `StaticEntitlementsProvider`, seeded from today's hardcoded - * `PLAN_LIMITS` (free/connect/pro/enterprise), defaulting single-tenant deployments to a - * generous/unlimited tier. The cloud product injects a provider that resolves - * quotas from its billing system (e.g. a Stripe subscription) per workspace and - * reports metered usage. + * Self-host uses `StaticEntitlementsProvider`: unlimited unless the operator + * sets `EngineConfig.limits` for their own instance. The hosted product injects + * a provider that resolves limits from its billing system per workspace, + * reports metered usage, and supplies the context its clients need to explain + * a refusal. */ export interface EntitlementsProvider { - /** Resolve the active plan limits for a workspace. */ - getLimits(workspace: Workspace): Promise; + /** Resolve the effective limits for the authenticated workspace row. */ + getLimits(workspace: Workspace): Promise; - /** Current usage for a metric, compared against the limit before allowing an action. */ + /** + * Current usage for a metric, compared against the limit before allowing an + * action. For a {@link RegistrationMetric} this is the number currently + * registered (see `countRegistrations`), not a per-period counter. + */ getUsage(workspaceId: string, metric: UsageMetric): Promise; /** @@ -36,4 +81,12 @@ export interface EntitlementsProvider { * UTC-month period. */ getUsageResetAt?(workspaceId: string, metric: UsageMetric): Promise; + + /** + * Host-owned context spread verbatim onto the top level of a + * `plan_limit_exceeded` refusal (for example the workspace's billing tier and + * an upgrade URL). The engine never interprets it. Optional; a failure or + * omission yields no extra fields. + */ + getLimitContext?(workspace: Workspace): Promise>; } diff --git a/packages/engine/src/ports/index.ts b/packages/engine/src/ports/index.ts index bd0d462e..53d2abd3 100644 --- a/packages/engine/src/ports/index.ts +++ b/packages/engine/src/ports/index.ts @@ -43,9 +43,11 @@ export type { } from './auth.js'; export type { EntitlementsProvider, - PlanLimits, + Limits, + RegistrationMetric, UsageMetric, } from './entitlements.js'; +export { UNLIMITED_LIMITS, resolveLimits } from './entitlements.js'; export { ENGINE_OBSERVATION_FACTS } from './telemetry.js'; export type { TelemetrySink, TelemetryEvent, EngineObservation, EngineObservationKind, EngineObservationFacts } from './telemetry.js'; export type { EngineLogRecord, EngineLogExport } from './logging.js'; @@ -58,7 +60,7 @@ import type { FileStorage } from './files.js'; import type { KeyValueStore } from './kv.js'; import type { EventQueue } from './event-queue.js'; import type { AuthProvider } from './auth.js'; -import type { EntitlementsProvider } from './entitlements.js'; +import type { EntitlementsProvider, Limits } from './entitlements.js'; import type { TelemetrySink } from './telemetry.js'; import type { EngineLogExport } from './logging.js'; @@ -94,6 +96,12 @@ export interface EngineDeps extends EnginePorts, EngineProviders { export interface EngineConfig { environment?: string; + /** + * Static limits for every workspace on this instance, applied by the default + * `StaticEntitlementsProvider` (absent keys are unlimited). Ignored when a + * host injects its own entitlements provider. + */ + limits?: Partial; appVersion?: string; appSemver?: string; sdkSemver?: string; diff --git a/packages/engine/src/ports/workspaceDelivery.ts b/packages/engine/src/ports/workspaceDelivery.ts index 0b00761f..98f0ef66 100644 --- a/packages/engine/src/ports/workspaceDelivery.ts +++ b/packages/engine/src/ports/workspaceDelivery.ts @@ -15,11 +15,12 @@ export interface WorkspaceDeliveryPolicyConfig { reserve?: number; workspaces?: Record; /** - * Server-only async resolver for dynamic, request-time plans (e.g. a - * KV-backed effective plan lookup). When present it takes precedence over the - * static `cap`/`workspaces` fields, so a host whose plan is async does not - * have to flatten every workspace into static config. Never source this from - * a client request. + * Server-only async resolver for dynamic, request-time policy (e.g. a + * KV-backed effective tier lookup). When present it takes precedence over the + * static `cap`/`workspaces` fields, so a host whose policy is async does not + * have to flatten every workspace into static config. `plan` is the + * workspace row's host-assigned key, opaque to the engine. Never source this + * from a client request. */ - resolve?: (workspace: { id: string; plan: string }) => Promise; + resolve?: (workspace: { id: string; plan: string | null }) => Promise; } diff --git a/packages/engine/src/providers/__tests__/static-entitlements.test.ts b/packages/engine/src/providers/__tests__/static-entitlements.test.ts index b49c5aa2..2fa2592a 100644 --- a/packages/engine/src/providers/__tests__/static-entitlements.test.ts +++ b/packages/engine/src/providers/__tests__/static-entitlements.test.ts @@ -1,34 +1,54 @@ import { describe, expect, it } from 'vitest'; -import { PLAN_LIMITS, StaticEntitlementsProvider } from '../static-entitlements.js'; -import type { Workspace } from '../../ports/index.js'; +import { StaticEntitlementsProvider } from '../static-entitlements.js'; +import { UNLIMITED_LIMITS, resolveLimits } from '../../ports/entitlements.js'; +import type { EngineDb, Workspace } from '../../ports/index.js'; -function workspace(plan: string): Workspace { - return { id: 'ws_1', plan } as Workspace; -} +const workspace = { id: 'ws_1', plan: 'anything' } as Workspace; describe('StaticEntitlementsProvider', () => { - it('resolves the connect profile for room workspaces', async () => { - const provider = new StaticEntitlementsProvider(undefined, 'free'); - const limits = await provider.getLimits(workspace('connect')); - expect(limits).toBe(PLAN_LIMITS.connect); - expect(limits.api_calls).toBe(10_000_000); - expect(limits.agents).toBe(10); + it('is unlimited by default, whatever the workspace row says', async () => { + const provider = new StaticEntitlementsProvider(); + await expect(provider.getLimits(workspace)).resolves.toEqual(UNLIMITED_LIMITS); + await expect(provider.getLimits({ id: 'ws_2', plan: 'other' } as Workspace)).resolves.toEqual(UNLIMITED_LIMITS); }); - it('sizes connect between free and pro on every metric', () => { - for (const metric of ['messages', 'agents', 'file_bytes', 'api_calls', 'rate_per_min'] as const) { - expect(PLAN_LIMITS.connect![metric]).toBeGreaterThan(PLAN_LIMITS.free![metric]); - } - for (const metric of ['messages', 'agents', 'file_bytes', 'rate_per_min'] as const) { - expect(PLAN_LIMITS.connect![metric]).toBeLessThan(PLAN_LIMITS.pro![metric]); - } - // Rooms poll constantly; connect's call budget is set to pro's for - // headroom, while its other limits stay below pro. - expect(PLAN_LIMITS.connect!.api_calls).toBeLessThanOrEqual(PLAN_LIMITS.pro!.api_calls); + it('applies the operator\'s configured limits over unlimited', async () => { + const provider = new StaticEntitlementsProvider(undefined, { devices: 5, rate_per_min: 600 }, {} as EngineDb); + await expect(provider.getLimits(workspace)).resolves.toEqual({ + ...UNLIMITED_LIMITS, + devices: 5, + rate_per_min: 600, + }); }); - it('falls back to the default plan for an unknown plan', async () => { - const provider = new StaticEntitlementsProvider(undefined, 'free'); - expect(await provider.getLimits(workspace('platinum'))).toBe(PLAN_LIMITS.free); + it('carries no tier names or numbers of its own', () => { + for (const value of Object.values(UNLIMITED_LIMITS)) expect(value).toBeNull(); + expect(Object.keys(UNLIMITED_LIMITS).sort()).toEqual([ + 'agents', 'api_calls', 'devices', 'file_bytes', 'humans', 'message_retention_days', 'messages', 'rate_per_min', + ]); + }); + + it('refuses registration caps it has no database to count against', () => { + expect(() => new StaticEntitlementsProvider(undefined, { devices: 1 })) + .toThrow(/registration caps \(devices\) require a database handle/); + // Metered-only limits need no database. + expect(() => new StaticEntitlementsProvider(undefined, { rate_per_min: 60 })).not.toThrow(); + }); + + it('reports zero metered usage without a key/value store', async () => { + const provider = new StaticEntitlementsProvider(); + await expect(provider.getUsage('ws_1', 'api_calls')).resolves.toBe(0); + }); +}); + +describe('resolveLimits', () => { + it('fills absent keys with unlimited', () => { + expect(resolveLimits()).toEqual(UNLIMITED_LIMITS); + expect(resolveLimits(null)).toEqual(UNLIMITED_LIMITS); + expect(resolveLimits({ agents: 3 })).toEqual({ ...UNLIMITED_LIMITS, agents: 3 }); + }); + + it('keeps an explicit null as unlimited', () => { + expect(resolveLimits({ agents: null, humans: 2 })).toEqual({ ...UNLIMITED_LIMITS, humans: 2 }); }); }); diff --git a/packages/engine/src/providers/static-entitlements.ts b/packages/engine/src/providers/static-entitlements.ts index ab3c6fb4..ae5845f2 100644 --- a/packages/engine/src/providers/static-entitlements.ts +++ b/packages/engine/src/providers/static-entitlements.ts @@ -1,43 +1,49 @@ -import type { EntitlementsProvider, PlanLimits, UsageMetric, Workspace } from '../ports/index.js'; +import type { EntitlementsProvider, Limits, UsageMetric, Workspace } from '../ports/index.js'; +import { resolveLimits } from '../ports/entitlements.js'; +import type { EngineDb } from '../ports/database.js'; import type { KeyValueStore } from '../ports/kv.js'; import { getUsageMetric, usagePeriodResetAt } from '../engine/usage.js'; +import { countRegistrations } from '../engine/planCapacity.js'; -/** - * The default plan limits, lifted from the original `middleware/planLimits.ts`. - * Self-host deployments typically run the `selfhost` tier (everything - * unlimited); the free/connect/pro/enterprise tiers are kept so the same code serves a - * tiered deployment if a workspace's `plan` column is set. - */ -export const PLAN_LIMITS: Record = { - free: { messages: 10_000, agents: 5, file_bytes: 100 * 1024 * 1024, api_calls: 100_000, rate_per_min: 300 }, - // Short-lived collaboration rooms (Relay Connect): a handful of seats that - // poll their inbox continuously, so api_calls dominates while messages stay - // low. Sized for a two-seat room polling every few seconds for a full month. - connect: { messages: 100_000, agents: 10, file_bytes: 1024 * 1024 * 1024, api_calls: 10_000_000, rate_per_min: 1200 }, - pro: { messages: 1_000_000, agents: 100, file_bytes: 50 * 1024 * 1024 * 1024, api_calls: 10_000_000, rate_per_min: 6000 }, - enterprise: { messages: Infinity, agents: Infinity, file_bytes: 500 * 1024 * 1024 * 1024, api_calls: Infinity, rate_per_min: 30000 }, - selfhost: { messages: Infinity, agents: Infinity, file_bytes: Infinity, api_calls: Infinity, rate_per_min: 30000 }, -}; +const REGISTRATION_METRICS = new Set(['devices', 'agents', 'humans']); +const REGISTRATION_CAPS = ['devices', 'agents', 'humans'] as const; /** - * Static, table-driven entitlements — the self-host default. Resolves limits - * from {@link PLAN_LIMITS} keyed by `workspace.plan`, and reads usage counters - * from the key/value store (the same `usage:::` keys the - * usage tracker increments, so both agree on the current billing period). The - * cloud product injects a billing-backed provider. + * Static entitlements — the self-host default. Every workspace gets the same + * limits: unlimited unless the operator sets some through + * `EngineConfig.limits`. Metered usage is read from the key/value store (the + * same `usage:::` keys the usage tracker increments, so + * both agree on the current billing period); registration metrics are counted + * live from the database. The hosted product injects its own provider. */ export class StaticEntitlementsProvider implements EntitlementsProvider { + private readonly limits: Limits; + constructor( private readonly kv?: KeyValueStore, - private readonly defaultPlan: string = 'selfhost', - ) {} + limits?: Partial | null, + private readonly db?: EngineDb, + ) { + this.limits = resolveLimits(limits); + // A cap the provider cannot count against would admit everything while + // looking configured; refuse the misconfiguration up front. + const uncountable = REGISTRATION_CAPS.filter((metric) => typeof this.limits[metric] === 'number'); + if (uncountable.length > 0 && !db) { + throw new Error( + `StaticEntitlementsProvider: registration caps (${uncountable.join(', ')}) require a database handle to count against`, + ); + } + } - async getLimits(workspace: Workspace): Promise { - const plan = workspace.plan || this.defaultPlan; - return PLAN_LIMITS[plan] ?? PLAN_LIMITS[this.defaultPlan]; + async getLimits(_workspace: Workspace): Promise { + return this.limits; } async getUsage(workspaceId: string, metric: UsageMetric): Promise { + if (REGISTRATION_METRICS.has(metric)) { + if (!this.db) return 0; + return countRegistrations(this.db, workspaceId, metric as 'devices' | 'agents' | 'humans'); + } if (!this.kv) return 0; try { return await getUsageMetric(this.kv, workspaceId, metric); diff --git a/packages/engine/src/routes/a2a.ts b/packages/engine/src/routes/a2a.ts index b2b582b1..6ea36077 100644 --- a/packages/engine/src/routes/a2a.ts +++ b/packages/engine/src/routes/a2a.ts @@ -15,6 +15,7 @@ import { import { getWorkspaceExpiry } from '../engine/workspace.js'; import { asCodedError, codedError, errorResponse, type CodedError } from '../lib/httpError.js'; import { rateLimit } from '../middleware/rateLimit.js'; +import { registrationCapacityGuard } from '../middleware/planLimits.js'; import { buildIdempotencyStorageKey, runIdempotent } from '../middleware/idempotency.js'; import { sha256Hex } from '../lib/crypto.js'; import { runInBackground } from './background.js'; @@ -264,6 +265,7 @@ a2aRoutes.post('/v1/a2a/register', requireAuth, rateLimit, async (c) => { authScheme: parsed.data.auth_scheme, authCredential: parsed.data.auth_credential, targetAgent: parsed.data.target_agent, + capacity: registrationCapacityGuard(c), }); return jsonCreated(c, { diff --git a/packages/engine/src/routes/agent.ts b/packages/engine/src/routes/agent.ts index 442f5d30..fb124102 100644 --- a/packages/engine/src/routes/agent.ts +++ b/packages/engine/src/routes/agent.ts @@ -12,6 +12,7 @@ import { } from '../auth/workspaceExpiry.js'; import { getWorkspaceExpiry } from '../engine/workspace.js'; import { rateLimit } from '../middleware/rateLimit.js'; +import { registrationCapacityGuard } from '../middleware/planLimits.js'; import * as agentEngine from '../engine/agent.js'; import { agentRetentionSchema, retainAgents } from '../engine/agentRetention.js'; import * as agentIdentityEngine from '../engine/agentIdentity.js'; @@ -314,6 +315,7 @@ agentRoutes.post( autoJoinGeneral, recoveryProofHash, workUnitId, + capacity: registrationCapacityGuard(c), }); if (skills?.length) { await directoryEngine.syncSourceAgentDirectoryEntry(db, workspace.id, { diff --git a/packages/engine/src/routes/node.ts b/packages/engine/src/routes/node.ts index bccc3151..b8ac6dea 100644 --- a/packages/engine/src/routes/node.ts +++ b/packages/engine/src/routes/node.ts @@ -5,6 +5,7 @@ import type { AppEnv } from '../env.js'; import { requireAuth, requireNodeToken, requireWorkspaceRead, requireWorkspaceKey } from '../middleware/auth.js'; import { rateLimit } from '../middleware/rateLimit.js'; import { parseIdempotencyKey } from '../middleware/idempotency.js'; +import { registrationCapacityGuard } from '../middleware/planLimits.js'; import { asCodedError, errorResponse } from '../lib/httpError.js'; import { isSafeExternalUrl } from '../lib/ssrf.js'; import { sha256Hex } from '../lib/crypto.js'; @@ -236,7 +237,7 @@ async function enrollNode(c: Context, data: z.infer