diff --git a/README.md b/README.md index c84750f..28e6fc1 100644 --- a/README.md +++ b/README.md @@ -37,6 +37,7 @@ app.route('/v1/agents', createAgentGateway({ authorizeConsumer: authorizeAgentAccess, recordUsage: usageStore.recordUsage, claimApiKeyRequest: createApiKeyRequestClaim(apiKeyStore), + apiKeyReservationLifecycle: apiKeyStore.reservations, settlePayment: createApiKeyUsageSettlement(apiKeyStore), x402: { operatorAddress: '0x…', @@ -82,8 +83,28 @@ Concurrent workers cannot claim the same slot, and a retry with the same request When `verifyApiKey` returns a minute or daily limit, configure `claimApiKeyRequest` or the request fails closed with `503`. The SQL store retains the current and previous UTC day, then prunes older claim rows every 256 accepted requests. The API-key store also records each usage settlement once and refuses a settlement that would exceed the spending limit. -This check runs after work completes, so it does not reserve funds before an in-flight request. -Use a payment authorization flow when the product requires a strict pre-run budget. +Finite-cap keys also reserve the full quoted customer charge before execution. +The quote uses the greater of service-token charges and the provider-cost ceiling, matching settlement pricing. +Reservations use whole cents; provider budgets remain USD, and token limits remain token counts. +Concurrent workers cannot reserve more than the remaining key cap. +Settlement records actual charges once and releases the unused portion of the matching reservation. +Only a pre-execution reservation can be released after failure. +An execution without a final receipt retains its reservation until an authoritative receipt reconciles it. +Reservation rows are separate from rate counters and never expire through rate-counter pruning. + +Existing installations must add the reservation table from `sqlApiKeyStoreSchemaStatements()` before adopting this version. +Its default name is `${table}_reservation`; `reservationTable` overrides it for custom schemas. +Wire `apiKeyReservationLifecycle: apiKeyStore.reservations` alongside the claim and settlement callbacks. +Do not delete executing reservations to recover capacity without first reconciling the underlying work. + +Capped execution requires `SandboxBox.prepareBudgetedPrompt`. +Preparation must start no compute and return either an unsupported result or a prepared stream that enforces every supplied limit. +Enforcement includes provider calls, retries, child calls, and tools throughout that stream. +The gateway uses only this prepared stream for capped requests and requires a complete, enforced usage receipt. +Missing support fails with `api_key.execution_budget_unsupported` before compute; explicitly uncapped keys keep their existing execution path. +A receipt ceiling alone does not prove that upstream provider spending was bounded. +Current remote agent-app chat adapters forward execution limits but do not implement this preparation contract. +They remain unsupported for capped execution until the maintained Sandbox backend enforces per-turn limits across its complete execution lifecycle. The usage store writes USD values as integer nanodollars instead of SQL floating-point values. Production requires either `x402.verifySigner` or `verifyApiKey`. diff --git a/package.json b/package.json index 343dfc2..bf6d930 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@tangle-network/agent-gateway", - "version": "0.9.0", + "version": "0.10.0", "packageManager": "pnpm@10.28.0", "engines": { "node": ">=22.13.0" diff --git a/src/api-key-budget.ts b/src/api-key-budget.ts new file mode 100644 index 0000000..512d2f3 --- /dev/null +++ b/src/api-key-budget.ts @@ -0,0 +1,64 @@ +import { apiKeySettlementCostCents, ApiKeyRequestClaimUnavailableError } from './api-keys' +import type { AuthorizedRequest } from './dispatch-types' +import type { GatewayConfig, GatewaySandboxContext, SandboxBox, SandboxExecutionBudget } from './types' + +export function apiKeyReservationQuote(authz: AuthorizedRequest, config: GatewayConfig): number | undefined { + const capped = authz.keyInfo?.spendingLimitCents !== undefined && authz.keyInfo.spendingLimitCents !== null + if (!capped) return undefined + if (!config.apiKeyReservationLifecycle) throw new ApiKeyRequestClaimUnavailableError('API key spending reservations are not configured') + const budget = authz.executionBudget + return apiKeySettlementCostCents(Math.max( + (budget.maxInputTokens + budget.maxOutputTokens + budget.maxReasoningTokens + budget.maxToolTokens) * authz.agent.pricePerTokenUsd, + budget.maxProviderCostUsd, + )) +} + +export async function prepareApiKeyPrompt( + box: SandboxBox, message: string, consumerId: string, systemPrompt: string | undefined, + executionBudget: SandboxExecutionBudget, signal: AbortSignal, sessionId?: string, + context?: GatewaySandboxContext, +) { + const promptOptions = { sessionId: sessionId ?? `consumer:${consumerId}`, systemPrompt, + maxOutputTokens: executionBudget.maxOutputTokens, executionBudget, signal } + if (!context?.apiKeyReservation) return { promptOptions, prepared: undefined } + if (!box.prepareBudgetedPrompt) throw new ApiKeyBudgetUnsupportedError('This execution adapter cannot enforce API key budgets') + const prepared = await box.prepareBudgetedPrompt(message, promptOptions) + if (prepared?.status === 'unsupported') throw new ApiKeyBudgetUnsupportedError(prepared.reason) + if (prepared?.status !== 'prepared' || typeof prepared.start !== 'function') { + throw new ApiKeyBudgetUnsupportedError('Execution adapter returned invalid budget preparation') + } + return { promptOptions, prepared } +} + +export class ApiKeyBudgetUnsupportedError extends Error { + readonly code = 'api_key.execution_budget_unsupported' + constructor(message: string) { super(message); this.name = 'ApiKeyBudgetUnsupportedError' } +} + +export function assertApiKeyRequestClaim( + claim: import('./types').ApiKeyRequestClaimResult, +): void { + if (typeof claim.allowed !== 'boolean') { + throw new ApiKeyRequestClaimUnavailableError('API key request claim is invalid') + } + for (const [name, value] of [ + ['minuteRemaining', claim.minuteRemaining], + ['dailyRemaining', claim.dailyRemaining], + ] as const) { + if (!Number.isSafeInteger(value) || value < 0) { + throw new ApiKeyRequestClaimUnavailableError(`API key request claim ${name} is invalid`) + } + } + for (const [name, value] of [ + ['minuteResetAt', claim.minuteResetAt], + ['dailyResetAt', claim.dailyResetAt], + ] as const) { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new ApiKeyRequestClaimUnavailableError(`API key request claim ${name} is invalid`) + } + } + if (!claim.allowed && claim.reason !== 'minute' && claim.reason !== 'daily' && claim.reason !== 'spending') { + throw new ApiKeyRequestClaimUnavailableError('API key request claim reason is invalid') + } +} + diff --git a/src/api-key-reservations-sql.ts b/src/api-key-reservations-sql.ts new file mode 100644 index 0000000..c697070 --- /dev/null +++ b/src/api-key-reservations-sql.ts @@ -0,0 +1,57 @@ +import type { SqlAdapter } from './a2a/task-store-sql' + +/** Reservations outlive rate-counter pruning. A usage row atomically closes its matching hold. */ +export class ApiKeyReservationsSql { + constructor( + private readonly db: SqlAdapter, + readonly table: string, + private readonly keys: string, + private readonly usage: string, + ) {} + + outstanding(excludeRequest = false): string { + return `COALESCE((SELECT SUM(r.reserved_cents) FROM ${this.table} r + WHERE r.key_id = k.id AND r.state <> 'released' + ${excludeRequest ? 'AND r.request_id <> ?' : ''} + AND NOT EXISTS (SELECT 1 FROM ${this.usage} u WHERE u.request_id = r.request_id)), 0)` + } + + async reserve(keyId: string, requestId: string, cents: number): Promise { + if (!Number.isSafeInteger(cents) || cents < 0) throw new TypeError('Reservation cents must be a non-negative safe integer') + await this.db.exec(`INSERT INTO ${this.table} (request_id, key_id, reserved_cents, state, created_at) + SELECT ?, k.id, ?, 'reserved', ? FROM ${this.keys} k + WHERE k.id = ? AND (k.expires_at IS NULL OR k.expires_at > ?) + AND (k.spending_limit_cents IS NULL OR k.spent_cents + + COALESCE((SELECT SUM(u.cost_cents) FROM ${this.usage} u WHERE u.key_id = k.id), 0) + + ${this.outstanding()} + ? <= k.spending_limit_cents) + ON CONFLICT(request_id) DO NOTHING`, [requestId, cents, Math.floor(Date.now() / 1000), keyId, Math.floor(Date.now() / 1000), cents]) + const row = (await this.db.query<{ key_id: string; reserved_cents: number; state: string }>( + `SELECT key_id, reserved_cents, state FROM ${this.table} WHERE request_id = ?`, [requestId], + ))[0] + if (!row) return false + if (row.key_id !== keyId || Number(row.reserved_cents) !== cents) throw new Error('API key reservation id was reused with different terms') + if (row.state !== 'reserved') throw new Error('API key reservation is no longer available for execution') + if ((await this.db.query(`SELECT request_id FROM ${this.usage} WHERE request_id = ?`, [requestId])).length) { + throw new Error('API key reservation is already settled') + } + return true + } + + async begin(keyId: string, requestId: string): Promise { + const result = await this.db.exec(`UPDATE ${this.table} SET state = 'executing' + WHERE request_id = ? AND key_id = ? AND state = 'reserved' + AND EXISTS (SELECT 1 FROM ${this.keys} k WHERE k.id = ? AND (k.expires_at IS NULL OR k.expires_at > ?) + AND (k.spending_limit_cents IS NULL OR k.spent_cents + + COALESCE((SELECT SUM(u.cost_cents) FROM ${this.usage} u WHERE u.key_id = k.id), 0) + + ${this.outstanding()} <= k.spending_limit_cents)) + AND NOT EXISTS (SELECT 1 FROM ${this.usage} WHERE request_id = ?)`, + [requestId, keyId, keyId, Math.floor(Date.now() / 1000), requestId]) + if (result.rowsAffected !== 1) throw new Error('API key reservation cannot begin execution') + } + + async release(keyId: string, requestId: string): Promise { + // An execution handoff is uncertain even when it produced no visible output. + await this.db.exec(`UPDATE ${this.table} SET state = 'released' + WHERE request_id = ? AND key_id = ? AND state = 'reserved'`, [requestId, keyId]) + } +} diff --git a/src/api-key-store-sql.ts b/src/api-key-store-sql.ts index 379a794..2d56952 100644 --- a/src/api-key-store-sql.ts +++ b/src/api-key-store-sql.ts @@ -1,4 +1,5 @@ import type { ApiKey, ApiKeyStore } from './api-keys' +import { ApiKeyReservationsSql } from './api-key-reservations-sql' import type { SqlAdapter } from './a2a/task-store-sql' import { requireSqlIdentifier } from './sql' import type { ApiKeyRequestClaimResult } from './types' @@ -43,6 +44,7 @@ export interface SqlApiKeyStoreOptions { table?: string usageTable?: string requestTable?: string + reservationTable?: string } const REQUEST_CLAIM_PRUNE_INTERVAL = 256 @@ -53,6 +55,7 @@ export function sqlApiKeyStoreSchemaStatements( const table = requireSqlIdentifier(options.table ?? 'agent_api_key') const usageTable = requireSqlIdentifier(options.usageTable ?? `${table}_usage`) const requestTable = requireSqlIdentifier(options.requestTable ?? `${table}_request`) + const reservationTable = requireSqlIdentifier(options.reservationTable ?? `${table}_reservation`) return [ `CREATE TABLE IF NOT EXISTS ${table} ( id TEXT PRIMARY KEY, @@ -93,6 +96,12 @@ export function sqlApiKeyStoreSchemaStatements( ON ${requestTable} (key_id, created_at)`, `CREATE INDEX IF NOT EXISTS idx_${requestTable}_day ON ${requestTable} (key_id, day_bucket)`, + `CREATE TABLE IF NOT EXISTS ${reservationTable} ( + request_id TEXT PRIMARY KEY, key_id TEXT NOT NULL, reserved_cents BIGINT NOT NULL, + state TEXT NOT NULL, created_at BIGINT NOT NULL, + FOREIGN KEY (key_id) REFERENCES ${table}(id) ON DELETE CASCADE + )`, + `CREATE INDEX IF NOT EXISTS idx_${reservationTable}_key ON ${reservationTable} (key_id)`, ] } @@ -141,6 +150,7 @@ export class SqlApiKeyStore implements ApiKeyStore { private readonly table: string private readonly usageTable: string private readonly requestTable: string + readonly reservations: ApiKeyReservationsSql constructor( private readonly db: SqlAdapter, @@ -149,8 +159,10 @@ export class SqlApiKeyStore implements ApiKeyStore { this.table = requireSqlIdentifier(options.table ?? 'agent_api_key') this.usageTable = requireSqlIdentifier(options.usageTable ?? `${this.table}_usage`) this.requestTable = requireSqlIdentifier(options.requestTable ?? `${this.table}_request`) - if (new Set([this.table, this.usageTable, this.requestTable]).size !== 3) { - throw new TypeError('API key, usage, and request table names must differ') + const reservationTable = requireSqlIdentifier(options.reservationTable ?? `${this.table}_reservation`) + this.reservations = new ApiKeyReservationsSql(db, reservationTable, this.table, this.usageTable) + if (new Set([this.table, this.usageTable, this.requestTable, reservationTable]).size !== 4) { + throw new TypeError('API key, usage, request, and reservation table names must differ') } } @@ -160,6 +172,7 @@ export class SqlApiKeyStore implements ApiKeyStore { table: this.table, usageTable: this.usageTable, requestTable: this.requestTable, + reservationTable: this.reservations.table, })) { await this.db.exec(statement) } @@ -256,7 +269,21 @@ export class SqlApiKeyStore implements ApiKeyStore { return result.rowsAffected > 0 } - async claimRequest( + async claimRequest(keyId: string, requestId: string, requestedAt = new Date(), reservationCents?: number): Promise { + const claim = await this.claimRateRequest(keyId, requestId, requestedAt) + if (!claim.allowed) return claim + if (reservationCents === undefined) { + const current = (await this.db.query<{ spending_limit_cents: number | null }>( + `SELECT spending_limit_cents FROM ${this.table} WHERE id = ?`, [keyId], + ))[0] + if (!current || current.spending_limit_cents !== null) throw new Error('Finite API key caps require a spending reservation quote') + return claim + } + const reserved = await this.reservations.reserve(keyId, requestId, reservationCents) + return reserved ? { ...claim, reservedCents: reservationCents } : { ...claim, allowed: false, reason: 'spending' } + } + + private async claimRateRequest( keyId: string, requestId: string, requestedAt = new Date(), @@ -352,10 +379,12 @@ export class SqlApiKeyStore implements ApiKeyStore { k.spending_limit_cents IS NULL OR k.spent_cents + COALESCE(( SELECT SUM(u.cost_cents) FROM ${this.usageTable} AS u WHERE u.key_id = k.id - ), 0) + ? <= k.spending_limit_cents + ), 0) + ${this.reservations.outstanding(true)} + ? <= k.spending_limit_cents ) + AND NOT EXISTS (SELECT 1 FROM ${this.reservations.table} r WHERE r.request_id = ? + AND (r.key_id <> k.id OR r.state = 'released' OR r.reserved_cents < ?)) ON CONFLICT(request_id) DO NOTHING`, - [usageRequestId, costCents, createdAt, keyId, costCents], + [usageRequestId, costCents, createdAt, keyId, usageRequestId, costCents, usageRequestId, costCents], ) if (result.rowsAffected === 1) return diff --git a/src/api-keys.ts b/src/api-keys.ts index 8578f77..dd46510 100644 --- a/src/api-keys.ts +++ b/src/api-keys.ts @@ -69,6 +69,7 @@ export interface ApiKeyStore { keyId: string, requestId: string, requestedAt?: Date, + reservationCents?: number, ): Promise } @@ -121,6 +122,7 @@ export function createApiKeyRequestClaim( input.keyInfo.keyId, input.requestId, input.requestedAt, + input.reservationCents, ) } @@ -163,7 +165,7 @@ export async function verifyApiKeyFromStore( authHeader: string, store: ApiKeyStore, prefix = 'ak_', -): Promise<{ key: ApiKey; keyId: string; consumerId: string; ownerId: string; scopes: string[]; rateLimitPerMinute: number; dailyLimit: number } | null> { +): Promise<{ key: ApiKey; keyId: string; consumerId: string; ownerId: string; scopes: string[]; rateLimitPerMinute: number; dailyLimit: number; spendingLimitCents: number | null } | null> { const bearerPrefix = `Bearer ${prefix}` if (!authHeader.startsWith(bearerPrefix)) return null @@ -186,6 +188,7 @@ export async function verifyApiKeyFromStore( scopes: key.scopes, rateLimitPerMinute: key.rateLimit, dailyLimit: key.dailyLimit, + spendingLimitCents: key.spendingLimitCents, } } diff --git a/src/dispatch-payment.ts b/src/dispatch-payment.ts index bff3712..caa2aa6 100644 --- a/src/dispatch-payment.ts +++ b/src/dispatch-payment.ts @@ -1,3 +1,4 @@ +import { apiKeyReservationQuote, assertApiKeyRequestClaim } from './api-key-budget' import { assertMppChargeOperation, mppPaymentOperationId, @@ -47,19 +48,22 @@ export async function claimPayment( if (!claimRequest) { if ( authz.keyInfo.rateLimitPerMinute !== undefined || - authz.keyInfo.dailyLimit !== undefined + authz.keyInfo.dailyLimit !== undefined || + (authz.keyInfo.spendingLimitCents !== undefined && authz.keyInfo.spendingLimitCents !== null) ) { throw new ApiKeyRequestClaimUnavailableError( 'API key request limits are not configured', ) } } else { + const reservationCents = apiKeyReservationQuote(authz, config) let claim try { claim = await claimRequest({ keyInfo: authz.keyInfo, requestId: authz.requestId, requestedAt: new Date(authz.startMs), + ...(reservationCents !== undefined ? { reservationCents } : {}), }) } catch (error) { if ( @@ -73,6 +77,10 @@ export async function claimPayment( } assertApiKeyRequestClaim(claim) if (!claim.allowed) throw new ApiKeyRequestLimitExceededError(claim) + if (reservationCents !== undefined) { + if (claim.reservedCents !== reservationCents) throw new ApiKeyRequestClaimUnavailableError('API key spending reservation is invalid') + authz.apiKeyReservedCents = reservationCents + } authz.rateLimitRemaining = Math.min( authz.rateLimitRemaining ?? claim.minuteRemaining, claim.minuteRemaining, @@ -264,39 +272,15 @@ export async function claimPayment( } } -function assertApiKeyRequestClaim( - claim: import('./types').ApiKeyRequestClaimResult, -): void { - if (typeof claim.allowed !== 'boolean') { - throw new ApiKeyRequestClaimUnavailableError('API key request claim is invalid') - } - for (const [name, value] of [ - ['minuteRemaining', claim.minuteRemaining], - ['dailyRemaining', claim.dailyRemaining], - ] as const) { - if (!Number.isSafeInteger(value) || value < 0) { - throw new ApiKeyRequestClaimUnavailableError(`API key request claim ${name} is invalid`) - } - } - for (const [name, value] of [ - ['minuteResetAt', claim.minuteResetAt], - ['dailyResetAt', claim.dailyResetAt], - ] as const) { - if (!Number.isSafeInteger(value) || value <= 0) { - throw new ApiKeyRequestClaimUnavailableError(`API key request claim ${name} is invalid`) - } - } - if (!claim.allowed && claim.reason !== 'minute' && claim.reason !== 'daily') { - throw new ApiKeyRequestClaimUnavailableError('API key request claim reason is invalid') - } -} - /** Release an owned operation when execution cannot produce a valid receipt. */ export async function releasePayment( authz: AuthorizedRequest, config: GatewayConfig, reason: string, ): Promise { + if (authz.apiKeyReservedCents !== undefined && authz.keyInfo) { + await config.apiKeyReservationLifecycle!.release(authz.keyInfo.keyId, authz.requestId) + } const ownsX402 = authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations @@ -356,6 +340,9 @@ export async function markPaymentExecutionStarted( authz: AuthorizedRequest, config: GatewayConfig, ): Promise { + if (authz.apiKeyReservedCents !== undefined && authz.keyInfo) { + await config.apiKeyReservationLifecycle!.begin(authz.keyInfo.keyId, authz.requestId) + } await updateExecutionLease(authz, config, true) } diff --git a/src/dispatch-sandbox.ts b/src/dispatch-sandbox.ts index 7b0f81b..596bf76 100644 --- a/src/dispatch-sandbox.ts +++ b/src/dispatch-sandbox.ts @@ -1,3 +1,4 @@ +import { prepareApiKeyPrompt } from './api-key-budget' import { redactSystemPromptFromOutput } from './filter' import type { A2ADispatchEvent, AuthorizedRequest } from './dispatch-types' import { @@ -23,6 +24,7 @@ export function buildGatewaySandboxContext( keyInfo: authz.keyInfo, requestId: authz.requestId, messages: authz.messages ?? [], + ...(authz.apiKeyReservedCents !== undefined ? { apiKeyReservation: { cents: authz.apiKeyReservedCents, executionBudget: authz.executionBudget } } : {}), ...(authz.threadId !== undefined ? { threadId: authz.threadId } : {}), } } @@ -121,37 +123,34 @@ export async function* dispatchSandboxStreamRich( const executionController = new AbortController() const forwardAbort = () => executionController.abort() if (signal?.aborted) return - signal?.addEventListener('abort', forwardAbort, { once: true }) - const executionBudget: SandboxExecutionBudget = { - maxInputTokens: maxInputTokens ?? maximumBillableInputTokens(agent, userMessage), - maxOutputTokens: outputLimit, - maxReasoningTokens: config.executionBudget?.maxReasoningTokens ?? outputLimit, - maxToolTokens: config.executionBudget?.maxToolTokens ?? outputLimit, - maxToolCalls: config.executionBudget?.maxToolCalls ?? 8, - maxProviderCostUsd: config.executionBudget?.maxProviderCostUsd ?? ( - (maxInputTokens ?? maximumBillableInputTokens(agent, userMessage)) + outputLimit + - (config.executionBudget?.maxReasoningTokens ?? outputLimit) + - (config.executionBudget?.maxToolTokens ?? outputLimit) - ) * agent.pricePerTokenUsd, - } - if (executionController.signal.aborted) return - await onExecutionStart?.() - if (executionController.signal.aborted) return let heartbeatError: unknown let heartbeatInFlight: Promise | undefined let heartbeatTimer: ReturnType | undefined let iterator: AsyncIterator | undefined try { + signal?.addEventListener('abort', forwardAbort, { once: true }) + const executionBudget: SandboxExecutionBudget = sandboxContext?.apiKeyReservation?.executionBudget ?? { + maxInputTokens: maxInputTokens ?? maximumBillableInputTokens(agent, userMessage), + maxOutputTokens: outputLimit, + maxReasoningTokens: config.executionBudget?.maxReasoningTokens ?? outputLimit, + maxToolTokens: config.executionBudget?.maxToolTokens ?? outputLimit, + maxToolCalls: config.executionBudget?.maxToolCalls ?? 8, + maxProviderCostUsd: config.executionBudget?.maxProviderCostUsd ?? ( + (maxInputTokens ?? maximumBillableInputTokens(agent, userMessage)) + outputLimit + + (config.executionBudget?.maxReasoningTokens ?? outputLimit) + + (config.executionBudget?.maxToolTokens ?? outputLimit) + ) * agent.pricePerTokenUsd, + } + const { prepared, promptOptions } = await prepareApiKeyPrompt(box, userMessage, consumerId, + agent.systemPrompt, executionBudget, executionController.signal, sessionId, sandboxContext) + if (prepared) requiresReceipt = true + if (executionController.signal.aborted) return + await onExecutionStart?.() + if (executionController.signal.aborted) return // This durable handoff is after sandbox acquisition and immediately before // the adapter call that may start paid work. await onSandboxStart?.() - const promptStream = box.streamPrompt(userMessage, { - sessionId: sessionId ?? `consumer:${consumerId}`, - systemPrompt: agent.systemPrompt, - maxOutputTokens: outputLimit, - executionBudget, - signal: executionController.signal, - }) + const promptStream = prepared ? prepared.start() : box.streamPrompt(userMessage, promptOptions) iterator = promptStream[Symbol.asyncIterator]() const heartbeatMs = onExecutionHeartbeat ? Math.max(100, Math.min( diff --git a/src/dispatch-types.ts b/src/dispatch-types.ts index b194db4..fd93fad 100644 --- a/src/dispatch-types.ts +++ b/src/dispatch-types.ts @@ -45,6 +45,7 @@ export interface AuthorizedRequest { executionBudget: SandboxExecutionBudget requiredPaymentAmount: bigint paymentPayload: Record | null + apiKeyReservedCents?: number paymentNonceKey?: string mppMethod?: string /** Live generic MPP credential. Never write it to the recovery store. */ diff --git a/src/index.ts b/src/index.ts index de33447..aedc2ca 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,3 +1,4 @@ +export { ApiKeyBudgetUnsupportedError } from './api-key-budget' export { createAgentGateway } from './middleware' export { maximumBillableInputTokens, reclaimPayment, SandboxStreamError } from './dispatch' export { @@ -130,6 +131,7 @@ export type { SandboxUsageReceipt, SandboxStreamEvent, SandboxBox, + SandboxPromptOptions, GatewaySandboxContext, ApiKeyGatewayConfig, CreateAgentGatewayConfig, diff --git a/src/middleware.ts b/src/middleware.ts index 05821f3..a44cdcb 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -1,3 +1,4 @@ +import { ApiKeyBudgetUnsupportedError } from './api-key-budget' import { Hono } from 'hono' import { isChatMessageArray } from './chat-input' @@ -673,9 +674,9 @@ async function completeChatCompletion( await reportCompletionError(obs, ctx, authz.consumerId, error) await releaseCompletionAfterFailure(authz, config, error, workObserved, usage) return c.json( - { error: { message: safeCompletionErrorMessage(error), type: 'server_error' } }, + { error: { message: safeCompletionErrorMessage(error), type: 'server_error', ...(error instanceof ApiKeyBudgetUnsupportedError ? { code: error.code } : {}) } }, { - status: 500, + status: error instanceof ApiKeyBudgetUnsupportedError ? 503 : 500, headers: completionHeaders(authz, false), }, ) @@ -783,7 +784,7 @@ function streamChatCompletions( ) { controller.enqueue( encoder.encode( - `data: ${JSON.stringify({ error: { message: safeMessage, type: 'server_error' } })}\n\n`, + `data: ${JSON.stringify({ error: { message: safeMessage, type: 'server_error', ...(err instanceof ApiKeyBudgetUnsupportedError ? { code: err.code } : {}) } })}\n\n`, ), ) } diff --git a/src/types.ts b/src/types.ts index 2d4cb94..cf3b474 100644 --- a/src/types.ts +++ b/src/types.ts @@ -187,6 +187,7 @@ export interface ApiKeyInfo { rateLimitPerMinute?: number /** Per-key daily limit override. */ dailyLimit?: number + spendingLimitCents?: number | null } /** Resolve the URL where a caller can obtain an API key for one agent. */ @@ -198,7 +199,8 @@ export type ApiKeyPurchaseUrl = export interface ApiKeyRequestClaimResult { allowed: boolean /** The exhausted policy when `allowed` is false. */ - reason?: 'minute' | 'daily' + reason?: 'minute' | 'daily' | 'spending' + reservedCents?: number minuteRemaining: number dailyRemaining: number minuteResetAt: number @@ -210,6 +212,7 @@ export interface ApiKeyRequestClaimInput { keyInfo: ApiKeyInfo requestId: string requestedAt: Date + reservationCents?: number } // --- Sandbox interface --- @@ -240,16 +243,24 @@ export interface SandboxStreamEvent { } } +export type SandboxPromptOptions = { + sessionId?: string + systemPrompt?: string + maxOutputTokens?: number + executionBudget?: SandboxExecutionBudget + signal?: AbortSignal +} + export interface SandboxBox { + /** Prepare without starting compute. The returned stream must enforce every supplied budget across retries and child calls. */ + prepareBudgetedPrompt?(message: string, opts: SandboxPromptOptions & { executionBudget: SandboxExecutionBudget }): Promise< + | { status: 'unsupported'; reason: string } + | { status: 'prepared'; start: () => AsyncIterable } + > + streamPrompt( message: string, - opts?: { - sessionId?: string - systemPrompt?: string - maxOutputTokens?: number - executionBudget?: SandboxExecutionBudget - signal?: AbortSignal - }, + opts?: SandboxPromptOptions, ): AsyncIterable } @@ -260,6 +271,7 @@ export interface GatewaySandboxContext { keyInfo: ApiKeyInfo | null requestId: string messages: ChatMessage[] + apiKeyReservation?: { cents: number; executionBudget: SandboxExecutionBudget } /** Stable UI conversation id when `conversationMode` is `thread`. */ threadId?: string } @@ -333,6 +345,11 @@ export interface GatewayConfig { * Atomically count an accepted API-key request before compute starts. * Required when `verifyApiKey` returns a minute or daily request limit. */ + apiKeyReservationLifecycle?: { + begin(keyId: string, requestId: string): Promise + release(keyId: string, requestId: string): Promise + } + claimApiKeyRequest?: ( input: ApiKeyRequestClaimInput, ) => Promise diff --git a/tests/api-key-reservations.test.ts b/tests/api-key-reservations.test.ts new file mode 100644 index 0000000..37386b8 --- /dev/null +++ b/tests/api-key-reservations.test.ts @@ -0,0 +1,239 @@ +import { DatabaseSync } from 'node:sqlite' +import { createHash } from 'node:crypto' +import { describe, expect, it, vi } from 'vitest' +import { SqlApiKeyStore } from '../src/api-key-store-sql' +import { createAgentGateway } from '../src/middleware' +import { createApiKeyRequestClaim, createApiKeyUsageSettlement, verifyApiKeyFromStore } from '../src/api-keys' +import type { SandboxBox, SandboxStreamEvent } from '../src/types' +import { prepareApiKeyPrompt, ApiKeyBudgetUnsupportedError } from '../src/api-key-budget' +import { dispatchSandboxStreamRich } from '../src/dispatch-sandbox' + +async function fixture(cap: number | null = 1) { + const db = new DatabaseSync(':memory:') + const adapter = { + async exec(sql: string, params: unknown[] = []) { return { rowsAffected: Number(db.prepare(sql).run(...params as never[]).changes) } }, + async query(sql: string, params: unknown[] = []) { return db.prepare(sql).all(...params as never[]) as T[] }, + } + const store = new SqlApiKeyStore(adapter) + await store.migrate() + const token = 'ak_synthetic_fixture' + const key = await store.create('owner', { name: 'fixture', keyHash: createHash('sha256').update(token).digest('hex'), keyPrefix: 'ak_', scopes: ['chat'], rateLimit: 60, dailyLimit: 100, spendingLimitCents: cap, expiresAt: null }) + return { db, store, adapter, token, key } +} + + +function gateway(store: SqlApiKeyStore, box: SandboxBox) { + return createAgentGateway({ + resolveAgent: async () => ({ id: 'agent', ownerId: 'owner', slug: 'agent', enabled: true, pricePerTokenUsd: 0, platformFeePercent: 0, sandboxEndpoint: null, remoteSandboxId: null, remoteBearerToken: null }), + authorizeConsumer: async () => ({ allow: true }), + verifyApiKey: header => verifyApiKeyFromStore(header, store), + claimApiKeyRequest: createApiKeyRequestClaim(store), apiKeyReservationLifecycle: store.reservations, + settlePayment: createApiKeyUsageSettlement(store), recordUsage: async () => {}, getSandbox: async () => box, + executionBudget: { maxProviderCostUsd: 0.01 }, a2a: false, + }) +} + +function request(app: ReturnType, token: string) { + return app.request('/agent/chat/completions', { method: 'POST', headers: { Authorization: `Bearer ${token}`, 'Content-Type': 'application/json' }, body: JSON.stringify({ messages: [{ role: 'user', content: 'x' }], stream: true }) }) +} + +describe('API key spending reservations', () => { + it.each(['preparation', 'execution-start'])('removes abort forwarding when %s fails', async (failure) => { + const controller = new AbortController() + const added = vi.spyOn(controller.signal, 'addEventListener') + const removed = vi.spyOn(controller.signal, 'removeEventListener') + const budget = { maxInputTokens: 1, maxOutputTokens: 1, maxReasoningTokens: 0, maxToolTokens: 0, maxToolCalls: 0, maxProviderCostUsd: 0.01 } + const box: SandboxBox = { + async *streamPrompt() { throw new Error('Unbounded execution is forbidden') }, + async prepareBudgetedPrompt() { + if (failure === 'preparation') throw new Error('Preparation failed') + return { status: 'prepared', start: async function* () { throw new Error('Execution must not start') } } + }, + } + const drain = async () => { + for await (const _event of dispatchSandboxStreamRich( + { id: 'agent', ownerId: 'owner', slug: 'agent', enabled: true, pricePerTokenUsd: 0, platformFeePercent: 0, sandboxEndpoint: null, remoteSandboxId: null, remoteBearerToken: null }, + 'x', 'apikey:key', { resolveAgent: async () => null, authorizeConsumer: async () => ({ allow: true }), getSandbox: async () => box, recordUsage: async () => {}, x402: { demoMode: true } }, + controller.signal, undefined, 1, async () => { throw new Error('Start failed') }, false, undefined, 1, undefined, + { consumerId: 'apikey:key', paymentMethod: 'apikey', keyInfo: null, requestId: 'request', messages: [], apiKeyReservation: { cents: 1, executionBudget: budget } }, + )) { /* Drain the stream until its expected failure. */ } + } + await expect(drain()).rejects.toThrow() + expect(added).toHaveBeenCalledOnce() + expect(removed).toHaveBeenCalledWith('abort', added.mock.calls[0][1]) + }) + + it.each([undefined, null, {}, { status: 'other' }, { status: 'prepared', start: true }])( + 'rejects malformed preparation %j without starting either execution path', async (result) => { + const { db, store, token } = await fixture() + let unboundedStarts = 0 + const app = gateway(store, { + async *streamPrompt() { unboundedStarts++; yield { type: 'finish' } }, + // Reflect models a JavaScript host that bypasses the return type. + prepareBudgetedPrompt: async () => Reflect.apply(() => result, undefined, []), + }) + try { + const response = await request(app, token) + expect(await response.text()).toContain('api_key.execution_budget_unsupported') + expect(unboundedStarts).toBe(0) + expect(db.prepare('SELECT state FROM agent_api_key_reservation').get()).toMatchObject({ state: 'released' }) + } finally { db.close() } + }, + ) + + it('adds reservation storage without rewriting existing keys or settled usage', async () => { + const { db, store, key } = await fixture(10) + try { + db.exec('DROP TABLE agent_api_key_reservation') + db.prepare('UPDATE agent_api_key SET spent_cents = 2 WHERE id = ?').run(key.id) + db.prepare('INSERT INTO agent_api_key_usage VALUES (?, ?, ?, ?)').run('old-receipt', key.id, 3, 1) + const before = await store.list('owner') + await store.migrate() + await store.migrate() + expect(await store.list('owner')).toEqual(before) + expect((await store.claimRequest(key.id, 'new', undefined, 6)).allowed).toBe(false) + expect((await store.claimRequest(key.id, 'fits', undefined, 5)).allowed).toBe(true) + } finally { db.close() } + }) + + it('rejects a missing quote when a key has acquired a finite cap', async () => { + const { db, store, key } = await fixture(null) + try { + await store.claimRequest(key.id, 'uncapped') + db.prepare('UPDATE agent_api_key SET spending_limit_cents = 1 WHERE id = ?').run(key.id) + await expect(store.claimRequest(key.id, 'newly-capped')).rejects.toThrow('require a spending reservation quote') + } finally { db.close() } + }) + + it('retains a handoff with no receipt even when the adapter fails before visible output', async () => { + const { db, store, token } = await fixture() + const app = gateway(store, { + async *streamPrompt() { throw new Error('Unbounded execution is forbidden') }, + async prepareBudgetedPrompt() { return { status: 'prepared', start: async function* () { throw new DOMException('Provider stream aborted', 'AbortError') } } }, + }) + try { + const response = await request(app, token) + await response.text() + expect(db.prepare('SELECT state FROM agent_api_key_reservation').get()).toMatchObject({ state: 'executing' }) + expect((await request(app, token)).status).toBe(429) + expect((await store.list('owner'))[0].spentCents).toBe(0) + } finally { db.close() } + }) + + it('releases an unsupported gateway reservation before compute', async () => { + const { db, store, token } = await fixture() + let started = 0 + const app = gateway(store, { async *streamPrompt() { started++; yield { type: 'finish' } } }) + try { + for (let n = 0; n < 2; n++) { + const response = await request(app, token) + expect(await response.text()).toContain('api_key.execution_budget_unsupported') + } + expect(started).toBe(0) + expect(db.prepare("SELECT count(*) AS n FROM agent_api_key_reservation WHERE state <> 'released'").get()).toMatchObject({ n: 0 }) + } finally { db.close() } + }) + + it('fails closed for unsupported capped hosts and leaves uncapped hosts usable', async () => { + let started = 0 + const box: SandboxBox = { async *streamPrompt() { started++; yield { type: 'finish' } } } + const budget = { maxInputTokens: 1, maxOutputTokens: 1, maxReasoningTokens: 0, maxToolTokens: 0, maxToolCalls: 0, maxProviderCostUsd: 0.01 } + const context = { consumerId: 'apikey:key', paymentMethod: 'apikey' as const, keyInfo: null, requestId: 'request', messages: [], apiKeyReservation: { cents: 1, executionBudget: budget } } + await expect(prepareApiKeyPrompt(box, 'x', 'apikey:key', undefined, budget, new AbortController().signal, undefined, context)).rejects.toBeInstanceOf(ApiKeyBudgetUnsupportedError) + expect(started).toBe(0) + const uncapped = await prepareApiKeyPrompt(box, 'x', 'apikey:key', undefined, budget, new AbortController().signal) + expect(uncapped.prepared).toBeUndefined() + for await (const _event of box.streamPrompt('x', uncapped.promptOptions)) { /* Drain the legacy adapter. */ } + expect(started).toBe(1) + }) + + it('prevents duplicate starts, settlement above the reservation, and revoked execution', async () => { + const { db, store, key } = await fixture(10) + try { + await store.claimRequest(key.id, 'one', undefined, 5) + const starts = await Promise.allSettled([store.reservations.begin(key.id, 'one'), store.reservations.begin(key.id, 'one')]) + expect(starts.filter(result => result.status === 'fulfilled')).toHaveLength(1) + await expect(store.recordUsage(key.id, 6, 'one')).rejects.toThrow() + expect((await store.list('owner'))[0].spentCents).toBe(0) + await store.claimRequest(key.id, 'two', undefined, 5) + await store.delete('owner', key.id) + await expect(store.reservations.begin(key.id, 'two')).rejects.toThrow() + } finally { db.close() } + }) + + it('admits only one competing full quote across independent stores', async () => { + const { db, store, adapter, key } = await fixture() + try { + const other = new SqlApiKeyStore(adapter) + const claims = await Promise.all([store.claimRequest(key.id, 'one', undefined, 1), other.claimRequest(key.id, 'two', undefined, 1)]) + expect(claims.filter(claim => claim.allowed)).toHaveLength(1) + expect(claims.find(claim => !claim.allowed)?.reason).toBe('spending') + } finally { db.close() } + }) + + it('reconciles once, rejects altered replay, and releases only unused pre-execution holds', async () => { + const { db, store, key } = await fixture(10) + try { + await store.claimRequest(key.id, 'one', undefined, 10) + await store.reservations.begin(key.id, 'one') + await store.reservations.release(key.id, 'one') + expect((await store.claimRequest(key.id, 'two', undefined, 1)).allowed).toBe(false) + await store.recordUsage(key.id, 3, 'one') + await store.recordUsage(key.id, 3, 'one') + await expect(store.recordUsage(key.id, 4, 'one')).rejects.toThrow('different usage') + expect((await store.claimRequest(key.id, 'two', undefined, 7)).allowed).toBe(true) + await store.reservations.release(key.id, 'two') + await store.reservations.release(key.id, 'two') + expect((await store.claimRequest(key.id, 'three', undefined, 7)).allowed).toBe(true) + await expect(store.reservations.begin(key.id, 'two')).rejects.toThrow() + expect((await store.list('owner'))[0].spentCents).toBe(3) + } finally { db.close() } + }) + + it('retains holds after rate-counter pruning and rechecks expiry before execution', async () => { + const { db, store, key } = await fixture() + try { + await store.claimRequest(key.id, 'one', undefined, 1) + db.exec('DELETE FROM agent_api_key_request') + expect((await store.claimRequest(key.id, 'two', undefined, 1)).allowed).toBe(false) + db.prepare('UPDATE agent_api_key SET expires_at = 1 WHERE id = ?').run(key.id) + await expect(store.reservations.begin(key.id, 'one')).rejects.toThrow() + await store.reservations.release(key.id, 'one') + expect(db.prepare('SELECT state FROM agent_api_key_reservation').get()).toMatchObject({ state: 'released' }) + } finally { db.close() } + }) + + it('runs an enforcing prepared adapter once and rejects competing provider work', async () => { + const { db, store, token } = await fixture() + let started = 0 + let release!: () => void + const gate = new Promise(resolve => { release = resolve }) + let entered!: () => void + const firstStarted = new Promise(resolve => { entered = resolve }) + const box: SandboxBox = { + async *streamPrompt() { throw new Error('Unbounded path must never execute') }, + async prepareBudgetedPrompt(_message, options) { + const cost = 0.01 + if (cost > options.executionBudget.maxProviderCostUsd) return { status: 'unsupported', reason: 'Provider minimum exceeds budget' } + return { status: 'prepared', start: async function* (): AsyncIterable { + started++; entered(); await gate + // This synthetic provider checks the full call price before it executes. + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid' } } + yield { type: 'sandbox.usage', data: { usage: { inputTokens: 0, outputTokens: 1, reasoningTokens: 0, toolTokens: 0, toolCallCount: 0, providerCostUsd: cost, budgetEnforced: true } } } + } } + }, + } + const app = gateway(store, box) + try { + const first = await request(app, token) + const body = first.text() + await firstStarted + const second = await request(app, token) + expect(second.status).toBe(429) + expect(started).toBe(1) + release() + expect(await body).toContain('paid') + expect((await store.list('owner'))[0].spentCents).toBe(1) + } finally { release(); db.close() } + }) +}) diff --git a/tests/api-key-store-sql.test.ts b/tests/api-key-store-sql.test.ts index fbc3ce3..3de3033 100644 --- a/tests/api-key-store-sql.test.ts +++ b/tests/api-key-store-sql.test.ts @@ -26,7 +26,7 @@ async function createKey( store: SqlApiKeyStore, userId = 'user-1', keyHash = 'hash-1', - spendingLimitCents = 500, + spendingLimitCents: number | null = 500, limits: { rateLimit?: number; dailyLimit?: number } = {}, ) { return store.create(userId, { @@ -121,7 +121,7 @@ describe('SqlApiKeyStore', () => { store, 'user-1', 'limited-hash', - 500, + null, { rateLimit: 2, dailyLimit: 3 }, ) const minuteOne = new Date('2026-09-02T12:00:10.000Z') @@ -154,7 +154,7 @@ describe('SqlApiKeyStore', () => { store, 'user-1', 'request-hash', - 500, + null, { rateLimit: 5, dailyLimit: 5 }, ) const now = new Date('2026-09-02T12:00:10.000Z') @@ -184,7 +184,7 @@ describe('SqlApiKeyStore', () => { store, 'user-1', 'retention-hash', - 500, + null, { rateLimit: 300, dailyLimit: 300 }, ) await store.claimRequest(key.id, 'old-request', new Date('2026-09-01T12:00:00.000Z'))