From 87577584af7df92fe73e7ef9246cb1226297009c Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 16:33:37 -0600 Subject: [PATCH 01/32] fix(payments): bound x402 requests before execution --- README.md | 9 ++ package.json | 2 +- src/a2a/handler.ts | 15 ++- src/dispatch.ts | 210 +++++++++++++++++++++++++++++---- src/middleware.ts | 69 +++++++++-- src/types.ts | 35 +++++- src/verify.ts | 46 +++++--- tests/integration-x402.test.ts | 19 +-- tests/middleware.test.ts | 142 +++++++++++++++++++++- tests/observer.test.ts | 2 +- tests/verify.test.ts | 36 ++++++ 11 files changed, 517 insertions(+), 68 deletions(-) diff --git a/README.md b/README.md index 373b569..b22ee9b 100644 --- a/README.md +++ b/README.md @@ -26,14 +26,23 @@ app.route('/v1/agents', createAgentGateway({ x402: { operatorAddress: '0x…', chainId: 3799, + currencyDecimals: 6, verifySigner: verifySpendAuthSignature, + authorizePayment: reserveSpendAuthorization, }, + defaultOutputTokens: 1024, + maxOutputTokens: 4096, verifyApiKey: (authHeader) => verifyApiKeyFromStore(authHeader, apiKeyStore), })) ``` `x402.verifySigner` is required for production. Set `x402.demoMode: true` only for local development and tests; that explicit mode also enables the built-in `sk_agent_*` demo key verifier. +Keep `verifySigner` free of side effects. +Use `authorizePayment` to reserve or claim funds after rate limits, content checks, and product authorization succeed. +Before it calls the verifier, the gateway requires the signed amount to cover filtered input plus the requested output limit. +The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. +An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. MPP is method-specific. Configure `mpp.verifySigner` for production MPP credentials; it receives the decoded JSON payload when available plus the original decoded credential, and returns the authenticated consumer ID or `null`. diff --git a/package.json b/package.json index 8b0434c..4bb2b9d 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@tangle-network/agent-gateway", - "version": "0.7.1", + "version": "0.7.2", "packageManager": "pnpm@10.28.0", "repository": { "type": "git", diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 1968085..8ef0a5f 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -18,6 +18,7 @@ import { type GatewayState, authenticateAndGuard, dispatchSandboxStreamRich, + estimateBillableInputTokens, estimateTokens, settleAndRecord, } from '../dispatch' @@ -169,7 +170,6 @@ async function handleMessageSend( const { authz, task } = guard let responseText = '' - let outputTokens = 0 let inputRequiredPrompt: string | undefined let inputRequiredSeen = false try { @@ -180,10 +180,10 @@ async function handleMessageSend( deps.config, undefined, task.id, + authz.maxOutputTokens, )) { if (event.kind === 'text') { responseText += event.delta - outputTokens += estimateTokens(event.delta) } else { inputRequiredSeen = true inputRequiredPrompt = event.prompt @@ -208,8 +208,8 @@ async function handleMessageSend( await settleAndRecord( authz.agent, authz, - estimateTokens(authz.userMessage), - outputTokens, + estimateBillableInputTokens(authz.agent, authz.userMessage), + estimateTokens(responseText), deps.config, deps.state.obs, ) @@ -250,8 +250,7 @@ async function handleMessageStream( const { authz, task } = guard const controller = cancels.register(task.id) - const inputTokens = estimateTokens(authz.userMessage) - let outputTokens = 0 + const inputTokens = estimateBillableInputTokens(authz.agent, authz.userMessage) let responseText = '' const stream = new ReadableStream({ @@ -282,10 +281,10 @@ async function handleMessageStream( deps.config, controller.signal, task.id, + authz.maxOutputTokens, )) { if (event.kind === 'text') { responseText += event.delta - outputTokens += estimateTokens(event.delta) const artifactEvent: TaskArtifactUpdateEvent = { kind: 'artifact-update', taskId: task.id, @@ -326,7 +325,7 @@ async function handleMessageStream( authz.agent, authz, inputTokens, - outputTokens, + estimateTokens(responseText), deps.config, deps.state.obs, ) diff --git a/src/dispatch.ts b/src/dispatch.ts index ebdecdd..aaba8a3 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -36,6 +36,8 @@ export interface GatewayState { globalRateLimit: { limit: number; windowSeconds: number } requiredScope: string maxLen: number + maxOutputTokens: number + defaultOutputTokens: number obs?: GatewayObserver } @@ -49,6 +51,44 @@ export interface AuthorizedRequest { rateLimitRemaining: number | undefined requestId: string startMs: number + maxOutputTokens: number +} + +function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { + if (!Number.isFinite(value) || value <= 0) { + throw new Error('agent pricePerTokenUsd must be a finite positive number') + } + const [mantissa, exponentText] = value.toString().toLowerCase().split('e') + const exponent = exponentText ? Number(exponentText) : 0 + const [whole, fraction = ''] = mantissa.split('.') + let numerator = BigInt(`${whole}${fraction}`) + let scale = fraction.length - exponent + if (scale < 0) { + numerator *= 10n ** BigInt(-scale) + scale = 0 + } + return { numerator, denominator: 10n ** BigInt(scale) } +} + +/** Exact base-unit reservation required to cover the request's token ceiling. */ +export function requiredX402Amount( + pricePerTokenUsd: number, + inputTokens: number, + maxOutputTokens: number, + currencyDecimals = 6, +): bigint { + if (!Number.isSafeInteger(inputTokens) || inputTokens < 0) { + throw new Error('input token estimate must be a non-negative safe integer') + } + if (!Number.isSafeInteger(maxOutputTokens) || maxOutputTokens <= 0) { + throw new Error('max output tokens must be a positive safe integer') + } + if (!Number.isInteger(currencyDecimals) || currencyDecimals < 0 || currencyDecimals > 18) { + throw new Error('x402 currencyDecimals must be an integer between 0 and 18') + } + const { numerator, denominator } = decimalFraction(pricePerTokenUsd) + const scaled = BigInt(inputTokens + maxOutputTokens) * numerator * 10n ** BigInt(currencyDecimals) + return (scaled + denominator - 1n) / denominator } /** @@ -67,6 +107,7 @@ export async function authenticateAndGuard( messages: ChatMessage[], config: GatewayConfig, state: GatewayState, + requestedMaxOutputTokens?: number, ): Promise { const startMs = Date.now() const requestId = generateRequestId() @@ -84,15 +125,78 @@ export async function authenticateAndGuard( ) } + const maxOutputTokens = requestedMaxOutputTokens ?? state.defaultOutputTokens + if ( + !Number.isInteger(maxOutputTokens) || + maxOutputTokens <= 0 || + maxOutputTokens > state.maxOutputTokens + ) { + return c.json( + { + error: { + message: `max_tokens must be an integer between 1 and ${state.maxOutputTokens}`, + type: 'invalid_request', + code: 'invalid_max_tokens', + }, + }, + 400, + ) + } + + // Price the exact filtered input that reaches the sandbox. This must happen + // before x402 verification because a production verifier can reserve funds. + const { messages: filtered, injectionWarnings } = filterConsumerMessagesStrict( + messages, + state.maxLen, + ) + const userMessage = filtered + .filter((m) => m.role === 'user') + .map((m) => m.content) + .join('\n\n') + if (!userMessage) { + return c.json( + { error: { message: 'No user message provided', type: 'invalid_request' } }, + 400, + ) + } + + let requiredPaymentAmount: bigint + try { + requiredPaymentAmount = requiredX402Amount( + agent.pricePerTokenUsd, + maximumBillableInputTokens(agent, userMessage), + maxOutputTokens, + config.x402.currencyDecimals, + ) + } catch { + return c.json( + { + error: { + message: 'Agent payment configuration is invalid', + type: 'server_error', + code: 'invalid_payment_configuration', + }, + }, + 503, + ) + } + // Payment / auth. const spendAuthHeader = c.req.header('X-Payment-Signature') const authHeader = c.req.header('Authorization') ?? '' let consumerId: string | null = null let paymentMethod: PaymentMethod = 'none' let keyInfo: ApiKeyInfo | null = null + let x402Payload: Record | null = null if (spendAuthHeader) { - const signer = await verifyX402(spendAuthHeader, config.x402, state.nonceStore) + const signer = await verifyX402( + spendAuthHeader, + config.x402, + state.nonceStore, + requiredPaymentAmount, + false, + ) if (!signer) { await state.obs?.onAuthFailure?.(ctx, { method: 'x402', @@ -105,6 +209,8 @@ export async function authenticateAndGuard( message: 'Invalid X-Payment-Signature', type: 'authentication_error', code: 'invalid_spend_auth', + required_amount: requiredPaymentAmount.toString(), + currency_decimals: config.x402.currencyDecimals ?? 6, }, }, { @@ -113,10 +219,17 @@ export async function authenticateAndGuard( }, ) } + x402Payload = JSON.parse(spendAuthHeader) as Record consumerId = signer paymentMethod = 'x402' } else if (isMppAuthEnabled(config) && authHeader.toLowerCase().startsWith('payment ')) { - const signer = await verifyMpp(authHeader, config.mpp!, config.x402, state.nonceStore) + const signer = await verifyMpp( + authHeader, + config.mpp!, + config.x402, + state.nonceStore, + requiredPaymentAmount, + ) if (!signer) { const realm = config.mpp!.realm const method = config.mpp!.method ?? 'blueprintevm' @@ -216,7 +329,9 @@ export async function authenticateAndGuard( operator: config.x402.operatorAddress, chain_id: config.x402.chainId, credits_address: config.x402.creditsAddress, - estimated_amount_per_request: '20000', + required_amount: requiredPaymentAmount.toString(), + currency_decimals: config.x402.currencyDecimals ?? 6, + max_output_tokens: maxOutputTokens, }, ...(isMppAuthEnabled(config) && config.mpp ? { mpp: { realm: config.mpp.realm, method: config.mpp.method ?? 'blueprintevm' } } @@ -270,11 +385,8 @@ export async function authenticateAndGuard( ) } - // Filter consumer messages — strip consumer-side system, length-cap, injection scan. - const { messages: filtered, injectionWarnings } = filterConsumerMessagesStrict( - messages, - state.maxLen, - ) + // Reject or report injection only after authentication so observer events + // retain the authenticated consumer identity. if (injectionWarnings.length > 0) { await state.obs?.onInjectionDetected?.(ctx, { consumerId: consumerId, @@ -294,17 +406,6 @@ export async function authenticateAndGuard( } } - const userMessage = filtered - .filter((m) => m.role === 'user') - .map((m) => m.content) - .join('\n\n') - if (!userMessage) { - return c.json( - { error: { message: 'No user message provided', type: 'invalid_request' } }, - 400, - ) - } - if (config.authorizeConsumer) { const authz = await config.authorizeConsumer(agent, { method: paymentMethod, @@ -326,6 +427,48 @@ export async function authenticateAndGuard( } } + if (paymentMethod === 'x402' && x402Payload) { + try { + if (config.x402.authorizePayment) { + const authorized = await config.x402.authorizePayment(x402Payload, { + requestId, + agentId: agent.id, + requiredAmount: requiredPaymentAmount, + maxOutputTokens, + }) + if (!authorized) throw new Error('payment authorization was rejected') + } + + const nonce = BigInt(String(x402Payload.nonce)) + const expiry = BigInt(String(x402Payload.expiry)) + const nonceKey = `${String(x402Payload.commitment)}:${nonce.toString()}` + if (!config.x402.authorizePayment && await state.nonceStore.hasSeen(nonceKey)) { + throw new Error('payment nonce was already consumed') + } + const ttl = Math.min(Number(expiry) - Math.floor(Date.now() / 1000), 3600) + await state.nonceStore.markSeen(nonceKey, Math.max(ttl, 60)) + } catch { + await state.obs?.onAuthFailure?.(ctx, { + method: 'x402', + code: 'payment_authorization_failed', + httpStatus: 402, + }) + return c.json( + { + error: { + message: 'Payment authorization failed', + type: 'payment_required', + code: 'payment_authorization_failed', + }, + }, + { + status: 402, + headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': requestId }, + }, + ) + } + } + return { agent, consumerId, @@ -335,6 +478,7 @@ export async function authenticateAndGuard( rateLimitRemaining: rl.remaining, requestId, startMs, + maxOutputTokens, } } @@ -353,6 +497,7 @@ export async function* dispatchSandboxStream( config: GatewayConfig, signal?: AbortSignal, sessionId?: string, + maxOutputTokens?: number, ): AsyncIterable { for await (const event of dispatchSandboxStreamRich( agent, @@ -361,6 +506,7 @@ export async function* dispatchSandboxStream( config, signal, sessionId, + maxOutputTokens, )) { if (event.kind === 'text') yield event.delta } @@ -394,11 +540,19 @@ export async function* dispatchSandboxStreamRich( config: GatewayConfig, signal?: AbortSignal, sessionId?: string, + maxOutputTokens?: number, ): AsyncIterable { const box = await config.getSandbox(agent) + const outputLimit = maxOutputTokens ?? config.defaultOutputTokens ?? 1024 + if (!Number.isSafeInteger(outputLimit) || outputLimit <= 0) { + throw new Error('max output tokens must be a positive safe integer') + } + let outputCharacters = 0 + const maxOutputCharacters = outputLimit * 4 const promptStream = box.streamPrompt(userMessage, { sessionId: sessionId ?? `consumer:${consumerId}`, systemPrompt: agent.systemPrompt, + maxOutputTokens: outputLimit, }) for await (const event of promptStream) { if (signal?.aborted) return @@ -407,10 +561,15 @@ export async function* dispatchSandboxStreamRich( event.data?.part?.type === 'text' && event.data.delta ) { + const remainingCharacters = maxOutputCharacters - outputCharacters + if (remainingCharacters <= 0) return + const boundedDelta = event.data.delta.slice(0, remainingCharacters) + outputCharacters += boundedDelta.length yield { kind: 'text', - delta: redactSystemPromptFromOutput(event.data.delta, agent.systemPrompt), + delta: redactSystemPromptFromOutput(boundedDelta, agent.systemPrompt), } + if (boundedDelta.length < event.data.delta.length) return continue } if (event.type === 'input-required' || event.data?.inputRequired) { @@ -484,3 +643,14 @@ export async function settleAndRecord( export function estimateTokens(text: string): number { return Math.ceil(text.length / 4) } + +/** Include the host-owned system prompt because the provider bills it too. */ +export function estimateBillableInputTokens(agent: AgentMeta, userMessage: string): number { + return estimateTokens(userMessage) + estimateTokens(agent.systemPrompt ?? '') +} + +/** A tokenizer cannot emit more tokens than the UTF-8 bytes it consumes. */ +export function maximumBillableInputTokens(agent: AgentMeta, userMessage: string): number { + const encoder = new TextEncoder() + return encoder.encode(userMessage).byteLength + encoder.encode(agent.systemPrompt ?? '').byteLength +} diff --git a/src/middleware.ts b/src/middleware.ts index fc7a4de..8155579 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -7,6 +7,7 @@ import { type GatewayState, authenticateAndGuard, dispatchSandboxStream, + estimateBillableInputTokens, estimateTokens, settleAndRecord, } from './dispatch' @@ -35,6 +36,28 @@ export function createAgentGateway(config: GatewayConfig) { 'For tests, set x402.demoMode: true explicitly.', ) } + const maxOutputTokens = config.maxOutputTokens ?? 4096 + const defaultOutputTokens = config.defaultOutputTokens ?? 1024 + if (!Number.isInteger(maxOutputTokens) || maxOutputTokens <= 0) { + throw new Error('createAgentGateway: maxOutputTokens must be a positive integer') + } + if ( + !Number.isInteger(defaultOutputTokens) || + defaultOutputTokens <= 0 || + defaultOutputTokens > maxOutputTokens + ) { + throw new Error( + 'createAgentGateway: defaultOutputTokens must be a positive integer no greater than maxOutputTokens', + ) + } + if ( + config.x402.currencyDecimals !== undefined && + (!Number.isInteger(config.x402.currencyDecimals) || + config.x402.currencyDecimals < 0 || + config.x402.currencyDecimals > 18) + ) { + throw new Error('createAgentGateway: x402.currencyDecimals must be an integer between 0 and 18') + } const gw = new Hono() const rateLimitStore: RateLimitStore = config.rateLimitStore ?? new MemoryRateLimitStore() const state: GatewayState = { @@ -43,6 +66,8 @@ export function createAgentGateway(config: GatewayConfig) { globalRateLimit: config.rateLimit ?? { limit: 60, windowSeconds: 60 }, requiredScope: config.requiredScope ?? 'chat', maxLen: config.maxMessageLength ?? 8000, + maxOutputTokens, + defaultOutputTokens, obs: config.observer, } const obs: GatewayObserver | undefined = state.obs @@ -125,7 +150,14 @@ export function createAgentGateway(config: GatewayConfig) { ) } - const guard = await authenticateAndGuard(c, slug, body.messages, config, state) + const guard = await authenticateAndGuard( + c, + slug, + body.messages, + config, + state, + body.max_tokens, + ) if (guard instanceof Response) return guard const authz = guard @@ -158,9 +190,17 @@ function streamChatCompletions( config: GatewayConfig, obs: GatewayObserver | undefined, ): Response { - const { agent, consumerId, paymentMethod, requestId, userMessage, rateLimitRemaining } = authz - const inputTokens = estimateTokens(userMessage) - let outputTokens = 0 + const { + agent, + consumerId, + paymentMethod, + requestId, + userMessage, + rateLimitRemaining, + maxOutputTokens, + } = authz + const inputTokens = estimateBillableInputTokens(agent, userMessage) + let outputText = '' const ctx: RequestContext = { requestId, agentSlug: agent.slug, @@ -171,7 +211,7 @@ function streamChatCompletions( async start(controller) { const encoder = new TextEncoder() const sendChunk = (delta: string) => { - outputTokens += estimateTokens(delta) + outputText += delta const chunk: ChatCompletionChunk = { id: `chatcmpl-${Date.now()}`, object: 'chat.completion.chunk', @@ -183,7 +223,15 @@ function streamChatCompletions( } try { - for await (const delta of dispatchSandboxStream(agent, userMessage, consumerId, config)) { + for await (const delta of dispatchSandboxStream( + agent, + userMessage, + consumerId, + config, + undefined, + undefined, + maxOutputTokens, + )) { sendChunk(delta) } @@ -197,7 +245,14 @@ function streamChatCompletions( controller.enqueue(encoder.encode(`data: ${JSON.stringify(done)}\n\n`)) controller.enqueue(encoder.encode('data: [DONE]\n\n')) - await settleAndRecord(agent, authz, inputTokens, outputTokens, config, obs) + await settleAndRecord( + agent, + authz, + inputTokens, + estimateTokens(outputText), + config, + obs, + ) } catch (err) { const rawMessage = err instanceof Error ? err.message : String(err) // Never expose stack traces / absolute paths from sandbox internals. diff --git a/src/types.ts b/src/types.ts index e3f2463..6b74680 100644 --- a/src/types.ts +++ b/src/types.ts @@ -84,8 +84,30 @@ export interface X402Config { rpcUrl?: string /** Demo mode: skip signature verification (default: false). NEVER enable in production. */ demoMode?: boolean - /** Production signer verification. Called with the raw SpendAuth payload. Return true if signature is valid. */ + /** + * Production signature verification. This callback must not reserve, claim, + * or mutate payment state. + */ verifySigner?: (payload: Record) => Promise + /** + * Reserve or claim the verified payment after all request checks pass and + * immediately before sandbox work starts. Return false to reject the call. + */ + authorizePayment?: ( + payload: Record, + context: { + requestId: string + agentId: string + requiredAmount: bigint + maxOutputTokens: number + }, + ) => Promise + /** + * Number of base-unit decimals used by the payment token. Defaults to 6. + * The gateway uses this value to reject a payment that cannot cover the + * request's maximum token charge before it calls `verifySigner`. + */ + currencyDecimals?: number } export interface MppConfig { @@ -178,7 +200,10 @@ export interface SandboxStreamEvent { } export interface SandboxBox { - streamPrompt(message: string, opts?: { sessionId?: string; systemPrompt?: string }): AsyncIterable + streamPrompt( + message: string, + opts?: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number }, + ): AsyncIterable } // --- Gateway config --- @@ -230,6 +255,12 @@ export interface GatewayConfig { /** Max message length in chars (default: 8000) */ maxMessageLength?: number + /** Maximum output token request the gateway accepts. Defaults to 4096. */ + maxOutputTokens?: number + + /** Output token limit used when a request omits `max_tokens`. Defaults to 1024. */ + defaultOutputTokens?: number + /** Required scope for chat endpoint (default: "chat"). API keys must include this scope. */ requiredScope?: string diff --git a/src/verify.ts b/src/verify.ts index ebe19c3..e8450b6 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -36,6 +36,8 @@ export async function verifyX402( spendAuthHeader: string, config: X402Config, nonceStore?: NonceStore, + minimumAmount = 1n, + markNonce = true, ): Promise { try { const raw = JSON.parse(spendAuthHeader) @@ -49,8 +51,10 @@ export async function verifyX402( // Reject expired payments if (expiry < BigInt(Math.floor(Date.now() / 1000))) return null - // Reject zero-amount payments - if (amount <= 0n) return null + // Reject payments that cannot cover the request's maximum charge. The + // check runs before the host verifier because that callback can reserve or + // settle funds as part of its production verification path. + if (amount < minimumAmount || minimumAmount <= 0n) return null const nonceKey = `${raw.commitment}:${nonce.toString()}` if (nonceStore && await nonceStore.hasSeen(nonceKey)) return null @@ -64,7 +68,7 @@ export async function verifyX402( // Check and mark only after the signature is accepted. Otherwise an // invalid request can burn a valid payer nonce and deny the real request. - if (nonceStore) { + if (nonceStore && markNonce) { // Mark seen with TTL matching the expiry window (max 1 hour) const ttl = Math.min(Number(expiry) - Math.floor(Date.now() / 1000), 3600) await nonceStore.markSeen(nonceKey, Math.max(ttl, 60)) @@ -92,6 +96,7 @@ export async function verifyMpp( config: MppConfig, x402Config: X402Config, nonceStore?: NonceStore, + minimumAmount = 1n, ): Promise { // MPP format: "Payment " const match = authHeader.match(/^Payment\s+(\S+)\s+(\S+)$/i) @@ -117,6 +122,27 @@ export async function verifyMpp( // Method-specific verifiers may accept a non-JSON credential format. } + // Validate common EVM fields before a production verifier can reserve or + // settle funds. BlueprinTEVM carries x402-equivalent token amounts and + // must cover the same request ceiling as the X-Payment-Signature path. + const operator = payload.operator ?? payload.to + if (operator !== undefined) { + if (typeof operator !== 'string' || operator.toLowerCase() !== x402Config.operatorAddress.toLowerCase()) { + return null + } + } + const paymentAmount = payload.amount ?? payload.value + if (paymentAmount !== undefined) { + const amount = BigInt(String(paymentAmount)) + if (amount <= 0n || (method === 'blueprintevm' && amount < minimumAmount)) return null + } else if (method === 'blueprintevm') { + return null + } + if (payload.nonce !== undefined) BigInt(String(payload.nonce)) + if (payload.expiry !== undefined && BigInt(String(payload.expiry)) < BigInt(Math.floor(Date.now() / 1000))) { + return null + } + const nonceKey = nonceStore && payload.nonce !== undefined ? `mpp:${method}:${String(payload.commitment ?? payload.from ?? 'unknown')}:${String(payload.nonce)}` @@ -138,20 +164,6 @@ export async function verifyMpp( } if (!consumerId) return null - // Validate common EVM fields when present. Method-specific verifiers own - // the complete credential contract for non-EVM methods. - const operator = payload.operator ?? payload.to - if (operator !== undefined) { - if (typeof operator !== 'string' || operator.toLowerCase() !== x402Config.operatorAddress.toLowerCase()) { - return null - } - } - if (payload.amount !== undefined && BigInt(String(payload.amount)) <= 0n) return null - if (payload.nonce !== undefined) BigInt(String(payload.nonce)) - if (payload.expiry !== undefined && BigInt(String(payload.expiry)) < BigInt(Math.floor(Date.now() / 1000))) { - return null - } - if (nonceStore && payload.nonce !== undefined) { const expiry = payload.expiry === undefined ? Math.floor(Date.now() / 1000) + 3600 diff --git a/tests/integration-x402.test.ts b/tests/integration-x402.test.ts index 62e2a66..c89b4e5 100644 --- a/tests/integration-x402.test.ts +++ b/tests/integration-x402.test.ts @@ -48,6 +48,7 @@ const CHAIN_ID = 3799 const OPERATOR_PRIVATE_KEY = generatePrivateKey() const OPERATOR_ADDRESS = privateKeyToAccount(OPERATOR_PRIVATE_KEY).address const CREDITS_ADDRESS: Hex = '0x00000000000000000000000000000000DeaDBeef' +const FUNDED_REQUEST_AMOUNT = 100_000n const domain = { name: 'ShieldedCredits', @@ -251,7 +252,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo it('happy path: consumer signs → gateway verifies signer address → sandbox streams → settlement fires', async () => { const spendAuth = await signSpendAuth({ consumerPrivateKey: harness.consumerPrivateKey, - amount: 20000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 1n, }) @@ -304,7 +305,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo primaryType: 'SpendAuth', message: { operator: otherOperator, // not our operator - amount: 20000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 1n, expiry: BigInt(Math.floor(Date.now() / 1000) + 600), }, @@ -314,7 +315,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo commitment: account.address, signature: tampered, operator: otherOperator, - amount: '20000', + amount: FUNDED_REQUEST_AMOUNT.toString(), nonce: '1', expiry: String(Math.floor(Date.now() / 1000) + 600), } @@ -345,7 +346,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo primaryType: 'SpendAuth', message: { operator: OPERATOR_ADDRESS, - amount: 20000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 2n, expiry: BigInt(Math.floor(Date.now() / 1000) + 600), }, @@ -356,7 +357,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo commitment: '0xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA', signature, operator: OPERATOR_ADDRESS, - amount: '20000', + amount: FUNDED_REQUEST_AMOUNT.toString(), nonce: '2', expiry: String(Math.floor(Date.now() / 1000) + 600), } @@ -379,7 +380,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo it('replay of the same signed SpendAuth is rejected across requests — regression: double-spend of one signed payment', async () => { const spendAuth = await signSpendAuth({ consumerPrivateKey: harness.consumerPrivateKey, - amount: 20000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 99n, }) const payloadStr = JSON.stringify(spendAuth) @@ -411,7 +412,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo it('expired SpendAuth is rejected — regression: forever-valid sigs enable drain attacks', async () => { const spendAuth = await signSpendAuth({ consumerPrivateKey: harness.consumerPrivateKey, - amount: 20000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 5n, expirySeconds: -10, // expired 10 seconds ago }) @@ -435,7 +436,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo // Alice's wallet const aliceAuth = await signSpendAuth({ consumerPrivateKey: alice.consumerPrivateKey, - amount: 10000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 1n, }) @@ -444,7 +445,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo const bobAddr = privateKeyToAccount(bobKey).address const bobAuth = await signSpendAuth({ consumerPrivateKey: bobKey, - amount: 10000n, + amount: FUNDED_REQUEST_AMOUNT, nonce: 1n, // same nonce, different commitment }) diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index 1b7e338..fc74e94 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -20,16 +20,17 @@ import { MemoryNonceStore } from '../src/nonce-store' import { MemoryRateLimitStore } from '../src/rate-limit' const operatorAddress = '0x1111111111111111111111111111111111111111' +const fundedRequestAmount = '100000' /** Sandbox that emits a fixed reply, captures the prompt + opts for assertion */ class StubSandbox implements SandboxBox { receivedPrompt: string | null = null - receivedOpts: { sessionId?: string; systemPrompt?: string } | undefined + receivedOpts: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number } | undefined constructor(private chunks: string[]) {} async *streamPrompt( message: string, - opts?: { sessionId?: string; systemPrompt?: string }, + opts?: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number }, ): AsyncIterable { this.receivedPrompt = message this.receivedOpts = opts @@ -137,7 +138,7 @@ function buildSpendAuth(overrides: Record = {}): string { return JSON.stringify({ commitment: '0xCommitmentAlice', signature: '0xSignatureBytes', - amount: '20000', + amount: fundedRequestAmount, nonce: String(Math.floor(Math.random() * 1e9)), operator: operatorAddress, expiry: String(now + 600), @@ -249,6 +250,84 @@ describe('POST /:slug/chat/completions — auth paths', () => { const body = await res.json() as { error: { payment_methods: string[]; x402: Record } } expect(body.error.payment_methods).toContain('x402') expect(body.error.x402.operator).toBe(operatorAddress) + expect(body.error.x402.required_amount).toBe('21020') + expect(body.error.x402.max_output_tokens).toBe(1024) + }) + + it('rejects an underfunded payment before the production verifier can reserve funds', async () => { + let verifierCalls = 0 + const { app } = buildHarness({ + x402: { + operatorAddress, + chainId: 3799, + verifySigner: async () => { + verifierCalls += 1 + return true + }, + }, + }) + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ amount: '21019' }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + + expect(res.status).toBe(402) + expect(verifierCalls).toBe(0) + }) + + it('enforces the paid max_tokens limit on the actual sandbox stream', async () => { + const { app, sandbox, usage } = buildHarness({}, ['abcdefgh', 'ijklmnop']) + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ amount: '580' }), + }, + body: JSON.stringify({ + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 2, + }), + }) + + expect(res.status).toBe(200) + const streamed = await readSse(res) + expect(streamed.combinedText).toBe('abcdefgh') + expect(sandbox.receivedOpts?.maxOutputTokens).toBe(2) + expect(usage[0]?.outputTokens).toBe(2) + }) + + it('rejects max_tokens above the configured ceiling before verification', async () => { + let verifierCalls = 0 + const { app } = buildHarness({ + maxOutputTokens: 8, + defaultOutputTokens: 4, + x402: { + operatorAddress, + chainId: 3799, + verifySigner: async () => { + verifierCalls += 1 + return true + }, + }, + }) + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth(), + }, + body: JSON.stringify({ + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 9, + }), + }) + + expect(res.status).toBe(400) + expect(verifierCalls).toBe(0) }) it('returns 402 with invalid_spend_auth on bad X-Payment-Signature — regression: silent bypass of failed sig', async () => { @@ -569,6 +648,63 @@ describe('POST /:slug/chat/completions — authorizeConsumer hook', () => { }) expect(getSandboxCalls).toBe(0) }) + + it('does not reserve x402 funds until every request check allows the call', async () => { + let paymentAuthorizations = 0 + const { app } = buildHarness({ + x402: { + operatorAddress, + chainId: 3799, + verifySigner: async () => true, + authorizePayment: async () => { + paymentAuthorizations += 1 + return true + }, + }, + authorizeConsumer: async () => ({ allow: false, reason: 'no', code: 'denied' }), + }) + + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ nonce: '5004' }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + + expect(res.status).toBe(403) + expect(paymentAuthorizations).toBe(0) + }) + + it('reserves one x402 payment immediately before allowed sandbox work', async () => { + let paymentAuthorizations = 0 + const { app } = buildHarness({ + x402: { + operatorAddress, + chainId: 3799, + verifySigner: async () => true, + authorizePayment: async () => { + paymentAuthorizations += 1 + return true + }, + }, + authorizeConsumer: async () => ({ allow: true }), + }) + + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ nonce: '5005' }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + + expect(res.status).toBe(200) + await readSse(res) + expect(paymentAuthorizations).toBe(1) + }) }) describe('createAgentGateway — production-config guard', () => { diff --git a/tests/observer.test.ts b/tests/observer.test.ts index 594aeab..ef539cd 100644 --- a/tests/observer.test.ts +++ b/tests/observer.test.ts @@ -140,7 +140,7 @@ function buildSpendAuth(overrides: Record = {}): string { return JSON.stringify({ commitment: '0xAlice', signature: '0xsig', - amount: '20000', + amount: '100000', nonce: String(Math.floor(Math.random() * 1e9)), operator: operatorAddress, expiry: String(now + 600), diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 80efb3e..7ceb70a 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -49,6 +49,21 @@ describe('verifyX402', () => { expect(await verifyX402(buildSpendAuth({ amount: '0' }), baseConfig)).toBeNull() }) + it('rejects a request-specific underpayment before calling the production verifier', async () => { + let calls = 0 + const config: X402Config = { + ...baseConfig, + demoMode: false, + verifySigner: async () => { + calls += 1 + return true + }, + } + + expect(await verifyX402(buildSpendAuth({ amount: '19999' }), config, undefined, 20000n)).toBeNull() + expect(calls).toBe(0) + }) + it('rejects expired payments — regression: forever-valid sigs enable drained-wallet attacks', async () => { const expired = buildSpendAuth({ expiry: String(Math.floor(Date.now() / 1000) - 10) }) expect(await verifyX402(expired, baseConfig)).toBeNull() @@ -156,6 +171,27 @@ describe('verifyMpp', () => { expect(seen[0].credential).toContain('commitment') }) + it('rejects an underfunded blueprintevm credential before its verifier can reserve funds', async () => { + let calls = 0 + const header = buildCredential({ + commitment: '0xAlice', + operator: operatorAddress, + amount: '19999', + nonce: '8', + expiry: String(Math.floor(Date.now() / 1000) + 600), + }) + const config: MppConfig = { + ...mppConfig, + verifySigner: async () => { + calls += 1 + return 'mpp:alice' + }, + } + + expect(await verifyMpp(header, config, { ...baseConfig, demoMode: false }, undefined, 20000n)).toBeNull() + expect(calls).toBe(0) + }) + it('rejects MPP in production when no method verifier is configured', async () => { const header = buildCredential({ commitment: '0xAlice', operator: operatorAddress }) const productionX402: X402Config = { ...baseConfig, demoMode: false } From 4f62515c5364442913d1a0514343156ef226a101 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 19:50:10 -0600 Subject: [PATCH 02/32] fix(payments): add atomic bounded payment lifecycle --- README.md | 3 + src/a2a/handler.ts | 366 +++++++++++++++++---- src/a2a/task-store-sql.ts | 26 ++ src/a2a/task-store.ts | 19 ++ src/dispatch.ts | 537 ++++++++++++++++++++++++++----- src/index.ts | 14 + src/middleware.ts | 109 ++++++- src/nonce-store.ts | 56 +++- src/payment-operations.ts | 316 ++++++++++++++++++ src/types.ts | 73 ++++- src/verify.ts | 54 +++- tests/a2a-long-horizon.test.ts | 32 ++ tests/a2a-payment-races.test.ts | 183 +++++++++++ tests/a2a.test.ts | 16 + tests/integration-x402.test.ts | 18 +- tests/kv-stores.test.ts | 18 ++ tests/middleware.test.ts | 48 ++- tests/observer.test.ts | 18 +- tests/payment-operations.test.ts | 433 +++++++++++++++++++++++++ tests/protocol-guards.test.ts | 210 ++++++++++++ tests/verify.test.ts | 19 ++ 21 files changed, 2372 insertions(+), 196 deletions(-) create mode 100644 src/payment-operations.ts create mode 100644 tests/a2a-payment-races.test.ts create mode 100644 tests/payment-operations.test.ts create mode 100644 tests/protocol-guards.test.ts diff --git a/README.md b/README.md index b22ee9b..d3f9d89 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,9 @@ app.route('/v1/agents', createAgentGateway({ Set `x402.demoMode: true` only for local development and tests; that explicit mode also enables the built-in `sk_agent_*` demo key verifier. Keep `verifySigner` free of side effects. Use `authorizePayment` to reserve or claim funds after rate limits, content checks, and product authorization succeed. +For production version 2, set `x402.paymentProtocolVersion: 2`, provide `paymentOperations`, and return its operation from `authorizePayment`. +The operation store owns claim, partial settle, release, and expiry reclaim. +Keep version 1 explicitly configured while old and new gateways coexist; shared nonce storage must reject a version 1 claim owned by a version 2 operation. Before it calls the verifier, the gateway requires the signed amount to cover filtered input plus the requested output limit. The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 8ef0a5f..b87be27 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -17,12 +17,14 @@ import { type AuthorizedRequest, type GatewayState, authenticateAndGuard, + claimPayment, dispatchSandboxStreamRich, - estimateBillableInputTokens, - estimateTokens, + releasePayment, + releasePaymentAfterFailure, settleAndRecord, } from '../dispatch' import type { GatewayConfig } from '../types' +import type { SandboxUsageReceipt } from '../types' import { buildAgentCard } from './agent-card' import { fail, ok, parseEnvelope } from './jsonrpc' import { @@ -68,6 +70,7 @@ const TERMINAL_STATES: ReadonlySet = new Set([ */ class CancelRegistry { private readonly controllers = new Map() + private readonly finalizing = new Set() register(taskId: string): AbortController { const c = new AbortController() @@ -77,9 +80,22 @@ class CancelRegistry { clear(taskId: string): void { this.controllers.delete(taskId) + this.finalizing.delete(taskId) + } + + beginFinalization(taskId: string): boolean { + const controller = this.controllers.get(taskId) + if (!controller || controller.signal.aborted || this.finalizing.has(taskId)) return false + this.finalizing.add(taskId) + return true + } + + isFinalizing(taskId: string): boolean { + return this.finalizing.has(taskId) } cancel(taskId: string): boolean { + if (this.finalizing.has(taskId)) return false const c = this.controllers.get(taskId) if (!c) return false c.abort() @@ -168,8 +184,17 @@ async function handleMessageSend( const guard = await guardMessageRequest(c, slug, req, deps) if (guard instanceof Response) return guard const { authz, task } = guard + const workingTask: Task = task.status.state === 'working' + ? task + : { ...task, status: { state: 'working', timestamp: nowIso() } } + if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { + await releaseOwnedPayment(authz, deps, 'A2A task changed before execution started') + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) + } let responseText = '' + let usage: SandboxUsageReceipt | undefined + let workObserved = false let inputRequiredPrompt: string | undefined let inputRequiredSeen = false try { @@ -184,15 +209,37 @@ async function handleMessageSend( )) { if (event.kind === 'text') { responseText += event.delta + workObserved = true + } else if (event.kind === 'activity') { + workObserved = true + } else if (event.kind === 'usage') { + usage = event.usage } else { inputRequiredSeen = true inputRequiredPrompt = event.prompt + workObserved = true } } } catch (err) { - const failed = withStatus(task, 'failed') - await deps.taskStore.put(failed) - await maybeDeliverPush(failed, deps) + await releaseOrRetainPayment( + authz, + deps, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, + ) + const currentTask = await deps.taskStore.get(task.id) + const failed = currentTask && isTerminal(currentTask.status.state) + ? currentTask + : withStatus(workingTask, 'failed') + try { + await deps.taskStore.put(failed) + await maybeDeliverPush(failed, deps) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } return c.json( fail( req.id, @@ -205,18 +252,38 @@ async function handleMessageSend( // Settle for the work done so far before short-circuiting on input-required. // The user has been charged for the partial response, which is the right // commercial behavior — the sandbox produced tokens. - await settleAndRecord( - authz.agent, - authz, - estimateBillableInputTokens(authz.agent, authz.userMessage), - estimateTokens(responseText), - deps.config, - deps.state.obs, - ) + try { + if (!usage) throw new Error('sandbox did not provide a usage receipt') + if (!await claimTaskFinalization(deps.taskStore, workingTask)) { + throw new Error('A2A task changed before payment settlement') + } + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) + } catch (err) { + await releaseOrRetainPayment( + authz, + deps, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, + ) + const currentTask = await deps.taskStore.get(task.id) + const failed = currentTask && isTerminal(currentTask.status.state) + ? currentTask + : withStatus(workingTask, 'failed') + try { + await deps.taskStore.put(failed) + await maybeDeliverPush(failed, deps) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) + } if (inputRequiredSeen) { const paused = withStatus( - task, + workingTask, 'input-required', inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, responseText @@ -228,7 +295,7 @@ async function handleMessageSend( return c.json(ok(req.id, paused)) } - const completed = withStatus(task, 'completed', undefined, [ + const completed = withStatus(workingTask, 'completed', undefined, [ responseTextToArtifact(responseText, `${task.id}-artifact-0`), ]) await deps.taskStore.put(completed) @@ -250,8 +317,9 @@ async function handleMessageStream( const { authz, task } = guard const controller = cancels.register(task.id) - const inputTokens = estimateBillableInputTokens(authz.agent, authz.userMessage) let responseText = '' + let usage: SandboxUsageReceipt | undefined + let workObserved = false const stream = new ReadableStream({ async start(ctrl) { @@ -268,12 +336,17 @@ async function handleMessageStream( status: { state: 'working', timestamp: nowIso() }, final: false, } - await deps.taskStore.put({ ...task, status: workingStatus.status }) - send(workingStatus) - + const workingTask: Task = task.status.state === 'working' + ? task + : { ...task, status: workingStatus.status } let inputRequiredPrompt: string | undefined let inputRequiredSeen = false try { + if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { + throw new Error('A2A task changed before execution started') + } + send(workingStatus) + for await (const event of dispatchSandboxStreamRich( authz.agent, authz.userMessage, @@ -285,6 +358,7 @@ async function handleMessageStream( )) { if (event.kind === 'text') { responseText += event.delta + workObserved = true const artifactEvent: TaskArtifactUpdateEvent = { kind: 'artifact-update', taskId: task.id, @@ -297,38 +371,67 @@ async function handleMessageStream( append: true, } send(artifactEvent) + } else if (event.kind === 'activity') { + workObserved = true + } else if (event.kind === 'usage') { + usage = event.usage } else { inputRequiredSeen = true inputRequiredPrompt = event.prompt + workObserved = true } } - // Caller aborted via tasks/cancel — emit canceled, do not settle. + // Caller aborted via tasks/cancel. Charge a complete receipt if one + // exists; otherwise retain ownership when output or hidden work was + // observed because releasing would make paid work free. if (controller.signal.aborted) { + if (usage) { + try { + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) + } catch (settlementError) { + console.error( + `[a2a] canceled task settlement retained for ${authz.requestId}:`, + settlementError instanceof Error ? settlementError.message : String(settlementError), + ) + } + } else { + await releaseOrRetainPayment(authz, deps, 'a2a task canceled', workObserved || usage !== undefined) + } const canceled = withStatus(task, 'canceled', undefined, [ responseTextToArtifact(responseText, `${task.id}-artifact-0`), ]) - await deps.taskStore.put(canceled) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: canceled.status, - final: true, - }) - await maybeDeliverPush(canceled, deps) + try { + await deps.taskStore.put(canceled) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: canceled.status, + final: true, + }) + await maybeDeliverPush(canceled, deps) + } catch (taskError) { + console.error( + `[a2a] failed to persist canceled task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } return } // Settle once for whatever the sandbox produced (full or partial). - await settleAndRecord( - authz.agent, - authz, - inputTokens, - estimateTokens(responseText), - deps.config, - deps.state.obs, - ) + if (!usage) throw new Error('sandbox did not provide a usage receipt') + if (!cancels.beginFinalization(task.id) || !await claimTaskFinalization(deps.taskStore, workingTask)) { + await releaseOrRetainPayment( + authz, + deps, + 'A2A task changed before payment settlement', + workObserved || usage !== undefined, + ) + return + } + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) if (inputRequiredSeen) { const paused = withStatus( @@ -377,27 +480,50 @@ async function handleMessageStream( }) await maybeDeliverPush(completed, deps) } catch (err) { - const failed = withStatus(task, 'failed') - await deps.taskStore.put(failed) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: failed.status, - final: true, - }) - await maybeDeliverPush(failed, deps) - await deps.state.obs?.onStreamError?.( - { - requestId: authz.requestId, - agentSlug: authz.agent.slug, - startMs: authz.startMs, - }, - { - consumerId: authz.consumerId, - errorMessage: err instanceof Error ? err.message : String(err), - }, + await releaseOrRetainPayment( + authz, + deps, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, ) + const currentTask = await deps.taskStore.get(task.id) + const failed = currentTask && isTerminal(currentTask.status.state) + ? currentTask + : withStatus(task, 'failed') + try { + await deps.taskStore.put(failed) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: failed.status, + final: true, + }) + await maybeDeliverPush(failed, deps) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + try { + await deps.state.obs?.onStreamError?.( + { + requestId: authz.requestId, + agentSlug: authz.agent.slug, + startMs: authz.startMs, + }, + { + consumerId: authz.consumerId, + errorMessage: err instanceof Error ? err.message : String(err), + }, + ) + } catch (observerError) { + console.error( + `[a2a] stream observer failed for ${authz.requestId}:`, + observerError instanceof Error ? observerError.message : String(observerError), + ) + } } finally { cancels.clear(task.id) ctrl.close() @@ -462,12 +588,30 @@ async function handleTasksCancel( ) } - const stillActive = cancels.cancel(task.id) + if (isTaskFinalizing(task) || cancels.isFinalizing(task.id)) { + return c.json( + fail( + req.id, + A2A_ERROR_CODES.TASK_NOT_CANCELABLE, + `task '${task.id}' is being finalized`, + ), + ) + } const canceled: Task = { ...task, status: { state: 'canceled', timestamp: nowIso() }, } - await deps.taskStore.put(canceled) + const transitioned = await compareAndSetTask(deps.taskStore, task, canceled) + if (!transitioned) { + const current = await deps.taskStore.get(task.id) + if (current && (isTerminal(current.status.state) || isTaskFinalizing(current))) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' changed before cancellation`), + ) + } + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'task changed before cancellation')) + } + const stillActive = cancels.cancel(task.id) // If a stream was active, it'll observe the abort and emit its own final // status-update AND fire its own push delivery; the dispatcher only fires @@ -682,7 +826,13 @@ async function guardMessageRequest( status: { state: 'working', timestamp: nowIso() }, history: [...(existing.history ?? []), appendedMessage], } - await deps.taskStore.put(continued) + if (!await compareAndSetTask(deps.taskStore, existing, continued)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${existing.id}' changed before continuation`), + ) + } + const claimError = await claimTaskPayment(c, req, continued, authz, deps) + if (claimError) return claimError return { authz, task: continued } } // Unknown taskId in params: fall through and mint a fresh task with that @@ -703,12 +853,106 @@ async function guardMessageRequest( status: { state: 'submitted', timestamp: nowIso() }, history: [initialMessage], } - await deps.taskStore.put(task) + if (!await createTask(deps.taskStore, task)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' already exists`), + ) + } + const claimError = await claimTaskPayment(c, req, task, authz, deps) + if (claimError) return claimError return { authz, task } } +async function claimTaskPayment( + c: Context, + req: JSONRPCRequest, + task: Task, + authz: AuthorizedRequest, + deps: A2AHandlerDeps, +): Promise { + try { + await claimPayment(authz, deps.config, deps.state) + return undefined + } catch { + await releaseOwnedPayment(authz, deps, 'payment authorization failed') + const failed = withStatus(task, 'failed') + try { + await deps.taskStore.put(failed) + await maybeDeliverPush(failed, deps) + } catch (taskError) { + console.error( + `[a2a] failed to persist payment-failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment authorization failed')) + } +} + +async function releaseOwnedPayment( + authz: AuthorizedRequest, + deps: A2AHandlerDeps, + reason: string, +): Promise { + try { + await releasePayment(authz, deps.config, reason) + } catch (releaseError) { + console.error( + `[a2a] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } +} + +async function releaseOrRetainPayment( + authz: AuthorizedRequest, + deps: A2AHandlerDeps, + reason: string, + workObserved: boolean, +): Promise { + try { + await releasePaymentAfterFailure(authz, deps.config, reason, workObserved) + } catch (releaseError) { + console.error( + `[a2a] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } +} + // ── Helpers ─────────────────────────────────────────────────────────────── +const FINALIZING_METADATA_KEY = 'gatewayFinalizing' + +async function createTask(taskStore: TaskStore, task: Task): Promise { + if (taskStore.createIfAbsent) return taskStore.createIfAbsent(task) + if (await taskStore.get(task.id)) return false + await taskStore.put(task) + return true +} + +async function compareAndSetTask(taskStore: TaskStore, expected: Task, next: Task): Promise { + if (taskStore.compareAndSet) return taskStore.compareAndSet(expected, next) + // Older adapters remain source-compatible, but their fallback is only + // process-safe. Durable adapters must implement compareAndSet. + const current = await taskStore.get(expected.id) + if (!current || JSON.stringify(current) !== JSON.stringify(expected)) return false + await taskStore.put(next) + return true +} + +async function claimTaskFinalization(taskStore: TaskStore, task: Task): Promise { + if (isTaskFinalizing(task)) return false + return compareAndSetTask(taskStore, task, { + ...task, + metadata: { ...(task.metadata ?? {}), [FINALIZING_METADATA_KEY]: true }, + }) +} + +function isTaskFinalizing(task: Task): boolean { + return task.metadata?.[FINALIZING_METADATA_KEY] === true +} + function isTerminal(state: Task['status']['state']): boolean { return ( state === 'completed' || diff --git a/src/a2a/task-store-sql.ts b/src/a2a/task-store-sql.ts index 5f941d5..857749a 100644 --- a/src/a2a/task-store-sql.ts +++ b/src/a2a/task-store-sql.ts @@ -166,6 +166,32 @@ export class SqlTaskStore implements TaskStore { } } + async createIfAbsent(task: Task): Promise { + const payload = JSON.stringify(task) + try { + const result = await this.db.exec( + `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at) VALUES (?, ?, ?, ?, ?)`, + [task.id, task.contextId, task.status.state, payload, Date.now()], + ) + return result.rowsAffected === 1 + } catch (error) { + // SQL dialects report duplicate primary keys as errors. Convert only a + // confirmed existing row into the protocol-level "already exists" result. + if (await this.get(task.id)) return false + throw error + } + } + + async compareAndSet(expected: Task, next: Task): Promise { + const expectedPayload = JSON.stringify(expected) + const payload = JSON.stringify(next) + const result = await this.db.exec( + `UPDATE ${this.table} SET context_id = ?, state = ?, payload = ?, updated_at = ? WHERE id = ? AND payload = ?`, + [next.contextId, next.status.state, payload, Date.now(), expected.id, expectedPayload], + ) + return result.rowsAffected === 1 + } + async delete(id: string): Promise { await this.db.exec(`DELETE FROM ${this.table} WHERE id = ?`, [id]) } diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index 0c24b3d..f320491 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -10,6 +10,10 @@ import type { Task } from './types' export interface TaskStore { get(id: string): Promise put(task: Task): Promise + /** Insert only when the task id is absent. Required for explicit A2A ids. */ + createIfAbsent?(task: Task): Promise + /** Replace only when the stored task still equals `expected`. */ + compareAndSet?(expected: Task, next: Task): Promise delete(id: string): Promise } @@ -32,6 +36,21 @@ export class InMemoryTaskStore implements TaskStore { this.entries.set(task.id, { task: clone(task), expiresAt: Date.now() + this.ttlMs }) } + async createIfAbsent(task: Task): Promise { + this.gc() + if (this.entries.has(task.id)) return false + this.entries.set(task.id, { task: clone(task), expiresAt: Date.now() + this.ttlMs }) + return true + } + + async compareAndSet(expected: Task, next: Task): Promise { + this.gc() + const entry = this.entries.get(expected.id) + if (!entry || JSON.stringify(entry.task) !== JSON.stringify(expected)) return false + this.entries.set(expected.id, { task: clone(next), expiresAt: Date.now() + this.ttlMs }) + return true + } + async delete(id: string): Promise { this.entries.delete(id) } diff --git a/src/dispatch.ts b/src/dispatch.ts index aaba8a3..737761f 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -14,17 +14,21 @@ import { filterConsumerMessagesStrict, redactSystemPromptFromOutput } from './fi import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' import { type RateLimitStore, checkRateLimit } from './rate-limit' import type { NonceStore } from './nonce-store' +import type { PaymentOperation } from './payment-operations' import type { AgentMeta, ApiKeyInfo, ChatMessage, GatewayConfig, PaymentMethod, + SandboxExecutionBudget, + SandboxUsageReceipt, } from './types' import { defaultVerifyApiKey, isApiKeyAuthEnabled, isMppAuthEnabled, + mppReplayNonceKey, verifyMpp, verifyX402, } from './verify' @@ -38,6 +42,10 @@ export interface GatewayState { maxLen: number maxOutputTokens: number defaultOutputTokens: number + maxReasoningTokens: number + maxToolTokens: number + maxToolCalls: number + maxProviderCostUsd?: number obs?: GatewayObserver } @@ -52,6 +60,11 @@ export interface AuthorizedRequest { requestId: string startMs: number maxOutputTokens: number + executionBudget: SandboxExecutionBudget + requiredPaymentAmount: bigint + paymentPayload: Record | null + paymentNonceKey?: string + paymentOperation?: PaymentOperation } function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { @@ -76,6 +89,9 @@ export function requiredX402Amount( inputTokens: number, maxOutputTokens: number, currencyDecimals = 6, + maxReasoningTokens = 0, + maxToolTokens = 0, + maxProviderCostUsd = 0, ): bigint { if (!Number.isSafeInteger(inputTokens) || inputTokens < 0) { throw new Error('input token estimate must be a non-negative safe integer') @@ -86,9 +102,27 @@ export function requiredX402Amount( if (!Number.isInteger(currencyDecimals) || currencyDecimals < 0 || currencyDecimals > 18) { throw new Error('x402 currencyDecimals must be an integer between 0 and 18') } + for (const [name, value] of [ + ['maxReasoningTokens', maxReasoningTokens], + ['maxToolTokens', maxToolTokens], + ] as const) { + if (!Number.isSafeInteger(value) || value < 0) throw new Error(`${name} must be a non-negative safe integer`) + } + if (!Number.isFinite(maxProviderCostUsd) || maxProviderCostUsd < 0) { + throw new Error('maxProviderCostUsd must be finite and non-negative') + } const { numerator, denominator } = decimalFraction(pricePerTokenUsd) - const scaled = BigInt(inputTokens + maxOutputTokens) * numerator * 10n ** BigInt(currencyDecimals) - return (scaled + denominator - 1n) / denominator + const tokenCount = inputTokens + maxOutputTokens + maxReasoningTokens + maxToolTokens + if (!Number.isSafeInteger(tokenCount)) throw new Error('token budget exceeds safe integer range') + const tokenScaled = BigInt(tokenCount) * + numerator * 10n ** BigInt(currencyDecimals) + const tokenAmount = (tokenScaled + denominator - 1n) / denominator + const provider = maxProviderCostUsd === 0 + ? { numerator: 0n, denominator: 1n } + : decimalFraction(maxProviderCostUsd) + const providerScaled = provider.numerator * 10n ** BigInt(currencyDecimals) + const providerAmount = (providerScaled + provider.denominator - 1n) / provider.denominator + return tokenAmount > providerAmount ? tokenAmount : providerAmount } /** @@ -143,8 +177,8 @@ export async function authenticateAndGuard( ) } - // Price the exact filtered input that reaches the sandbox. This must happen - // before x402 verification because a production verifier can reserve funds. + // Quote the maximum UTF-8 input plus every hidden provider cost before + // verification. The verifier must remain read-only at this point. const { messages: filtered, injectionWarnings } = filterConsumerMessagesStrict( messages, state.maxLen, @@ -161,12 +195,29 @@ export async function authenticateAndGuard( } let requiredPaymentAmount: bigint + const maxInputTokens = maximumBillableInputTokens(agent, userMessage) + const maxReasoningTokens = state.maxReasoningTokens + const maxToolTokens = state.maxToolTokens + const maxToolCalls = state.maxToolCalls + const maxProviderCostUsd = state.maxProviderCostUsd ?? + (maxInputTokens + maxOutputTokens + maxReasoningTokens + maxToolTokens) * agent.pricePerTokenUsd + const executionBudget: SandboxExecutionBudget = { + maxInputTokens, + maxOutputTokens, + maxReasoningTokens, + maxToolTokens, + maxToolCalls, + maxProviderCostUsd, + } try { requiredPaymentAmount = requiredX402Amount( agent.pricePerTokenUsd, - maximumBillableInputTokens(agent, userMessage), + maxInputTokens, maxOutputTokens, config.x402.currencyDecimals, + maxReasoningTokens, + maxToolTokens, + maxProviderCostUsd, ) } catch { return c.json( @@ -188,6 +239,7 @@ export async function authenticateAndGuard( let paymentMethod: PaymentMethod = 'none' let keyInfo: ApiKeyInfo | null = null let x402Payload: Record | null = null + let paymentNonceKey: string | undefined if (spendAuthHeader) { const signer = await verifyX402( @@ -220,6 +272,7 @@ export async function authenticateAndGuard( ) } x402Payload = JSON.parse(spendAuthHeader) as Record + paymentNonceKey = `${String(x402Payload.commitment).toLowerCase()}:${BigInt(String(x402Payload.nonce)).toString()}` consumerId = signer paymentMethod = 'x402' } else if (isMppAuthEnabled(config) && authHeader.toLowerCase().startsWith('payment ')) { @@ -229,6 +282,7 @@ export async function authenticateAndGuard( config.x402, state.nonceStore, requiredPaymentAmount, + false, ) if (!signer) { const realm = config.mpp!.realm @@ -257,6 +311,7 @@ export async function authenticateAndGuard( } consumerId = signer paymentMethod = 'mpp' + paymentNonceKey = mppReplayNonceKey(authHeader) } else if (authHeader.startsWith('Bearer ')) { const verify = config.verifyApiKey ?? (config.x402.demoMode ? defaultVerifyApiKey : null) if (!verify || !isApiKeyAuthEnabled(config)) { @@ -351,12 +406,6 @@ export async function authenticateAndGuard( ) } - await state.obs?.onPaymentVerified?.(ctx, { - method: paymentMethod, - consumerId: consumerId, - keyId: keyInfo?.keyId, - }) - // Rate limit. const effectiveRateLimit = keyInfo?.rateLimitPerMinute ? { limit: keyInfo.rateLimitPerMinute, windowSeconds: 60 } @@ -427,48 +476,6 @@ export async function authenticateAndGuard( } } - if (paymentMethod === 'x402' && x402Payload) { - try { - if (config.x402.authorizePayment) { - const authorized = await config.x402.authorizePayment(x402Payload, { - requestId, - agentId: agent.id, - requiredAmount: requiredPaymentAmount, - maxOutputTokens, - }) - if (!authorized) throw new Error('payment authorization was rejected') - } - - const nonce = BigInt(String(x402Payload.nonce)) - const expiry = BigInt(String(x402Payload.expiry)) - const nonceKey = `${String(x402Payload.commitment)}:${nonce.toString()}` - if (!config.x402.authorizePayment && await state.nonceStore.hasSeen(nonceKey)) { - throw new Error('payment nonce was already consumed') - } - const ttl = Math.min(Number(expiry) - Math.floor(Date.now() / 1000), 3600) - await state.nonceStore.markSeen(nonceKey, Math.max(ttl, 60)) - } catch { - await state.obs?.onAuthFailure?.(ctx, { - method: 'x402', - code: 'payment_authorization_failed', - httpStatus: 402, - }) - return c.json( - { - error: { - message: 'Payment authorization failed', - type: 'payment_required', - code: 'payment_authorization_failed', - }, - }, - { - status: 402, - headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': requestId }, - }, - ) - } - } - return { agent, consumerId, @@ -479,7 +486,150 @@ export async function authenticateAndGuard( requestId, startMs, maxOutputTokens, + executionBudget, + requiredPaymentAmount, + paymentPayload: x402Payload, + paymentNonceKey, + } +} + +/** Claim payment ownership after every request guard has accepted the call. */ +export async function claimPayment( + authz: AuthorizedRequest, + config: GatewayConfig, + state: GatewayState, +): Promise { + if (authz.paymentMethod === 'x402' && authz.paymentPayload) { + const context = { + requestId: authz.requestId, + agentId: authz.agent.id, + requiredAmount: authz.requiredPaymentAmount, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + } + let operation: PaymentOperation | undefined + if (config.x402.authorizePayment) { + const result = await config.x402.authorizePayment(authz.paymentPayload, context) + if (!result) throw new Error('payment authorization was rejected') + if (typeof result !== 'boolean') { + operation = result + // Retain the durable owner even if a later compatibility check fails. + // The caller can then release or reclaim the operation instead of + // losing its recovery handle. + authz.paymentOperation = operation + } + else if (config.x402.paymentProtocolVersion === 2) { + throw new Error('version 2 payment authorization did not return an operation') + } else if (authz.paymentNonceKey) { + const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + if (!claimed) throw new Error('payment nonce was already consumed') + } + } else if (config.x402.paymentOperations) { + operation = await config.x402.paymentOperations.claimPayment(authz.paymentPayload, context) + } else if (authz.paymentNonceKey) { + const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + if (!claimed) throw new Error('payment nonce was already consumed') + } + if (operation && operation.protocolVersion !== 2) { + throw new Error('payment operation protocol version mismatch') + } + if (operation && !config.x402.paymentOperations) { + throw new Error('durable payment operations are required to settle a claimed operation') + } + if (operation && authz.paymentNonceKey) { + const claimed = await claimNonce( + state.nonceStore, + authz.paymentNonceKey, + authz.paymentPayload, + operation.operationId, + ) + if (!claimed) { + try { + await config.x402.paymentOperations!.releasePayment(operation, 'shared payment nonce was already owned') + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } + throw new Error('payment nonce was already consumed') + } + } + authz.paymentOperation = operation + } else if (authz.paymentMethod === 'mpp' && authz.paymentNonceKey) { + const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) + if (!claimed) throw new Error('payment nonce was already consumed') + } + + try { + await state.obs?.onPaymentVerified?.( + { + requestId: authz.requestId, + agentSlug: authz.agent.slug, + startMs: authz.startMs, + }, + { + method: authz.paymentMethod, + consumerId: authz.consumerId, + keyId: authz.keyInfo?.keyId, + }, + ) + } catch (error) { + // Observability must not turn a durable claim into a stranded payment. + console.error( + '[agent-gateway] payment observer failed for ' + authz.requestId + ':', + error instanceof Error ? error.message : String(error), + ) + } +} + +/** Release an owned operation when execution cannot produce a valid receipt. */ +export async function releasePayment( + authz: AuthorizedRequest, + config: GatewayConfig, + reason: string, +): Promise { + if (!authz.paymentOperation || !config.x402.paymentOperations) return + await config.x402.paymentOperations.releasePayment(authz.paymentOperation, reason) +} + +/** + * Release only when no sandbox work was observed. Once output or a receipt + * exists, retain the owner for settlement or background recovery. + */ +export async function releasePaymentAfterFailure( + authz: AuthorizedRequest, + config: GatewayConfig, + reason: string, + workObserved: boolean, +): Promise { + if (workObserved) { + console.error( + `[agent-gateway] retaining payment ownership after sandbox work for ${authz.requestId}: ${reason}`, + ) + return } + await releasePayment(authz, config, reason) +} + +export async function reclaimPayment( + operationId: string, + config: GatewayConfig, +): Promise { + if (!config.x402.paymentOperations) throw new Error('durable payment operations are not configured') + return config.x402.paymentOperations.reclaimPayment(operationId) +} + +async function claimNonce( + nonceStore: NonceStore, + nonceKey: string, + payload: Record, + ownerId?: string, +): Promise { + const expiry = Number(payload.expiry ?? Math.floor(Date.now() / 1000) + 3600) + const ttl = Math.min(expiry - Math.floor(Date.now() / 1000), 3600) + if (!Number.isFinite(ttl) || ttl <= 0) return false + return nonceStore.claim(nonceKey, Math.max(ttl, 60), ownerId) } /** @@ -521,6 +671,8 @@ export async function* dispatchSandboxStream( export type A2ADispatchEvent = | { kind: 'text'; delta: string } | { kind: 'input-required'; prompt?: string } + | { kind: 'activity' } + | { kind: 'usage'; usage: SandboxUsageReceipt } /** * Like `dispatchSandboxStream` but yields a discriminated union so callers can @@ -547,38 +699,96 @@ export async function* dispatchSandboxStreamRich( if (!Number.isSafeInteger(outputLimit) || outputLimit <= 0) { throw new Error('max output tokens must be a positive safe integer') } - let outputCharacters = 0 - const maxOutputCharacters = outputLimit * 4 + let outputBytes = 0 + // Bound untrusted adapters while the final receipt is pending. The receipt + // remains authoritative for token count, so over-limit output is never sent. + const maxOutputBytes = outputLimit * 4 + if (!Number.isSafeInteger(maxOutputBytes)) { + throw new Error('max output token bound exceeds safe integer range') + } + const encoder = new TextEncoder() + let usageParts: Partial = {} + let observedReasoningTokens = 0 + let observedToolTokens = 0 + let observedToolCalls = 0 + const bufferedTextParts: string[] = [] + const executionBudget: SandboxExecutionBudget = { + 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 ?? + (maximumBillableInputTokens(agent, userMessage) + outputLimit + + (config.executionBudget?.maxReasoningTokens ?? outputLimit) + + (config.executionBudget?.maxToolTokens ?? outputLimit)) * agent.pricePerTokenUsd, + } const promptStream = box.streamPrompt(userMessage, { sessionId: sessionId ?? `consumer:${consumerId}`, systemPrompt: agent.systemPrompt, maxOutputTokens: outputLimit, + executionBudget, }) for await (const event of promptStream) { if (signal?.aborted) return + if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) + if (event.data?.reasoning?.tokens !== undefined) { + observedReasoningTokens += nonNegativeSafeInteger(event.data.reasoning.tokens, 'reasoning tokens') + yield { kind: 'activity' } + } + if (event.data?.tool) { + observedToolCalls += 1 + observedToolTokens += + nonNegativeSafeInteger(event.data.tool.inputTokens ?? 0, 'tool input tokens') + + nonNegativeSafeInteger(event.data.tool.outputTokens ?? 0, 'tool output tokens') + yield { kind: 'activity' } + } + enforceUsageBudget(withObservedUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + ), executionBudget) if ( event.type === 'message.part.updated' && event.data?.part?.type === 'text' && event.data.delta ) { - const remainingCharacters = maxOutputCharacters - outputCharacters - if (remainingCharacters <= 0) return - const boundedDelta = event.data.delta.slice(0, remainingCharacters) - outputCharacters += boundedDelta.length - yield { - kind: 'text', - delta: redactSystemPromptFromOutput(boundedDelta, agent.systemPrompt), + const remainingBytes = maxOutputBytes - outputBytes + if (remainingBytes <= 0) throw new Error('sandbox exceeded max output tokens') + const bounded = truncateUtf8(event.data.delta, remainingBytes, encoder) + const boundedDelta = bounded.text + outputBytes += bounded.bytes + bufferedTextParts.push(boundedDelta) + yield { kind: 'activity' } + if (bounded.truncated) { + throw new Error('sandbox exceeded max output tokens') } - if (boundedDelta.length < event.data.delta.length) return continue } if (event.type === 'input-required' || event.data?.inputRequired) { + const usage = finalizeUsage( + withObservedUsage(usageParts, observedReasoningTokens, observedToolTokens, observedToolCalls), + executionBudget, + ) + for (const bufferedText of bufferedTextParts) { + yield { kind: 'text', delta: redactSystemPromptFromOutput(bufferedText, agent.systemPrompt) } + } yield { kind: 'input-required', prompt: event.data?.inputRequired?.prompt } // Terminal for the sandbox stream — sandbox SHOULD stop emitting until // the gateway dispatches a continuation message with the new user input. + yield { kind: 'usage', usage } return } } + const usage = finalizeUsage( + withObservedUsage(usageParts, observedReasoningTokens, observedToolTokens, observedToolCalls), + executionBudget, + ) + for (const bufferedText of bufferedTextParts) { + yield { kind: 'text', delta: redactSystemPromptFromOutput(bufferedText, agent.systemPrompt) } + } + yield { kind: 'usage', usage } } /** @@ -589,12 +799,14 @@ export async function* dispatchSandboxStreamRich( export async function settleAndRecord( agent: AgentMeta, authz: AuthorizedRequest, - inputTokens: number, - outputTokens: number, + usage: SandboxUsageReceipt, config: GatewayConfig, obs: GatewayObserver | undefined, ): Promise { - const totalCost = (inputTokens + outputTokens) * agent.pricePerTokenUsd + const tokenCost = ( + usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens + ) * agent.pricePerTokenUsd + const totalCost = Math.max(tokenCost, usage.providerCostUsd) const ownerEarned = totalCost * (1 - agent.platformFeePercent) const platformFee = totalCost * agent.platformFeePercent const usageEvent = { @@ -603,23 +815,39 @@ export async function settleAndRecord( agentSlug: agent.slug, consumerId: authz.consumerId, paymentMethod: authz.paymentMethod, - inputTokens, - outputTokens, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + reasoningTokens: usage.reasoningTokens, + toolTokens: usage.toolTokens, + toolCallCount: usage.toolCallCount, + providerCostUsd: usage.providerCostUsd, totalCostUsd: totalCost, ownerEarnedUsd: ownerEarned, platformFeeUsd: platformFee, durationMs: Date.now() - authz.startMs, } - await config.recordUsage(usageEvent) const ctx: RequestContext = { requestId: authz.requestId, agentSlug: agent.slug, startMs: authz.startMs, } - await obs?.onRequestComplete?.(ctx, usageEvent) - if (config.settlePayment) { - await config - .settlePayment( + try { + if (authz.paymentOperation && config.x402.paymentOperations) { + const amount = actualX402Amount( + agent.pricePerTokenUsd, + usage.inputTokens, + usage.outputTokens, + usage.reasoningTokens, + usage.toolTokens, + config.x402.currencyDecimals, + usage.providerCostUsd, + ) + authz.paymentOperation = await config.x402.paymentOperations.settlePayment( + authz.paymentOperation, + { amount, totalCostUsd: totalCost, usage }, + ) + } else if (config.settlePayment) { + await config.settlePayment( { method: authz.paymentMethod, consumerId: authz.consumerId, @@ -627,16 +855,161 @@ export async function settleAndRecord( }, totalCost, ) - .catch(async (err) => { - const msg = err instanceof Error ? err.message : String(err) - console.error(`[agent-gateway] settlement failed for ${authz.consumerId}: ${msg}`) - await obs?.onSettlementError?.(ctx, { - consumerId: authz.consumerId, - method: authz.paymentMethod, - errorMessage: msg, - }) - }) + } + await config.recordUsage(usageEvent) + await obs?.onRequestComplete?.(ctx, usageEvent) + } catch (err) { + const msg = err instanceof Error ? err.message : String(err) + console.error(`[agent-gateway] settlement failed for ${authz.consumerId}: ${msg}`) + await obs?.onSettlementError?.(ctx, { + consumerId: authz.consumerId, + method: authz.paymentMethod, + errorMessage: msg, + }) + throw err + } +} + +function mergeUsage( + current: Partial, + update: Partial, +): Partial { + const merged = { ...current, ...update } + for (const key of ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount', 'providerCostUsd'] as const) { + const value = update[key] + if (value !== undefined && (!Number.isFinite(value) || value < 0)) { + throw new Error(`sandbox usage field ${key} is invalid`) + } + if (value !== undefined && current[key] !== undefined) { + // Usage events are cumulative receipts. Never let a later partial or + // final event erase spend observed earlier in the same execution. + merged[key] = Math.max(current[key]!, value) + } + } + if (current.budgetEnforced === false || update.budgetEnforced === false) { + merged.budgetEnforced = false + } + return merged +} + +function withObservedUsage( + usage: Partial, + reasoningTokens: number, + toolTokens: number, + toolCallCount: number, +): Partial { + return { + ...usage, + ...(usage.reasoningTokens !== undefined || reasoningTokens > 0 + ? { reasoningTokens: Math.max(usage.reasoningTokens ?? 0, reasoningTokens) } + : {}), + ...(usage.toolTokens !== undefined || toolTokens > 0 + ? { toolTokens: Math.max(usage.toolTokens ?? 0, toolTokens) } + : {}), + ...(usage.toolCallCount !== undefined || toolCallCount > 0 + ? { toolCallCount: Math.max(usage.toolCallCount ?? 0, toolCallCount) } + : {}), + } +} + +function nonNegativeSafeInteger(value: number, name: string): number { + if (!Number.isSafeInteger(value) || value < 0) { + throw new Error(`sandbox ${name} is invalid`) + } + return value +} + +function enforceUsageBudget( + usage: Partial, + budget: SandboxExecutionBudget, +): void { + if (usage.inputTokens !== undefined && usage.inputTokens > budget.maxInputTokens) { + throw new Error('sandbox exceeded max input tokens') + } + if (usage.outputTokens !== undefined && usage.outputTokens > budget.maxOutputTokens) { + throw new Error('sandbox exceeded max output tokens') + } + if (usage.reasoningTokens !== undefined && usage.reasoningTokens > budget.maxReasoningTokens) { + throw new Error('sandbox exceeded max reasoning tokens') + } + if (usage.toolTokens !== undefined && usage.toolTokens > budget.maxToolTokens) { + throw new Error('sandbox exceeded max tool tokens') + } + if (usage.toolCallCount !== undefined && usage.toolCallCount > budget.maxToolCalls) { + throw new Error('sandbox exceeded max tool calls') + } + if (usage.providerCostUsd !== undefined && usage.providerCostUsd > budget.maxProviderCostUsd) { + throw new Error('sandbox exceeded max provider cost') + } +} + +function finalizeUsage( + parts: Partial, + budget: SandboxExecutionBudget, +): SandboxUsageReceipt { + const fields = ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount', 'providerCostUsd', 'budgetEnforced'] as const + if (fields.some((field) => parts[field] === undefined)) { + throw new Error('sandbox did not provide a complete usage receipt') + } + const usage = parts as SandboxUsageReceipt + for (const field of ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount'] as const) { + if (!Number.isSafeInteger(usage[field]) || usage[field] < 0) { + throw new Error(`sandbox usage field ${field} is invalid`) + } + } + if (!Number.isFinite(usage.providerCostUsd) || usage.providerCostUsd < 0) { + throw new Error('sandbox usage provider cost is invalid') + } + if (typeof usage.budgetEnforced !== 'boolean') { + throw new Error('sandbox usage budget flag is invalid') + } + if (!Number.isSafeInteger( + usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens, + )) { + throw new Error('sandbox usage token total exceeds safe integer range') + } + enforceUsageBudget(usage, budget) + if (!usage.budgetEnforced) throw new Error('sandbox did not enforce the execution budget') + return usage +} + +function actualX402Amount( + pricePerTokenUsd: number, + inputTokens: number, + outputTokens: number, + reasoningTokens: number, + toolTokens: number, + currencyDecimals = 6, + providerCostUsd = 0, +): bigint { + const { numerator, denominator } = decimalFraction(pricePerTokenUsd) + const scaled = BigInt(inputTokens + outputTokens + reasoningTokens + toolTokens) * + numerator * 10n ** BigInt(currencyDecimals) + const tokenAmount = (scaled + denominator - 1n) / denominator + const provider = providerCostUsd === 0 + ? { numerator: 0n, denominator: 1n } + : decimalFraction(providerCostUsd) + const providerScaled = provider.numerator * 10n ** BigInt(currencyDecimals) + const providerAmount = (providerScaled + provider.denominator - 1n) / provider.denominator + return tokenAmount > providerAmount ? tokenAmount : providerAmount +} + +function truncateUtf8( + value: string, + maxBytes: number, + encoder: TextEncoder, +): { text: string; bytes: number; truncated: boolean } { + const bytes = encoder.encode(value).byteLength + if (bytes <= maxBytes) return { text: value, bytes, truncated: false } + let text = '' + let used = 0 + for (const character of value) { + const characterBytes = encoder.encode(character).byteLength + if (used + characterBytes > maxBytes) break + text += character + used += characterBytes } + return { text, bytes: used, truncated: true } } /** Token estimate matching the existing chat-completions handler (4 chars ≈ 1 token). */ diff --git a/src/index.ts b/src/index.ts index 6f8ee80..2684d1c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,11 +1,23 @@ export { createAgentGateway } from './middleware' +export { reclaimPayment } from './dispatch' export { verifyX402, verifyMpp, defaultVerifyApiKey, isApiKeyAuthEnabled, isMppAuthEnabled, + mppReplayNonceKey, } from './verify' +export { + PAYMENT_PROTOCOL_VERSION, + MemoryPaymentOperations, + type MemoryPaymentOperationsOptions, + type PaymentAuthorizationContext, + type PaymentOperation, + type PaymentOperationState, + type PaymentOperations, + type PaymentSettlementInput, +} from './payment-operations' export { filterConsumerMessages, filterConsumerMessagesStrict, @@ -56,6 +68,8 @@ export type { PaymentResult, ApiKeyInfo, GatewayUsageEvent, + SandboxExecutionBudget, + SandboxUsageReceipt, SandboxStreamEvent, SandboxBox, GatewayConfig, diff --git a/src/middleware.ts b/src/middleware.ts index 8155579..8497115 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -6,9 +6,10 @@ import { type AuthorizedRequest, type GatewayState, authenticateAndGuard, - dispatchSandboxStream, - estimateBillableInputTokens, - estimateTokens, + claimPayment, + dispatchSandboxStreamRich, + releasePayment, + releasePaymentAfterFailure, settleAndRecord, } from './dispatch' import { MemoryNonceStore } from './nonce-store' @@ -38,11 +39,11 @@ export function createAgentGateway(config: GatewayConfig) { } const maxOutputTokens = config.maxOutputTokens ?? 4096 const defaultOutputTokens = config.defaultOutputTokens ?? 1024 - if (!Number.isInteger(maxOutputTokens) || maxOutputTokens <= 0) { + if (!Number.isSafeInteger(maxOutputTokens) || maxOutputTokens <= 0) { throw new Error('createAgentGateway: maxOutputTokens must be a positive integer') } if ( - !Number.isInteger(defaultOutputTokens) || + !Number.isSafeInteger(defaultOutputTokens) || defaultOutputTokens <= 0 || defaultOutputTokens > maxOutputTokens ) { @@ -58,6 +59,32 @@ export function createAgentGateway(config: GatewayConfig) { ) { throw new Error('createAgentGateway: x402.currencyDecimals must be an integer between 0 and 18') } + const executionBudget = config.executionBudget + for (const [name, value] of [ + ['maxReasoningTokens', executionBudget?.maxReasoningTokens ?? maxOutputTokens], + ['maxToolTokens', executionBudget?.maxToolTokens ?? maxOutputTokens], + ['maxToolCalls', executionBudget?.maxToolCalls ?? 8], + ] as const) { + if (!Number.isSafeInteger(value) || value < 0) { + throw new Error(`createAgentGateway: executionBudget.${name} must be a non-negative safe integer`) + } + } + if ( + executionBudget?.maxProviderCostUsd !== undefined && + (!Number.isFinite(executionBudget.maxProviderCostUsd) || executionBudget.maxProviderCostUsd < 0) + ) { + throw new Error('createAgentGateway: executionBudget.maxProviderCostUsd must be finite and non-negative') + } + if (config.x402.paymentOperations && config.x402.paymentProtocolVersion === undefined) { + throw new Error('createAgentGateway: paymentProtocolVersion must be explicit when durable payment operations are configured') + } + if (config.x402.paymentProtocolVersion === 2 && + (!config.x402.paymentOperations || config.x402.paymentOperations.protocolVersion !== 2)) { + throw new Error('createAgentGateway: payment protocol version 2 requires durable payment operations') + } + if (config.x402.paymentProtocolVersion === 1 && config.x402.paymentOperations) { + throw new Error('createAgentGateway: version 1 cannot be combined with version 2 payment operations') + } const gw = new Hono() const rateLimitStore: RateLimitStore = config.rateLimitStore ?? new MemoryRateLimitStore() const state: GatewayState = { @@ -68,6 +95,10 @@ export function createAgentGateway(config: GatewayConfig) { maxLen: config.maxMessageLength ?? 8000, maxOutputTokens, defaultOutputTokens, + maxReasoningTokens: config.executionBudget?.maxReasoningTokens ?? maxOutputTokens, + maxToolTokens: config.executionBudget?.maxToolTokens ?? maxOutputTokens, + maxToolCalls: config.executionBudget?.maxToolCalls ?? 8, + maxProviderCostUsd: config.executionBudget?.maxProviderCostUsd, obs: config.observer, } const obs: GatewayObserver | undefined = state.obs @@ -160,6 +191,36 @@ export function createAgentGateway(config: GatewayConfig) { ) if (guard instanceof Response) return guard const authz = guard + try { + await claimPayment(authz, config, state) + } catch { + try { + await releasePayment(authz, config, 'payment authorization failed') + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } + await obs?.onAuthFailure?.( + { + requestId: authz.requestId, + agentSlug: authz.agent.slug, + startMs: authz.startMs, + }, + { method: authz.paymentMethod, code: 'payment_authorization_failed', httpStatus: 402 }, + ) + return c.json( + { + error: { + message: 'Payment authorization failed', + type: 'payment_required', + code: 'payment_authorization_failed', + }, + }, + { status: 402, headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': authz.requestId } }, + ) + } return streamChatCompletions(c, authz, config, obs) }) @@ -178,6 +239,10 @@ export function createAgentGateway(config: GatewayConfig) { return gw } +// Consumers that bind the gateway through a package boundary can fail closed +// when an old binary ignores the version 2 operation contract. +Object.assign(createAgentGateway, { paymentProtocolVersion: 2 as const }) + /** * Drain the sandbox stream into an OpenAI-shaped SSE response, settle the * payment, fire observer hooks. Identical pre-refactor behavior, just lifted @@ -199,8 +264,9 @@ function streamChatCompletions( rateLimitRemaining, maxOutputTokens, } = authz - const inputTokens = estimateBillableInputTokens(agent, userMessage) let outputText = '' + let usage: import('./types').SandboxUsageReceipt | undefined + let workObserved = false const ctx: RequestContext = { requestId, agentSlug: agent.slug, @@ -223,7 +289,7 @@ function streamChatCompletions( } try { - for await (const delta of dispatchSandboxStream( + for await (const event of dispatchSandboxStreamRich( agent, userMessage, consumerId, @@ -232,9 +298,16 @@ function streamChatCompletions( undefined, maxOutputTokens, )) { - sendChunk(delta) + if (event.kind === 'text') { + sendChunk(event.delta) + workObserved = true + } + if (event.kind === 'activity') workObserved = true + if (event.kind === 'usage') usage = event.usage } + if (!usage) throw new Error('sandbox did not provide a usage receipt') + const done: ChatCompletionChunk = { id: `chatcmpl-${Date.now()}`, object: 'chat.completion.chunk', @@ -248,8 +321,7 @@ function streamChatCompletions( await settleAndRecord( agent, authz, - inputTokens, - estimateTokens(outputText), + usage, config, obs, ) @@ -260,7 +332,22 @@ function streamChatCompletions( rawMessage.includes('/') || rawMessage.includes('\\') ? 'Internal agent error' : rawMessage - await obs?.onStreamError?.(ctx, { consumerId, errorMessage: rawMessage }) + try { + await obs?.onStreamError?.(ctx, { consumerId, errorMessage: rawMessage }) + } catch (observerError) { + console.error( + `[agent-gateway] stream observer failed for ${requestId}:`, + observerError instanceof Error ? observerError.message : String(observerError), + ) + } + try { + await releasePaymentAfterFailure(authz, config, rawMessage, workObserved || usage !== undefined) + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } controller.enqueue( encoder.encode( `data: ${JSON.stringify({ error: { message: safeMessage, type: 'server_error' } })}\n\n`, diff --git a/src/nonce-store.ts b/src/nonce-store.ts index 3abb85e..326cae2 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -6,6 +6,11 @@ export interface NonceStore { /** Check if nonce has been seen. Returns true if already used (reject). */ hasSeen(nonce: string): Promise + /** + * Atomically claim a nonce. An owner id makes a retry by the same payment + * operation idempotent while a legacy claim still fails closed. + */ + claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise /** Mark nonce as used. TTL = how long to remember it (seconds). */ markSeen(nonce: string, ttlSeconds: number): Promise } @@ -16,22 +21,33 @@ export interface NonceStore { /** In-memory nonce store with automatic eviction. Use in tests or single-worker deploys. */ export class MemoryNonceStore implements NonceStore { - private seen = new Map() // nonce → expiresAt + private seen = new Map() private lastEviction = Date.now() async hasSeen(nonce: string): Promise { this.evictExpired() - const expiresAt = this.seen.get(nonce) - if (!expiresAt) return false - if (expiresAt < Date.now()) { + const entry = this.seen.get(nonce) + if (!entry) return false + if (entry.expiresAt < Date.now()) { this.seen.delete(nonce) return false } return true } + async claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise { + this.evictExpired() + const now = Date.now() + const entry = this.seen.get(nonce) + if (entry !== undefined && entry.expiresAt >= now) { + return ownerId !== undefined && entry.ownerId === ownerId + } + this.seen.set(nonce, { expiresAt: now + ttlSeconds * 1000, ownerId }) + return true + } + async markSeen(nonce: string, ttlSeconds: number): Promise { - this.seen.set(nonce, Date.now() + ttlSeconds * 1000) + this.seen.set(nonce, { expiresAt: Date.now() + ttlSeconds * 1000 }) this.evictExpired() } @@ -40,8 +56,8 @@ export class MemoryNonceStore implements NonceStore { // Evict at most every 60 seconds to avoid O(n) on every request if (now - this.lastEviction < 60_000) return this.lastEviction = now - for (const [nonce, expiresAt] of this.seen) { - if (expiresAt < now) this.seen.delete(nonce) + for (const [nonce, entry] of this.seen) { + if (entry.expiresAt < now) this.seen.delete(nonce) } } } @@ -58,6 +74,8 @@ export class MemoryNonceStore implements NonceStore { export interface KVNamespace { get(key: string, options?: { type?: 'text' | 'json' }): Promise put(key: string, value: string, options?: { expirationTtl?: number }): Promise + /** Required for version 2 gateway claims. A plain KV put is not atomic. */ + putIfAbsent?(key: string, value: string, options?: { expirationTtl?: number }): Promise delete(key: string): Promise } @@ -67,12 +85,8 @@ export interface KVNamespace { * Why this exists: MemoryNonceStore works on a single worker instance, but * Cloudflare routes requests across multiple isolates. Without shared state, * an attacker could retry a replayed nonce against a different isolate and - * have it accepted. This implementation uses Workers KV with native TTL so - * the nonce automatically expires at payment-expiry time. - * - * TTL precision: KV is eventually consistent (propagation ~60s). For x402 - * with 10-minute expiry windows this is fine — by the time KV propagates, - * the payment itself would be expired anyway. + * have it accepted. Version 2 requires an atomic binding for payment claims. + * This store keeps the nonce only for legacy replay protection. * * Usage: * const nonceStore = new KvNonceStore(env.NONCE_KV, 'x402') @@ -90,6 +104,22 @@ export class KvNonceStore implements NonceStore { return value !== null } + async claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise { + if (!this.kv.putIfAbsent) { + throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') + } + const ttl = Math.max(ttlSeconds, 60) + const key = this.key(nonce) + const value = ownerId ?? '1' + const existing = await this.kv.get(key) + if (existing !== null) return ownerId !== undefined && existing === ownerId + const inserted = await this.kv.putIfAbsent(key, value, { expirationTtl: ttl }) + if (inserted || ownerId === undefined) return inserted + // Another isolate may have won between get and putIfAbsent. Re-read so a + // retry by the same durable operation remains idempotent. + return (await this.kv.get(key)) === ownerId + } + async markSeen(nonce: string, ttlSeconds: number): Promise { // KV minimum TTL is 60 seconds const ttl = Math.max(ttlSeconds, 60) diff --git a/src/payment-operations.ts b/src/payment-operations.ts new file mode 100644 index 0000000..8c412cf --- /dev/null +++ b/src/payment-operations.ts @@ -0,0 +1,316 @@ +import type { SandboxUsageReceipt } from './types' + +/** Version negotiated by gateways that use durable payment operations. */ +export const PAYMENT_PROTOCOL_VERSION = 2 as const + +export type PaymentOperationState = + | 'claiming' + | 'claimed' + | 'settling' + | 'settled' + | 'releasing' + | 'released' + | 'reclaimable' + | 'reclaimed' + +/** Durable ownership of one signed payment authorization. */ +export interface PaymentOperation { + protocolVersion: typeof PAYMENT_PROTOCOL_VERSION + operationId: string + nonceKey: string + authorizationId: string + reservedAmount: bigint + settledAmount: bigint + refundAmount: bigint + expiresAt: number + state: PaymentOperationState +} + +export interface PaymentAuthorizationContext { + requestId: string + agentId: string + requiredAmount: bigint + maxOutputTokens: number + executionBudget: { + maxInputTokens: number + maxOutputTokens: number + maxReasoningTokens: number + maxToolTokens: number + maxToolCalls: number + maxProviderCostUsd: number + } +} + +export interface PaymentSettlementInput { + amount: bigint + totalCostUsd: number + usage: SandboxUsageReceipt +} + +/** + * One payment lifecycle shared by every payment-backed gateway surface. + * Implementations must persist the operation before external side effects. + */ +export interface PaymentOperations { + readonly protocolVersion: typeof PAYMENT_PROTOCOL_VERSION + claimPayment( + payload: Record, + context: PaymentAuthorizationContext, + ): Promise + settlePayment( + operation: PaymentOperation, + input: PaymentSettlementInput, + ): Promise + releasePayment(operation: PaymentOperation, reason: string): Promise + reclaimPayment(operationId: string): Promise +} + +export interface MemoryPaymentOperationsOptions { + now?: () => number + onClaim?: (operation: PaymentOperation) => Promise + onSettle?: (operation: PaymentOperation, input: PaymentSettlementInput) => Promise + onRelease?: (operation: PaymentOperation, reason: string) => Promise + onReclaim?: (operation: PaymentOperation) => Promise +} + +/** Small atomic implementation used by single-process deployments and tests. */ +export class MemoryPaymentOperations implements PaymentOperations { + readonly protocolVersion = PAYMENT_PROTOCOL_VERSION + private readonly operations = new Map() + private readonly settleFlights = new Map>() + private readonly releaseFlights = new Map>() + private readonly now: () => number + + constructor(private readonly options: MemoryPaymentOperationsOptions = {}) { + this.now = options.now ?? (() => Math.floor(Date.now() / 1000)) + } + + async claimPayment( + payload: Record, + context: PaymentAuthorizationContext, + ): Promise { + const nonceKey = paymentNonceKey(payload) + const operationId = `x402:${nonceKey}` + const existing = this.operations.get(operationId) + if (existing) { + if (existing.state === 'claiming') throw new Error('payment operation is being recovered') + throw new Error('payment operation was already claimed') + } + + const reservedAmount = unsignedAmount(payload.amount) + if (reservedAmount < context.requiredAmount) { + throw new Error('payment authorization is below the request ceiling') + } + const expiresAt = unsignedAmount(payload.expiry, 'expiry') + if (expiresAt > BigInt(Number.MAX_SAFE_INTEGER)) throw new Error('payment expiry is too large') + if (Number(expiresAt) <= this.now()) throw new Error('payment authorization has expired') + + const operation: PaymentOperation = { + protocolVersion: PAYMENT_PROTOCOL_VERSION, + operationId, + nonceKey, + authorizationId: typeof payload.authHash === 'string' ? payload.authHash : operationId, + reservedAmount, + settledAmount: 0n, + refundAmount: reservedAmount, + expiresAt: Number(expiresAt), + state: 'claiming', + } + // The map write is synchronous. No second caller can observe a free nonce + // between the uniqueness check and ownership write. + this.operations.set(operationId, operation) + try { + await this.options.onClaim?.(operation) + const claimed = { ...operation, state: 'claimed' as const } + this.operations.set(operationId, claimed) + return claimed + } catch (error) { + // Keep the durable `claiming` row. The external authorization may have + // committed before its acknowledgement was lost, so expiry recovery + // must own the decision to reclaim it. + throw error + } + } + + async settlePayment( + operation: PaymentOperation, + input: PaymentSettlementInput, + ): Promise { + const current = this.requireCurrent(operation) + if (current.state === 'settled') { + if (current.settledAmount !== input.amount) throw new Error('payment operation was settled twice') + return current + } + if (current.state === 'settling') { + if (current.settledAmount !== input.amount) throw new Error('payment operation has a different pending settlement') + const flight = this.settleFlights.get(current.operationId) + if (flight) return flight + return this.recoverSettlement(current, input) + } else if (current.state !== 'claimed') { + throw new Error(`cannot settle payment in state ${current.state}`) + } + validateSettlement(current, input) + const settling = { + ...current, + state: 'settling' as const, + settledAmount: input.amount, + refundAmount: current.reservedAmount - input.amount, + } + this.operations.set(current.operationId, settling) + const flight = this.runSettlement(settling, input) + this.settleFlights.set(current.operationId, flight) + try { + return await flight + } finally { + if (this.settleFlights.get(current.operationId) === flight) { + this.settleFlights.delete(current.operationId) + } + } + } + + async releasePayment(operation: PaymentOperation, reason: string): Promise { + const current = this.requireCurrent(operation) + if (current.state === 'released') return current + if (current.state === 'releasing') { + const flight = this.releaseFlights.get(current.operationId) + if (flight) return flight + return this.runRelease(current, reason) + } + if (current.state !== 'claimed') throw new Error(`cannot release payment in state ${current.state}`) + const releasing = { ...current, state: 'releasing' as const } + this.operations.set(current.operationId, releasing) + return this.runRelease(releasing, reason) + } + + async reclaimPayment(operationId: string): Promise { + const current = this.operations.get(operationId) + if (!current) throw new Error('payment operation was not found') + if (current.state === 'reclaimed') return current + if (current.state === 'settling') { + const flight = this.settleFlights.get(operationId) + if (flight) return flight + return this.recoverSettlement(current, { + amount: current.settledAmount, + totalCostUsd: 0, + usage: { + inputTokens: 0, + outputTokens: 0, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0, + budgetEnforced: true, + }, + }) + } + if (current.state === 'releasing') { + if (!this.options.onReclaim) throw new Error('payment release recovery is not configured') + await this.options.onReclaim(current) + const released = { ...current, state: 'released' as const } + this.operations.set(operationId, released) + return released + } + if (current.state !== 'claiming' && current.state !== 'claimed' && current.state !== 'reclaimable') { + throw new Error(`cannot reclaim payment in state ${current.state}`) + } + if (current.expiresAt > this.now()) throw new Error('payment operation has not expired') + const reclaimable = { ...current, state: 'reclaimable' as const } + this.operations.set(operationId, reclaimable) + await this.options.onReclaim?.(reclaimable) + const reclaimed = { ...reclaimable, state: 'reclaimed' as const } + this.operations.set(operationId, reclaimed) + return reclaimed + } + + private async runSettlement( + settling: PaymentOperation, + input: PaymentSettlementInput, + ): Promise { + await this.options.onSettle?.(settling, input) + const settled = { + ...settling, + state: 'settled' as const, + settledAmount: input.amount, + refundAmount: settling.reservedAmount - input.amount, + } + this.operations.set(settling.operationId, settled) + return settled + } + + private async recoverSettlement( + settling: PaymentOperation, + input: PaymentSettlementInput, + ): Promise { + if (!this.options.onReclaim) throw new Error('payment settlement recovery is not configured') + await this.options.onReclaim(settling) + const settled = { + ...settling, + state: 'settled' as const, + settledAmount: input.amount, + refundAmount: settling.reservedAmount - input.amount, + } + this.operations.set(settling.operationId, settled) + return settled + } + + private async runRelease( + releasing: PaymentOperation, + reason: string, + ): Promise { + const flight = (async () => { + await this.options.onRelease?.(releasing, reason) + const released = { ...releasing, state: 'released' as const } + this.operations.set(releasing.operationId, released) + return released + })() + this.releaseFlights.set(releasing.operationId, flight) + try { + return await flight + } finally { + if (this.releaseFlights.get(releasing.operationId) === flight) { + this.releaseFlights.delete(releasing.operationId) + } + } + } + + get(operationId: string): PaymentOperation | undefined { + const operation = this.operations.get(operationId) + return operation ? { ...operation } : undefined + } + + private requireCurrent(operation: PaymentOperation): PaymentOperation { + const current = this.operations.get(operation.operationId) + if (!current) throw new Error('payment operation was not found') + if (current.protocolVersion !== operation.protocolVersion) { + throw new Error('payment operation protocol version mismatch') + } + return current + } +} + +export function paymentNonceKey(payload: Record): string { + const commitment = String(payload.commitment ?? '').toLowerCase() + const nonce = unsignedAmount(payload.nonce, 'nonce').toString() + if (!commitment) throw new Error('payment commitment is required') + return `${commitment}:${nonce}` +} + +function unsignedAmount(value: unknown, name = 'amount'): bigint { + const raw = typeof value === 'string' + ? value + : typeof value === 'number' && Number.isSafeInteger(value) + ? String(value) + : '' + if (!/^\d+$/.test(raw)) throw new Error(`${name} is not an unsigned integer`) + return BigInt(raw) +} + +function validateSettlement(operation: PaymentOperation, input: PaymentSettlementInput): void { + if (input.amount < 0n || input.amount > operation.reservedAmount) { + throw new Error('settled amount must be between zero and the reserved amount') + } + if (!Number.isFinite(input.totalCostUsd) || input.totalCostUsd < 0) { + throw new Error('settlement cost must be finite and non-negative') + } + if (!input.usage.budgetEnforced) throw new Error('sandbox usage receipt is not budget-enforced') +} diff --git a/src/types.ts b/src/types.ts index 6b74680..abdc2ec 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,3 +1,9 @@ +import type { + PaymentAuthorizationContext, + PaymentOperation, + PaymentOperations, +} from './payment-operations' + // --- Agent resolution --- export interface AgentMeta { @@ -73,6 +79,15 @@ export interface AgentMeta { export type PaymentMethod = 'x402' | 'mpp' | 'apikey' | 'none' +export interface SandboxExecutionBudget { + maxInputTokens: number + maxOutputTokens: number + maxReasoningTokens: number + maxToolTokens: number + maxToolCalls: number + maxProviderCostUsd: number +} + export interface X402Config { /** Ethereum operator address for SpendAuth verification */ operatorAddress: string @@ -84,24 +99,27 @@ export interface X402Config { rpcUrl?: string /** Demo mode: skip signature verification (default: false). NEVER enable in production. */ demoMode?: boolean + /** Protocol version for new durable payment operations. Version 1 remains supported for mixed deploys. */ + paymentProtocolVersion?: 1 | 2 /** * Production signature verification. This callback must not reserve, claim, * or mutate payment state. */ - verifySigner?: (payload: Record) => Promise + verifySigner?: ( + payload: Record, + context?: { protocolVersion: 1 | 2; requestId?: string }, + ) => Promise /** - * Reserve or claim the verified payment after all request checks pass and - * immediately before sandbox work starts. Return false to reject the call. + * Claim the verified payment after all request checks pass and immediately + * before sandbox work starts. Version 2 returns durable operation ownership. + * A boolean return is the version 1 mixed-deploy compatibility path. */ authorizePayment?: ( payload: Record, - context: { - requestId: string - agentId: string - requiredAmount: bigint - maxOutputTokens: number - }, - ) => Promise + context: PaymentAuthorizationContext, + ) => Promise + /** Version 2 operation store. It owns claim, settle, release, and reclaim. */ + paymentOperations?: PaymentOperations /** * Number of base-unit decimals used by the payment token. Defaults to 6. * The gateway uses this value to reject a payment that cannot cover the @@ -173,6 +191,10 @@ export interface GatewayUsageEvent { paymentMethod: PaymentMethod inputTokens: number outputTokens: number + reasoningTokens: number + toolTokens: number + toolCallCount: number + providerCostUsd: number totalCostUsd: number ownerEarnedUsd: number platformFeeUsd: number @@ -196,13 +218,34 @@ export interface SandboxStreamEvent { * the caller (rendered as the input-required message body). */ inputRequired?: { prompt?: string } + /** Provider receipt fields. The final event must include every field. */ + usage?: Partial + /** Tool or reasoning events may carry hidden usage without visible text. */ + tool?: { name?: string; inputTokens?: number; outputTokens?: number } + reasoning?: { tokens?: number } } } +export interface SandboxUsageReceipt { + inputTokens: number + outputTokens: number + reasoningTokens: number + toolTokens: number + toolCallCount: number + providerCostUsd: number + /** True only when the provider/adapter enforced every supplied budget. */ + budgetEnforced: boolean +} + export interface SandboxBox { streamPrompt( message: string, - opts?: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number }, + opts?: { + sessionId?: string + systemPrompt?: string + maxOutputTokens?: number + executionBudget?: SandboxExecutionBudget + }, ): AsyncIterable } @@ -261,6 +304,14 @@ export interface GatewayConfig { /** Output token limit used when a request omits `max_tokens`. Defaults to 1024. */ defaultOutputTokens?: number + /** Hidden provider spend limits included in the pre-execution payment quote. */ + executionBudget?: { + maxReasoningTokens?: number + maxToolTokens?: number + maxToolCalls?: number + maxProviderCostUsd?: number + } + /** Required scope for chat endpoint (default: "chat"). API keys must include this scope. */ requiredScope?: string diff --git a/src/verify.ts b/src/verify.ts index e8450b6..17e7b6d 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -1,6 +1,36 @@ import type { X402Config, MppConfig, ApiKeyInfo, GatewayConfig } from './types' import type { NonceStore } from './nonce-store' +/** Return the canonical opaque nonce key used by the final payment claim. */ +export function mppReplayNonceKey(authHeader: string): string | undefined { + const match = authHeader.match(/^Payment\s+(\S+)\s+(\S+)$/i) + if (!match) return undefined + const [, method, credentialB64] = match + try { + const decoded = Buffer.from(credentialB64, 'base64url').toString('utf-8') + const credential = JSON.parse(decoded) as Record + const nested = credential.payload + const payload = nested && typeof nested === 'object' && !Array.isArray(nested) + ? nested as Record + : credential + return canonicalMppNonceKey(method, payload) + } catch { + return undefined + } +} + +function canonicalMppNonceKey(method: string, payload: Record): string | undefined { + if (payload.nonce === undefined) return undefined + const nonce = BigInt(String(payload.nonce)).toString() + const commitment = payload.commitment + // BlueprinTEVM carries the same SpendAuth identity as x402. Keep one + // namespace so a credential cannot cross the two HTTP transports. + if (method.toLowerCase() === 'blueprintevm' && typeof commitment === 'string' && commitment.length > 0) { + return `${commitment.toLowerCase()}:${nonce}` + } + return `mpp:${method.toLowerCase()}:${String(payload.commitment ?? payload.from ?? 'unknown').toLowerCase()}:${nonce}` +} + /** Pure capability checks shared by discovery and every request protocol. */ export function isApiKeyAuthEnabled( config: Pick, @@ -56,11 +86,13 @@ export async function verifyX402( // settle funds as part of its production verification path. if (amount < minimumAmount || minimumAmount <= 0n) return null - const nonceKey = `${raw.commitment}:${nonce.toString()}` + const nonceKey = `${String(raw.commitment).toLowerCase()}:${nonce.toString()}` if (nonceStore && await nonceStore.hasSeen(nonceKey)) return null if (config.verifySigner) { - const verified = await config.verifySigner(raw) + const verified = await config.verifySigner(raw, { + protocolVersion: config.paymentProtocolVersion ?? (config.paymentOperations ? 2 : 1), + }) if (!verified) return null } else if (!config.demoMode) { return null @@ -71,7 +103,8 @@ export async function verifyX402( if (nonceStore && markNonce) { // Mark seen with TTL matching the expiry window (max 1 hour) const ttl = Math.min(Number(expiry) - Math.floor(Date.now() / 1000), 3600) - await nonceStore.markSeen(nonceKey, Math.max(ttl, 60)) + const claimed = await nonceStore.claim(nonceKey, Math.max(ttl, 60)) + if (!claimed) return null } return raw.commitment @@ -97,6 +130,7 @@ export async function verifyMpp( x402Config: X402Config, nonceStore?: NonceStore, minimumAmount = 1n, + markNonce = true, ): Promise { // MPP format: "Payment " const match = authHeader.match(/^Payment\s+(\S+)\s+(\S+)$/i) @@ -143,17 +177,16 @@ export async function verifyMpp( return null } - const nonceKey = - nonceStore && payload.nonce !== undefined - ? `mpp:${method}:${String(payload.commitment ?? payload.from ?? 'unknown')}:${String(payload.nonce)}` - : null + const nonceKey = nonceStore ? canonicalMppNonceKey(method, payload) ?? null : null if (nonceKey && await nonceStore!.hasSeen(nonceKey)) return null let consumerId: string | null = null if (config.verifySigner) { consumerId = await config.verifySigner(payload, { method, credential: decoded }) } else if (method === 'blueprintevm' && x402Config.verifySigner && payload.commitment) { - const verified = await x402Config.verifySigner(payload) + const verified = await x402Config.verifySigner(payload, { + protocolVersion: x402Config.paymentProtocolVersion ?? (x402Config.paymentOperations ? 2 : 1), + }) consumerId = verified ? String(payload.commitment) : null } else if (x402Config.demoMode) { const identity = payload.commitment ?? payload.from @@ -164,12 +197,13 @@ export async function verifyMpp( } if (!consumerId) return null - if (nonceStore && payload.nonce !== undefined) { + if (nonceStore && payload.nonce !== undefined && markNonce) { const expiry = payload.expiry === undefined ? Math.floor(Date.now() / 1000) + 3600 : Number(payload.expiry) const ttl = Math.min(expiry - Math.floor(Date.now() / 1000), 3600) - await nonceStore.markSeen(nonceKey!, Math.max(ttl, 60)) + const claimed = await nonceStore.claim(nonceKey!, Math.max(ttl, 60)) + if (!claimed) return null } return consumerId diff --git a/tests/a2a-long-horizon.test.ts b/tests/a2a-long-horizon.test.ts index 1358118..966b736 100644 --- a/tests/a2a-long-horizon.test.ts +++ b/tests/a2a-long-horizon.test.ts @@ -53,9 +53,25 @@ class InputRequiringSandbox implements SandboxBox { const seq = this.sequences[this.callIdx] this.callIdx += 1 if (!seq) throw new Error(`InputRequiringSandbox: out of canned sequences at call ${this.callIdx}`) + let output = '' for (const delta of seq.chunks) { + output += delta yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta } } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: Math.ceil(output.length / 4), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: (1 + Math.ceil(output.length / 4)) * 0.00002, + budgetEnforced: true, + }, + }, + } if (seq.pause) { yield { type: 'input-required', data: { inputRequired: { prompt: seq.pause.prompt } } } } @@ -65,9 +81,25 @@ class InputRequiringSandbox implements SandboxBox { class StubSandbox implements SandboxBox { constructor(private chunks: string[]) {} async *streamPrompt(): AsyncIterable { + let output = '' for (const delta of this.chunks) { + output += delta yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta } } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: Math.ceil(output.length / 4), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: (1 + Math.ceil(output.length / 4)) * 0.00002, + budgetEnforced: true, + }, + }, + } } } diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts new file mode 100644 index 0000000..4221d01 --- /dev/null +++ b/tests/a2a-payment-races.test.ts @@ -0,0 +1,183 @@ +import { Hono } from 'hono' +import { describe, expect, it } from 'vitest' + +import { InMemoryTaskStore } from '../src/a2a/task-store' +import { createAgentGateway } from '../src/middleware' +import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryPaymentOperations } from '../src/payment-operations' +import type { AgentMeta, GatewayConfig, SandboxBox } from '../src/types' + +const operatorAddress = '0x1111111111111111111111111111111111111111' +const commitment = `0x${'ab'.repeat(32)}` + +const agent: AgentMeta = { + id: 'agent-a2a-races', + ownerId: 'owner', + slug: 'a2a-races', + systemPrompt: '', + pricePerTokenUsd: 0.000001, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +function paymentHeader(nonce: string): string { + return JSON.stringify({ + commitment, + signature: '0xsig', + operator: operatorAddress, + amount: '1000000000', + nonce, + expiry: String(Math.floor(Date.now() / 1000) + 300), + }) +} + +function message(text: string, taskId: string) { + return { + kind: 'message', + role: 'user', + taskId, + contextId: 'ctx-race', + messageId: `message-${text}`, + parts: [{ kind: 'text', text }], + } +} + +function usage() { + return { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.000002, + budgetEnforced: true, + } +} + +describe('A2A payment ownership races', () => { + it('allows only one concurrent continuation to claim and settle a task', async () => { + const taskStore = new InMemoryTaskStore() + await taskStore.put({ + kind: 'task', + id: 'task-continuation', + contextId: 'ctx-race', + status: { state: 'input-required', timestamp: new Date().toISOString() }, + history: [message('initial', 'task-continuation')], + }) + let runs = 0 + let settlements = 0 + const box: SandboxBox = { + async *streamPrompt() { + runs += 1 + await new Promise((resolve) => setTimeout(resolve, 5)) + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'done' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + settlePayment: async () => { settlements += 1 }, + verifyApiKey: async () => ({ consumerId: 'consumer', keyId: 'key', scopes: ['chat'] }), + x402: { operatorAddress, chainId: 1, demoMode: true }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const request = (text: string) => app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', Authorization: 'Bearer sk_agent_race' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: text, + method: 'message/send', + params: { message: message(text, 'task-continuation') }, + }), + }) + + const [first, second] = await Promise.all([request('one'), request('two')]) + const bodies = await Promise.all([first.json(), second.json()]) as Array<{ + result?: { status?: { state?: string }; history?: unknown[] } + error?: { code?: number } + }> + expect(bodies.filter((body) => body.result?.status?.state === 'completed')).toHaveLength(1) + expect(bodies.filter((body) => body.error?.code === -32602)).toHaveLength(1) + expect(runs).toBe(1) + expect(settlements).toBe(1) + const finalTask = await taskStore.get('task-continuation') + expect(finalTask?.history).toHaveLength(2) + }) + + it('retains ownership when cancellation follows delivered output without a receipt', async () => { + const taskStore = new InMemoryTaskStore() + const operations = new MemoryPaymentOperations() + let outputSeen!: () => void + const outputReady = new Promise((resolve) => { outputSeen = resolve }) + let releaseSandbox!: () => void + const sandboxReleased = new Promise((resolve) => { releaseSandbox = resolve }) + const box: SandboxBox = { + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + outputSeen() + await sandboxReleased + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const streamPromise = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('77'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/stream', + params: { message: message('run', 'task-cancel') }, + }), + }) + const stream = await streamPromise + const reader = stream.body!.getReader() + await outputReady + + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-cancel' }, + }), + }) + expect(cancel.status).toBe(200) + releaseSandbox() + while (!(await reader.read()).done) { + // Drain the stream so its recovery path runs. + } + + expect(operations.get(`x402:${commitment}:77`)?.state).toBe('claimed') + const canceled = await taskStore.get('task-cancel') + expect(canceled?.status.state).toBe('canceled') + }) +}) diff --git a/tests/a2a.test.ts b/tests/a2a.test.ts index 355bddd..a51c1da 100644 --- a/tests/a2a.test.ts +++ b/tests/a2a.test.ts @@ -37,10 +37,26 @@ class StubSandbox implements SandboxBox { private opts: { delayMs?: number } = {}, ) {} async *streamPrompt(): AsyncIterable { + let output = '' for (const delta of this.chunks) { if (this.opts.delayMs) await new Promise((r) => setTimeout(r, this.opts.delayMs)) + output += delta yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta } } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: Math.ceil(output.length / 4), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: (1 + Math.ceil(output.length / 4)) * 0.00002, + budgetEnforced: true, + }, + }, + } } } diff --git a/tests/integration-x402.test.ts b/tests/integration-x402.test.ts index c89b4e5..9bd4f7d 100644 --- a/tests/integration-x402.test.ts +++ b/tests/integration-x402.test.ts @@ -48,7 +48,7 @@ const CHAIN_ID = 3799 const OPERATOR_PRIVATE_KEY = generatePrivateKey() const OPERATOR_ADDRESS = privateKeyToAccount(OPERATOR_PRIVATE_KEY).address const CREDITS_ADDRESS: Hex = '0x00000000000000000000000000000000DeaDBeef' -const FUNDED_REQUEST_AMOUNT = 100_000n +const FUNDED_REQUEST_AMOUNT = 1_000_000n const domain = { name: 'ShieldedCredits', @@ -145,9 +145,25 @@ async function verifySignerOnChain(payload: Record): Promise { + let output = '' for (const delta of this.chunks) { + output += delta yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta } } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: Math.ceil(output.length / 4), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: (1 + Math.ceil(output.length / 4)) * 0.00005, + budgetEnforced: true, + }, + }, + } } } diff --git a/tests/kv-stores.test.ts b/tests/kv-stores.test.ts index 4079895..55b5825 100644 --- a/tests/kv-stores.test.ts +++ b/tests/kv-stores.test.ts @@ -23,6 +23,12 @@ class StubKV implements NonceKV, RlKV { this.store.set(key, { value, expiresAt: this.now() + ttl * 1000 }) } + async putIfAbsent(key: string, value: string, options?: { expirationTtl?: number }): Promise { + if (await this.get(key) !== null) return false + await this.put(key, value, options) + return true + } + async delete(key: string): Promise { this.store.delete(key) } @@ -79,6 +85,18 @@ describe('KvNonceStore', () => { // Isolate B sees the same nonce as used — no cross-isolate bypass expect(await isolateB.hasSeen('shared-nonce')).toBe(true) }) + + it('keeps same-owner claims idempotent when isolates race on the atomic insert', async () => { + const sharedKv = new StubKV() + const isolateA = new KvNonceStore(sharedKv) + const isolateB = new KvNonceStore(sharedKv) + const results = await Promise.all([ + isolateA.claim('shared-operation', 300, 'operation-1'), + isolateB.claim('shared-operation', 300, 'operation-1'), + ]) + expect(results).toEqual([true, true]) + expect(await isolateA.claim('shared-operation', 300, 'operation-2')).toBe(false) + }) }) describe('KvRateLimitStore', () => { diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index fc74e94..1515c46 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -18,27 +18,49 @@ import type { } from '../src/types' import { MemoryNonceStore } from '../src/nonce-store' import { MemoryRateLimitStore } from '../src/rate-limit' +import { MemoryPaymentOperations } from '../src/payment-operations' const operatorAddress = '0x1111111111111111111111111111111111111111' -const fundedRequestAmount = '100000' +const fundedRequestAmount = '1000000' /** Sandbox that emits a fixed reply, captures the prompt + opts for assertion */ class StubSandbox implements SandboxBox { receivedPrompt: string | null = null - receivedOpts: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number } | undefined + receivedOpts: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number; executionBudget?: unknown } | undefined constructor(private chunks: string[]) {} async *streamPrompt( message: string, - opts?: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number }, + opts?: { sessionId?: string; systemPrompt?: string; maxOutputTokens?: number; executionBudget?: unknown }, ): AsyncIterable { this.receivedPrompt = message this.receivedOpts = opts + let remaining = (opts?.maxOutputTokens ?? 1024) * 4 + let output = '' for (const delta of this.chunks) { + const bounded = delta.slice(0, remaining) + remaining -= bounded.length + output += bounded + if (!bounded) break yield { type: 'message.part.updated', - data: { part: { type: 'text' }, delta }, + data: { part: { type: 'text' }, delta: bounded }, } + if (bounded.length < delta.length) break + } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: Math.ceil(output.length / 4), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: (1 + Math.ceil(output.length / 4)) * 0.00002, + budgetEnforced: true, + }, + }, } } } @@ -250,7 +272,7 @@ describe('POST /:slug/chat/completions — auth paths', () => { const body = await res.json() as { error: { payment_methods: string[]; x402: Record } } expect(body.error.payment_methods).toContain('x402') expect(body.error.x402.operator).toBe(operatorAddress) - expect(body.error.x402.required_amount).toBe('21020') + expect(body.error.x402.required_amount).toBe('184861') expect(body.error.x402.max_output_tokens).toBe(1024) }) @@ -285,7 +307,7 @@ describe('POST /:slug/chat/completions — auth paths', () => { method: 'POST', headers: { 'Content-Type': 'application/json', - 'X-Payment-Signature': buildSpendAuth({ amount: '580' }), + 'X-Payment-Signature': buildSpendAuth({ amount: '1000000' }), }, body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }], @@ -734,6 +756,20 @@ describe('createAgentGateway — production-config guard', () => { x402: { operatorAddress, chainId: 3799, verifySigner: async () => true }, })).not.toThrow() }) + + it('requires an explicit version when durable payment operations are configured', () => { + expect(() => createAgentGateway({ + resolveAgent: async () => null, + getSandbox: async () => ({ async *streamPrompt() { /* unused */ } }), + recordUsage: async () => { /* unused */ }, + x402: { + operatorAddress, + chainId: 3799, + demoMode: true, + paymentOperations: new MemoryPaymentOperations(), + }, + })).toThrow(/paymentProtocolVersion must be explicit/) + }) }) describe('POST /:slug/chat/completions — error safety', () => { diff --git a/tests/observer.test.ts b/tests/observer.test.ts index ef539cd..edacaa9 100644 --- a/tests/observer.test.ts +++ b/tests/observer.test.ts @@ -111,9 +111,25 @@ describe('ConsoleObserver', () => { class StubSandbox implements SandboxBox { constructor(private chunks: string[]) {} async *streamPrompt(): AsyncIterable { + let output = '' for (const delta of this.chunks) { + output += delta yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta } } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: Math.ceil(output.length / 4), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: (1 + Math.ceil(output.length / 4)) * 0.00002, + budgetEnforced: true, + }, + }, + } } } @@ -140,7 +156,7 @@ function buildSpendAuth(overrides: Record = {}): string { return JSON.stringify({ commitment: '0xAlice', signature: '0xsig', - amount: '100000', + amount: '1000000', nonce: String(Math.floor(Math.random() * 1e9)), operator: operatorAddress, expiry: String(now + 600), diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts new file mode 100644 index 0000000..2d9158b --- /dev/null +++ b/tests/payment-operations.test.ts @@ -0,0 +1,433 @@ +import { describe, expect, it } from 'vitest' +import { Hono } from 'hono' +import { + dispatchSandboxStreamRich, + maximumBillableInputTokens, + requiredX402Amount, +} from '../src/dispatch' +import { createAgentGateway } from '../src/middleware' +import { + MemoryPaymentOperations, + type PaymentAuthorizationContext, +} from '../src/payment-operations' +import { MemoryNonceStore } from '../src/nonce-store' +import type { + AgentMeta, + GatewayConfig, + SandboxBox, + SandboxStreamEvent, +} from '../src/types' + +const agent: AgentMeta = { + id: 'agent-payment-tests', + ownerId: 'owner', + slug: 'payment-tests', + systemPrompt: '', + pricePerTokenUsd: 0.00002, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +function payload(amount = '1000', nonce = '1', expiry = String(Math.floor(Date.now() / 1000) + 300)) { + return { + commitment: `0x${'ab'.repeat(32)}`, + signature: '0xsig', + operator: '0x1', + amount, + nonce, + expiry, + } +} + +function context(requiredAmount = 500n): PaymentAuthorizationContext { + return { + requestId: 'request-1', + agentId: agent.id, + requiredAmount, + maxOutputTokens: 4, + executionBudget: { + maxInputTokens: 10, + maxOutputTokens: 4, + maxReasoningTokens: 4, + maxToolTokens: 4, + maxToolCalls: 2, + maxProviderCostUsd: 1, + }, + } +} + +describe('version 2 payment operations', () => { + it('has one atomic owner and refunds the unused reservation', async () => { + const operations = new MemoryPaymentOperations() + const results = await Promise.allSettled([ + operations.claimPayment(payload(), context()), + operations.claimPayment(payload(), { ...context(), requestId: 'request-2' }), + ]) + expect(results.filter((result) => result.status === 'fulfilled')).toHaveLength(1) + + const owner = results.find((result): result is PromiseFulfilledResult>> => result.status === 'fulfilled')!.value + const settled = await operations.settlePayment(owner, { + amount: 200n, + totalCostUsd: 0.2, + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 1, + toolTokens: 1, + toolCallCount: 1, + providerCostUsd: 0.2, + budgetEnforced: true, + }, + }) + expect(settled.state).toBe('settled') + expect(settled.settledAmount).toBe(200n) + expect(settled.refundAmount).toBe(800n) + await expect(operations.settlePayment(owner, { + amount: 1001n, + totalCostUsd: 1, + usage: { ...settledUsage(), budgetEnforced: true }, + })).rejects.toThrow() + }) + + it('reclaims a claimed operation after expiry', async () => { + let now = 100 + let recoveredCalls = 0 + const operations = new MemoryPaymentOperations({ + now: () => now, + onReclaim: async () => { recoveredCalls += 1 }, + }) + const owner = await operations.claimPayment(payload('1000', '2', '200'), context()) + now = 201 + const reclaimed = await operations.reclaimPayment(owner.operationId) + expect(reclaimed.state).toBe('reclaimed') + expect(reclaimed.refundAmount).toBe(1000n) + expect(recoveredCalls).toBe(1) + }) + + it('recovers a release after a worker crash between state and side effect', async () => { + let fail = true + const operations = new MemoryPaymentOperations({ + onRelease: async () => { + if (fail) throw new Error('worker crashed during release') + }, + onReclaim: async () => undefined, + }) + const owner = await operations.claimPayment(payload('1000', '11'), context()) + + await expect(operations.releasePayment(owner, 'sandbox failed')).rejects.toThrow('worker crashed') + expect(operations.get(owner.operationId)?.state).toBe('releasing') + fail = false + const recovered = await operations.reclaimPayment(owner.operationId) + expect(recovered.state).toBe('released') + }) + + it('runs one settlement side effect for concurrent retries', async () => { + let effects = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { + effects += 1 + await new Promise((resolve) => setTimeout(resolve, 1)) + }, + }) + const owner = await operations.claimPayment(payload('1000', '13'), context()) + const input = { + amount: 200n, + totalCostUsd: 0.2, + usage: settledUsage(), + } + const settled = await Promise.all([ + operations.settlePayment(owner, input), + operations.settlePayment(owner, input), + ]) + expect(effects).toBe(1) + expect(settled[0].state).toBe('settled') + expect(settled[1].state).toBe('settled') + }) + + it('does not close a release when recovery has no refund proof', async () => { + const operations = new MemoryPaymentOperations({ + onRelease: async () => { throw new Error('acknowledgement lost') }, + }) + const owner = await operations.claimPayment(payload('1000', '14'), context()) + await expect(operations.releasePayment(owner, 'sandbox failed')).rejects.toThrow() + await expect(operations.reclaimPayment(owner.operationId)).rejects.toThrow('recovery is not configured') + expect(operations.get(owner.operationId)?.state).toBe('releasing') + }) + + it('reclaims a claim that crashed after durable ownership but before completion', async () => { + let now = 100 + let recoveries = 0 + const operations = new MemoryPaymentOperations({ + now: () => now, + onClaim: async () => { throw new Error('worker crashed during claim') }, + onReclaim: async () => { recoveries += 1 }, + }) + + await expect(operations.claimPayment(payload('1000', '12', '200'), context())).rejects.toThrow('worker crashed') + expect(operations.get('x402:0x' + 'ab'.repeat(32) + ':12')?.state).toBe('claiming') + now = 201 + const recovered = await operations.reclaimPayment('x402:0x' + 'ab'.repeat(32) + ':12') + expect(recovered.state).toBe('reclaimed') + expect(recoveries).toBe(1) + }) +}) + +describe('bounded request pricing and sandbox receipts', () => { + it('quotes conservative UTF-8 input and hidden provider costs', () => { + expect(maximumBillableInputTokens({ ...agent, systemPrompt: '' }, '😀')).toBe(4) + expect(requiredX402Amount(0.000001, 4, 2, 6, 3, 5, 0.00002)).toBe(20n) + }) + + it.each([ + ['honors max output', false], + ['ignores max output', true], + ])('passes and enforces maxOutputTokens when a Sandbox implementation %s', async (_name, ignoresLimit) => { + let received: Record | undefined + const box: SandboxBox = { + async *streamPrompt(_message, opts) { + received = opts as unknown as Record + yield { + type: 'message.part.updated', + data: { part: { type: 'text' }, delta: ignoresLimit ? '01234567890123456789' : 'abcd' }, + } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: ignoresLimit ? 6 : 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.00002, + budgetEnforced: !ignoresLimit, + }, + }, + } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + executionBudget: { + maxReasoningTokens: 2, + maxToolTokens: 2, + maxToolCalls: 1, + maxProviderCostUsd: 1, + }, + } + const run = async () => { + const events = [] + for await (const event of dispatchSandboxStreamRich(agent, 'hi', 'consumer', config, undefined, undefined, 4)) { + events.push(event) + } + return events + } + if (ignoresLimit) await expect(run()).rejects.toThrow(/max output|budget/) + else expect(await run()).toHaveLength(3) + expect(received?.maxOutputTokens).toBe(4) + }) + + it('rejects hidden reasoning, tool, and provider spend beyond the receipt budget', async () => { + const box: SandboxBox = { + async *streamPrompt() { + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 3, + toolTokens: 3, + toolCallCount: 2, + providerCostUsd: 2, + budgetEnforced: true, + }, + }, + } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + executionBudget: { + maxReasoningTokens: 2, + maxToolTokens: 2, + maxToolCalls: 1, + maxProviderCostUsd: 1, + }, + } + const run = async () => { + for await (const _event of dispatchSandboxStreamRich(agent, 'hi', 'consumer', config, undefined, undefined, 4)) { + // The receipt is intentionally consumed only to exercise the generator. + } + } + await expect(run()).rejects.toThrow(/reasoning|tool|provider/) + }) + + it('bounds output before delivery and retains payment ownership after an over-limit stream', async () => { + let releases = 0 + const operations = new MemoryPaymentOperations({ onRelease: async () => { releases += 1 } }) + const box: SandboxBox = { + async *streamPrompt() { + yield { + type: 'message.part.updated', + data: { part: { type: 'text' }, delta: '0123456789abcdefghijklmnop' }, + } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + maxOutputTokens: 4, + defaultOutputTokens: 4, + x402: { + operatorAddress: '0x1', + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const response = await app.request('/v1/agents/payment-tests/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': JSON.stringify({ + ...payload('1000000000', '15'), + operator: '0x1', + }), + }, + body: JSON.stringify({ max_tokens: 4, messages: [{ role: 'user', content: 'hi' }] }), + }) + const wire = await response.text() + expect(wire).toContain('sandbox exceeded max output tokens') + expect(wire).not.toContain('0123456789abcdefghijklmnop') + expect(releases).toBe(0) + expect(operations.get(`x402:${'0x' + 'ab'.repeat(32)}:15`)?.state).toBe('claimed') + }) + + it('retains hidden usage when a final provider receipt omits it', async () => { + const box: SandboxBox = { + async *streamPrompt() { + yield { type: 'task.reasoning', data: { reasoning: { tokens: 1 } } } + yield { type: 'task.tool.updated', data: { tool: { inputTokens: 1, outputTokens: 0 } } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.1, + budgetEnforced: true, + }, + }, + } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + executionBudget: { + maxReasoningTokens: 2, + maxToolTokens: 2, + maxToolCalls: 1, + maxProviderCostUsd: 1, + }, + } + const run = async () => { + const events = [] + for await (const _event of dispatchSandboxStreamRich(agent, 'hi', 'consumer', config, undefined, undefined, 4)) { + events.push(_event) + } + return events.find((event) => event.kind === 'usage') + } + await expect(run()).resolves.toMatchObject({ + kind: 'usage', + usage: { reasoningTokens: 1, toolTokens: 1, toolCallCount: 1 }, + }) + }) + + it('does not let a lower final receipt erase earlier hidden spend or a failed budget flag', async () => { + const box: SandboxBox = { + async *streamPrompt() { + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 5, + toolTokens: 4, + toolCallCount: 2, + providerCostUsd: 0.9, + budgetEnforced: false, + }, + }, + } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.1, + budgetEnforced: true, + }, + }, + } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + executionBudget: { + maxReasoningTokens: 10, + maxToolTokens: 10, + maxToolCalls: 10, + maxProviderCostUsd: 1, + }, + } + const run = async () => { + const events = [] + for await (const event of dispatchSandboxStreamRich(agent, 'hi', 'consumer', config, undefined, undefined, 4)) { + events.push(event) + } + return events.find((event) => event.kind === 'usage') + } + await expect(run()).rejects.toThrow(/budget|provider|reasoning|tool/) + }) +}) + +function settledUsage() { + return { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0, + budgetEnforced: true, + } +} diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts new file mode 100644 index 0000000..da37cbc --- /dev/null +++ b/tests/protocol-guards.test.ts @@ -0,0 +1,210 @@ +import { Hono } from 'hono' +import { describe, expect, it } from 'vitest' +import { InMemoryTaskStore } from '../src/a2a/task-store' +import { A2A_ERROR_CODES, type Task } from '../src/a2a/types' +import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryPaymentOperations } from '../src/payment-operations' +import { MemoryRateLimitStore } from '../src/rate-limit' +import { createAgentGateway } from '../src/middleware' +import { verifyX402 } from '../src/verify' +import type { AgentMeta, GatewayConfig, SandboxBox } from '../src/types' + +const operatorAddress = '0x1111111111111111111111111111111111111111' +const commitment = `0x${'ab'.repeat(32)}` + +const agent: AgentMeta = { + id: 'agent-guards', + ownerId: 'owner', + slug: 'guards', + pricePerTokenUsd: 0.00002, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +function paymentHeader(nonce = '1'): string { + return JSON.stringify({ + commitment, + signature: '0xsignature', + amount: '1000000', + nonce, + operator: operatorAddress, + expiry: String(Math.floor(Date.now() / 1000) + 600), + }) +} + +function box(): SandboxBox { + return { + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'ok' } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.00002, + budgetEnforced: true, + }, + }, + } + }, + } +} + +describe('final payment boundary protocol guards', () => { + it('does not claim payment before the host authorization guard', async () => { + let claims = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box(), + recordUsage: async () => undefined, + authorizeConsumer: async () => ({ allow: false, reason: 'blocked', code: 'blocked' }), + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + authorizePayment: async () => { + claims += 1 + return true + }, + }, + nonceStore: new MemoryNonceStore(), + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const response = await app.request('/v1/agents/guards/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('2') }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + expect(response.status).toBe(403) + expect(claims).toBe(0) + }) + + it('checks an A2A task state before claiming a payment', async () => { + const taskStore = new InMemoryTaskStore() + const terminal: Task = { + kind: 'task', + id: 'existing-task', + contextId: 'ctx', + status: { state: 'completed', timestamp: new Date().toISOString() }, + history: [], + } + await taskStore.put(terminal) + let claims = 0 + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => box(), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + authorizePayment: async () => { + claims += 1 + return true + }, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore }, + })) + + const response = await app.request('/v1/agents/guards', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('3') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: terminal.id, + parts: [{ kind: 'text', text: 'retry' }], + }, + }, + }), + }) + const body = await response.json() as { error?: { code: number } } + expect(body.error?.code).toBe(A2A_ERROR_CODES.INVALID_PARAMS) + expect(claims).toBe(0) + }) + + it('shares one canonical nonce authority across mixed verifier versions', async () => { + const store = new MemoryNonceStore() + const config = { operatorAddress, chainId: 1, demoMode: true } + const header = paymentHeader('4') + const upperHeader = header.replace(commitment, commitment.toUpperCase()) + + await expect(verifyX402(header, config, store, 1n, false)).resolves.toBe(commitment) + await expect(verifyX402(upperHeader, config, store, 1n, false)).resolves.toBe(commitment.toUpperCase()) + await expect(verifyX402(header, config, store, 1n, true)).resolves.toBe(commitment) + expect(await store.claim(`${commitment}:4`, 60)).toBe(false) + }) + + it('lets only one mixed-version gateway own a payment at the final boundary', async () => { + const nonceStore = new MemoryNonceStore() + const operations = new MemoryPaymentOperations() + let enteredLegacy!: () => void + const legacyEntered = new Promise((resolve) => { enteredLegacy = resolve }) + let releaseLegacy!: () => void + const legacyRelease = new Promise((resolve) => { releaseLegacy = resolve }) + const shared = { + resolveAgent: async () => agent, + getSandbox: async () => box(), + recordUsage: async () => undefined, + nonceStore, + } + const legacy = new Hono() + legacy.route('/v1/agents', createAgentGateway({ + ...shared, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => { + enteredLegacy() + await legacyRelease + return true + }, + }, + })) + const modern = new Hono() + modern.route('/v1/agents', createAgentGateway({ + ...shared, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + })) + + const request = (app: Hono) => app.request('/v1/agents/guards/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('5') }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'mixed deploy' }] }), + }) + const legacyResponsePromise = request(legacy) + await legacyEntered + const modernResponse = await request(modern) + await modernResponse.text() + releaseLegacy() + const legacyResponse = await legacyResponsePromise + await legacyResponse.text() + + expect(modernResponse.status).toBe(200) + expect(legacyResponse.status).toBe(402) + expect(operations.get(`x402:${commitment}:5`)?.state).toBe('settled') + }) +}) diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 7ceb70a..5742673 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -171,6 +171,25 @@ describe('verifyMpp', () => { expect(seen[0].credential).toContain('commitment') }) + it('shares the x402 nonce authority with equivalent BlueprinTEVM credentials', async () => { + const nonceStore = new MemoryNonceStore() + const payload = { + commitment: '0xAlice', + signature: '0xSignatureBytes', + operator: operatorAddress, + amount: '20000', + nonce: '01', + expiry: String(Math.floor(Date.now() / 1000) + 600), + } + expect(await verifyX402(JSON.stringify({ ...payload, nonce: '1' }), baseConfig, nonceStore)).toBe('0xAlice') + expect(await verifyMpp( + buildCredential({ ...payload, commitment: '0xALICE' }), + mppConfig, + baseConfig, + nonceStore, + )).toBeNull() + }) + it('rejects an underfunded blueprintevm credential before its verifier can reserve funds', async () => { let calls = 0 const header = buildCredential({ From 0ba929e43feca54f9f7bb92d0da471adfe94eea9 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 21:08:05 -0600 Subject: [PATCH 03/32] fix(payments): harden reservation recovery and task access --- README.md | 5 + docs/a2a-long-horizon.md | 10 +- src/a2a/handler.ts | 70 ++++++++++++++ src/a2a/push-notifications.ts | 4 + src/a2a/types.ts | 1 + src/dispatch.ts | 87 +++++++++++++---- src/middleware.ts | 24 +++-- src/payment-operations.ts | 63 +++++++++++-- src/types.ts | 26 ++++-- tests/a2a-long-horizon.test.ts | 25 +++++ tests/a2a-payment-races.test.ts | 4 +- tests/a2a.test.ts | 27 ++++++ tests/middleware.test.ts | 76 +++++++++++++++ tests/payment-operations.test.ts | 156 +++++++++++++++++++++++++++++++ tests/protocol-guards.test.ts | 8 +- 15 files changed, 537 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index d3f9d89..5cdbad7 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,8 @@ Keep version 1 explicitly configured while old and new gateways coexist; shared Before it calls the verifier, the gateway requires the signed amount to cover filtered input plus the requested output limit. The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. +Sandbox adapters should emit a complete `sandbox.usage` receipt. +Version 2 payment operations reject missing receipts; legacy adapters use visible-token estimates only. MPP is method-specific. Configure `mpp.verifySigner` for production MPP credentials; it receives the decoded JSON payload when available plus the original decoded credential, and returns the authenticated consumer ID or `null`. @@ -58,6 +60,9 @@ Wire protocol handlers only translate their request and response shapes. ## A2A protocol The gateway speaks Google's A2A protocol alongside its OpenAI-compatible surface: discovery via `.well-known/agent.json`, JSON-RPC 2.0 dispatch for `message/send`, `message/stream`, `tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and the four `tasks/pushNotificationConfig/*` methods. Long-horizon agents — durable tasks across worker restarts, webhook delivery on terminal state, `input-required` pauses with multi-turn continuation — are documented in [`docs/a2a-long-horizon.md`](./docs/a2a-long-horizon.md). +Production A2A task control requires `a2a.authorizeTaskAccess`; explicit demo mode is the local-test exception. +Custom production task stores must implement atomic `createIfAbsent` and `compareAndSet` methods. +Push destinations must use HTTPS without URL credentials. ## Tier diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 6b95a19..3a8c112 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -13,7 +13,14 @@ All four are gated on configuration — they cost nothing for agents that don't ## Durable tasks (SqlTaskStore) -By default `GatewayConfig.a2a.taskStore` is in-memory: fast, zero-config, fine for tests and single-machine deployments. Production deployments swap in `SqlTaskStore` against any SQL store — D1, postgres, sqlite, libSQL, Turso — via a 2-method `SqlAdapter` shim. +By default `GatewayConfig.a2a.taskStore` is in-memory: fast, zero-config, fine for tests and single-machine deployments. +Production deployments swap in `SqlTaskStore` against any SQL store — D1, postgres, sqlite, libSQL, Turso — via a `SqlAdapter` shim. +Custom production task stores must implement both atomic methods, `createIfAbsent` and `compareAndSet`. +The gateway rejects a custom store without those methods because a read-then-write fallback can run paid work twice across workers. + +Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. +The hook receives the task and request headers so the application can enforce task ownership. +Explicit `x402.demoMode` permits these methods without the hook for local tests only. ### D1 (Cloudflare Workers) @@ -164,6 +171,7 @@ When `pushStore` is set, the agent card advertises `capabilities.pushNotificatio ``` Get / list / delete mirror standard CRUD via the same method namespace. +Push destinations must use HTTPS and must not include URL credentials. ### Webhook receiver shape diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index b87be27..36a4ba9 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -559,6 +559,8 @@ async function handleTasksGet( fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), ) } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError return c.json(ok(req.id, task)) } @@ -578,6 +580,8 @@ async function handleTasksCancel( fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), ) } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError if (isTerminal(task.status.state)) { return c.json( fail( @@ -651,6 +655,8 @@ async function handleTasksResubscribe( fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), ) } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError const final = isTerminal(task.status.state) || task.status.state === 'input-required' const event: TaskStatusUpdateEvent = { kind: 'status-update', @@ -702,6 +708,11 @@ async function handlePushSet( if (!task) { return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.taskId}' not found`)) } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError + if (!isHttpsUrl(params.pushNotificationConfig.url)) { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url must use https')) + } await deps.pushStore.set(params.taskId, params.pushNotificationConfig) const stored = await deps.pushStore.get(params.taskId, params.pushNotificationConfig.id) return c.json(ok(req.id, { taskId: params.taskId, pushNotificationConfig: stored })) @@ -721,6 +732,12 @@ async function handlePushGet( fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id and params.pushNotificationConfigId required'), ) } + const task = await deps.taskStore.get(params.id) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError const cfg = await deps.pushStore.get(params.id, params.pushNotificationConfigId) if (!cfg) { return c.json( @@ -746,6 +763,12 @@ async function handlePushList( if (!params || typeof params.id !== 'string') { return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) } + const task = await deps.taskStore.get(params.id) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError const configs = await deps.pushStore.list(params.id) return c.json(ok(req.id, configs.map((cfg) => ({ taskId: params.id, pushNotificationConfig: cfg })))) } @@ -764,6 +787,12 @@ async function handlePushDelete( fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id and params.pushNotificationConfigId required'), ) } + const task = await deps.taskStore.get(params.id) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await authorizeTaskAccess(c, req, task, deps) + if (accessError) return accessError await deps.pushStore.delete(params.id, params.pushNotificationConfigId) return c.json(ok(req.id, null)) } @@ -924,6 +953,47 @@ async function releaseOrRetainPayment( const FINALIZING_METADATA_KEY = 'gatewayFinalizing' +async function authorizeTaskAccess( + c: Context, + req: JSONRPCRequest, + task: Task, + deps: A2AHandlerDeps, +): Promise { + const authorize = deps.config.a2a?.authorizeTaskAccess + if (!authorize && deps.config.x402.demoMode) return undefined + if (!authorize) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task access authorization is not configured'), + 403, + ) + } + let allowed = false + try { + allowed = await authorize(task, { + method: req.method, + agentSlug: c.req.param('slug') ?? '', + authorization: c.req.header('Authorization') ?? '', + paymentSignature: c.req.header('X-Payment-Signature') ?? '', + }) + } catch (error) { + console.error( + `[a2a] task access authorization failed for ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + } + if (allowed) return undefined + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task access denied'), 403) +} + +function isHttpsUrl(value: string): boolean { + try { + const url = new URL(value) + return url.protocol === 'https:' && url.username === '' && url.password === '' + } catch { + return false + } +} + async function createTask(taskStore: TaskStore, task: Task): Promise { if (taskStore.createIfAbsent) return taskStore.createIfAbsent(task) if (await taskStore.get(task.id)) return false diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index 85e6043..36ba3d5 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -248,6 +248,10 @@ export async function deliverPushNotifications(args: { let result: PushDeliveryResult try { + const url = new URL(config.url) + if (url.protocol !== 'https:' || url.username !== '' || url.password !== '') { + throw new Error('push notification URL must use https without credentials') + } const res = await fetcher(config.url, { method: 'POST', headers, body }) result = { taskId: args.task.id, diff --git a/src/a2a/types.ts b/src/a2a/types.ts index 41cb504..9e02e63 100644 --- a/src/a2a/types.ts +++ b/src/a2a/types.ts @@ -51,6 +51,7 @@ export const A2A_ERROR_CODES = { CONTENT_TYPE_NOT_SUPPORTED: -32005, INVALID_AGENT_RESPONSE: -32006, AUTHENTICATED_EXTENDED_CARD_NOT_CONFIGURED: -32007, + TASK_ACCESS_DENIED: -32008, } as const // ── Message parts ──────────────────────────────────────────────────────── diff --git a/src/dispatch.ts b/src/dispatch.ts index 737761f..86a3643 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -68,8 +68,8 @@ export interface AuthorizedRequest { } function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { - if (!Number.isFinite(value) || value <= 0) { - throw new Error('agent pricePerTokenUsd must be a finite positive number') + if (!Number.isFinite(value) || value < 0) { + throw new Error('agent pricePerTokenUsd must be a finite non-negative number') } const [mantissa, exponentText] = value.toString().toLowerCase().split('e') const exponent = exponentText ? Number(exponentText) : 0 @@ -509,6 +509,13 @@ export async function claimPayment( } let operation: PaymentOperation | undefined if (config.x402.authorizePayment) { + // Version 1 has no durable operation to release if another request wins + // the shared nonce while this callback is still running. Claim first so + // an external reserve or charge cannot happen for a losing request. + const legacyClaimed = config.x402.paymentProtocolVersion !== 2 && authz.paymentNonceKey + ? await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + : undefined + if (legacyClaimed === false) throw new Error('payment nonce was already consumed') const result = await config.x402.authorizePayment(authz.paymentPayload, context) if (!result) throw new Error('payment authorization was rejected') if (typeof result !== 'boolean') { @@ -520,7 +527,7 @@ export async function claimPayment( } else if (config.x402.paymentProtocolVersion === 2) { throw new Error('version 2 payment authorization did not return an operation') - } else if (authz.paymentNonceKey) { + } else if (authz.paymentNonceKey && legacyClaimed === undefined) { const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) if (!claimed) throw new Error('payment nonce was already consumed') } @@ -711,7 +718,7 @@ export async function* dispatchSandboxStreamRich( let observedReasoningTokens = 0 let observedToolTokens = 0 let observedToolCalls = 0 - const bufferedTextParts: string[] = [] + let legacyOutputText = '' const executionBudget: SandboxExecutionBudget = { maxInputTokens: maximumBillableInputTokens(agent, userMessage), maxOutputTokens: outputLimit, @@ -757,23 +764,27 @@ export async function* dispatchSandboxStreamRich( const remainingBytes = maxOutputBytes - outputBytes if (remainingBytes <= 0) throw new Error('sandbox exceeded max output tokens') const bounded = truncateUtf8(event.data.delta, remainingBytes, encoder) - const boundedDelta = bounded.text - outputBytes += bounded.bytes - bufferedTextParts.push(boundedDelta) - yield { kind: 'activity' } if (bounded.truncated) { + yield { kind: 'activity' } throw new Error('sandbox exceeded max output tokens') } + outputBytes += bounded.bytes + legacyOutputText += bounded.text + yield { kind: 'activity' } + yield { kind: 'text', delta: redactSystemPromptFromOutput(bounded.text, agent.systemPrompt) } continue } if (event.type === 'input-required' || event.data?.inputRequired) { - const usage = finalizeUsage( - withObservedUsage(usageParts, observedReasoningTokens, observedToolTokens, observedToolCalls), + const usage = completeUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + userMessage, + legacyOutputText, executionBudget, + config.x402.paymentOperations !== undefined, ) - for (const bufferedText of bufferedTextParts) { - yield { kind: 'text', delta: redactSystemPromptFromOutput(bufferedText, agent.systemPrompt) } - } yield { kind: 'input-required', prompt: event.data?.inputRequired?.prompt } // Terminal for the sandbox stream — sandbox SHOULD stop emitting until // the gateway dispatches a continuation message with the new user input. @@ -781,13 +792,16 @@ export async function* dispatchSandboxStreamRich( return } } - const usage = finalizeUsage( - withObservedUsage(usageParts, observedReasoningTokens, observedToolTokens, observedToolCalls), + const usage = completeUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + userMessage, + legacyOutputText, executionBudget, + config.x402.paymentOperations !== undefined, ) - for (const bufferedText of bufferedTextParts) { - yield { kind: 'text', delta: redactSystemPromptFromOutput(bufferedText, agent.systemPrompt) } - } yield { kind: 'usage', usage } } @@ -832,6 +846,9 @@ export async function settleAndRecord( startMs: authz.startMs, } try { + // Record attribution before settlement. Marketplace adapters use this + // row to map the request to its agent and consumer before deducting funds. + await config.recordUsage(usageEvent) if (authz.paymentOperation && config.x402.paymentOperations) { const amount = actualX402Amount( agent.pricePerTokenUsd, @@ -856,7 +873,6 @@ export async function settleAndRecord( totalCost, ) } - await config.recordUsage(usageEvent) await obs?.onRequestComplete?.(ctx, usageEvent) } catch (err) { const msg = err instanceof Error ? err.message : String(err) @@ -973,6 +989,39 @@ function finalizeUsage( return usage } +function completeUsage( + parts: Partial, + reasoningTokens: number, + toolTokens: number, + toolCallCount: number, + userMessage: string, + outputText: string, + budget: SandboxExecutionBudget, + requiresReceipt: boolean, +): SandboxUsageReceipt { + const observed = withObservedUsage(parts, reasoningTokens, toolTokens, toolCallCount) + if ( + !requiresReceipt && + Object.keys(parts).length === 0 && + reasoningTokens === 0 && + toolTokens === 0 && + toolCallCount === 0 + ) { + // Preserve the pre-receipt SandboxBox contract for legacy API-key + // adapters. Durable payment operations must use provider-enforced usage. + return { + inputTokens: estimateTokens(userMessage), + outputTokens: estimateTokens(outputText), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0, + budgetEnforced: false, + } + } + return finalizeUsage(observed, budget) +} + function actualX402Amount( pricePerTokenUsd: number, inputTokens: number, diff --git a/src/middleware.ts b/src/middleware.ts index 8497115..5262c78 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -231,6 +231,14 @@ export function createAgentGateway(config: GatewayConfig) { // settleAndRecord, so every security and billing guarantee applies uniformly // regardless of which protocol the caller used. const taskStore = config.a2a?.taskStore ?? new InMemoryTaskStore() + if ( + config.a2a?.taskStore && + (!taskStore.createIfAbsent || !taskStore.compareAndSet) + ) { + throw new Error( + 'createAgentGateway: custom A2A taskStore must implement atomic createIfAbsent and compareAndSet', + ) + } const pushStore = config.a2a?.pushStore const a2a = createA2AHandlers({ config, state, taskStore, pushStore }) gw.get('/:slug/.well-known/agent.json', a2a.handleAgentCard) @@ -308,6 +316,14 @@ function streamChatCompletions( if (!usage) throw new Error('sandbox did not provide a usage receipt') + await settleAndRecord( + agent, + authz, + usage, + config, + obs, + ) + const done: ChatCompletionChunk = { id: `chatcmpl-${Date.now()}`, object: 'chat.completion.chunk', @@ -317,14 +333,6 @@ function streamChatCompletions( } controller.enqueue(encoder.encode(`data: ${JSON.stringify(done)}\n\n`)) controller.enqueue(encoder.encode('data: [DONE]\n\n')) - - await settleAndRecord( - agent, - authz, - usage, - config, - obs, - ) } catch (err) { const rawMessage = err instanceof Error ? err.message : String(err) // Never expose stack traces / absolute paths from sandbox internals. diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 8c412cf..d35f919 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -79,6 +79,7 @@ export class MemoryPaymentOperations implements PaymentOperations { private readonly operations = new Map() private readonly settleFlights = new Map>() private readonly releaseFlights = new Map>() + private readonly reclaimFlights = new Map>() private readonly now: () => number constructor(private readonly options: MemoryPaymentOperationsOptions = {}) { @@ -145,7 +146,17 @@ export class MemoryPaymentOperations implements PaymentOperations { if (current.settledAmount !== input.amount) throw new Error('payment operation has a different pending settlement') const flight = this.settleFlights.get(current.operationId) if (flight) return flight - return this.recoverSettlement(current, input) + const reclaimFlight = this.reclaimFlights.get(current.operationId) + if (reclaimFlight) return reclaimFlight + const recovery = this.recoverSettlement(current, input) + this.settleFlights.set(current.operationId, recovery) + try { + return await recovery + } finally { + if (this.settleFlights.get(current.operationId) === recovery) { + this.settleFlights.delete(current.operationId) + } + } } else if (current.state !== 'claimed') { throw new Error(`cannot settle payment in state ${current.state}`) } @@ -174,6 +185,8 @@ export class MemoryPaymentOperations implements PaymentOperations { if (current.state === 'releasing') { const flight = this.releaseFlights.get(current.operationId) if (flight) return flight + const reclaimFlight = this.reclaimFlights.get(current.operationId) + if (reclaimFlight) return reclaimFlight return this.runRelease(current, reason) } if (current.state !== 'claimed') throw new Error(`cannot release payment in state ${current.state}`) @@ -189,7 +202,9 @@ export class MemoryPaymentOperations implements PaymentOperations { if (current.state === 'settling') { const flight = this.settleFlights.get(operationId) if (flight) return flight - return this.recoverSettlement(current, { + const recoveryFlight = this.reclaimFlights.get(operationId) + if (recoveryFlight) return recoveryFlight + const recovery = this.recoverSettlement(current, { amount: current.settledAmount, totalCostUsd: 0, usage: { @@ -202,24 +217,42 @@ export class MemoryPaymentOperations implements PaymentOperations { budgetEnforced: true, }, }) + this.reclaimFlights.set(operationId, recovery) + try { + return await recovery + } finally { + if (this.reclaimFlights.get(operationId) === recovery) this.reclaimFlights.delete(operationId) + } } if (current.state === 'releasing') { if (!this.options.onReclaim) throw new Error('payment release recovery is not configured') - await this.options.onReclaim(current) - const released = { ...current, state: 'released' as const } - this.operations.set(operationId, released) - return released + const releaseFlight = this.releaseFlights.get(operationId) + if (releaseFlight) return releaseFlight + const flight = this.reclaimFlights.get(operationId) + if (flight) return flight + const recovery = this.runReclaim(current, 'released') + this.reclaimFlights.set(operationId, recovery) + try { + return await recovery + } finally { + if (this.reclaimFlights.get(operationId) === recovery) this.reclaimFlights.delete(operationId) + } } if (current.state !== 'claiming' && current.state !== 'claimed' && current.state !== 'reclaimable') { throw new Error(`cannot reclaim payment in state ${current.state}`) } + const flight = this.reclaimFlights.get(operationId) + if (flight) return flight if (current.expiresAt > this.now()) throw new Error('payment operation has not expired') const reclaimable = { ...current, state: 'reclaimable' as const } this.operations.set(operationId, reclaimable) - await this.options.onReclaim?.(reclaimable) - const reclaimed = { ...reclaimable, state: 'reclaimed' as const } - this.operations.set(operationId, reclaimed) - return reclaimed + const recovery = this.runReclaim(reclaimable, 'reclaimed') + this.reclaimFlights.set(operationId, recovery) + try { + return await recovery + } finally { + if (this.reclaimFlights.get(operationId) === recovery) this.reclaimFlights.delete(operationId) + } } private async runSettlement( @@ -253,6 +286,16 @@ export class MemoryPaymentOperations implements PaymentOperations { return settled } + private async runReclaim( + operation: PaymentOperation, + finalState: 'released' | 'reclaimed', + ): Promise { + await this.options.onReclaim?.(operation) + const recovered = { ...operation, state: finalState as PaymentOperationState } + this.operations.set(operation.operationId, recovered) + return recovered + } + private async runRelease( releasing: PaymentOperation, reason: string, diff --git a/src/types.ts b/src/types.ts index abdc2ec..ce7e7aa 100644 --- a/src/types.ts +++ b/src/types.ts @@ -218,7 +218,7 @@ export interface SandboxStreamEvent { * the caller (rendered as the input-required message body). */ inputRequired?: { prompt?: string } - /** Provider receipt fields. The final event must include every field. */ + /** Provider receipt fields. Version 2 operations require every field. */ usage?: Partial /** Tool or reasoning events may carry hidden usage without visible text. */ tool?: { name?: string; inputTokens?: number; outputTokens?: number } @@ -285,10 +285,10 @@ export interface GatewayConfig { verifyApiKey?: (authHeader: string) => Promise /** - * Settle payment after successful response. - * For x402: call ShieldedCredits.claimPayment() - * For API key: deduct from spending limit - * Default: no-op (demo mode). + * Settle a legacy payment after usage attribution is recorded. + * Version 2 x402 operations use `x402.paymentOperations` instead. + * For API keys, deduct from the spending limit. + * Default: no-op in explicit demo mode. */ settlePayment?: (payment: PaymentResult, cost: number) => Promise @@ -344,8 +344,22 @@ export interface GatewayConfig { * Auth + rate-limit + injection-filter + authorization all share the * same pipeline as the OpenAI-compat path. `taskStore` defaults to * `InMemoryTaskStore`; swap in D1/postgres/DO for durable deployments. - */ + */ a2a?: { + /** + * Authorize reads, cancellation, resubscription, and push configuration + * for an existing task. Production control methods fail closed when this + * hook is absent; explicit demo mode permits local tests. + */ + authorizeTaskAccess?: ( + task: import('./a2a/types').Task, + context: { + method: string + agentSlug: string + authorization: string + paymentSignature: string + }, + ) => Promise /** * Where tasks live. Defaults to `InMemoryTaskStore`; swap in * `SqlTaskStore` (D1, postgres, sqlite, libSQL) for durability across diff --git a/tests/a2a-long-horizon.test.ts b/tests/a2a-long-horizon.test.ts index 966b736..580ebd3 100644 --- a/tests/a2a-long-horizon.test.ts +++ b/tests/a2a-long-horizon.test.ts @@ -361,6 +361,31 @@ describe('A2A — push notification config RPCs', () => { expect(await pushStore.get(task.id, cfg.id)).toBeUndefined() }) + it('rejects non-HTTPS push destinations', async () => { + const { app } = buildHarness() + const sendRes = await postJsonRpc( + app, + { jsonrpc: '2.0', id: 1, method: 'message/send', params: { message: textMessage('hi') } }, + apiKeyHeader(), + ) + const task = ((await sendRes.json()) as JSONRPCSuccessResponse).result + const res = await postJsonRpc( + app, + { + jsonrpc: '2.0', + id: 2, + method: 'tasks/pushNotificationConfig/set', + params: { + taskId: task.id, + pushNotificationConfig: { id: 'cfg-http', url: 'http://localhost/hook' }, + }, + }, + apiKeyHeader(), + ) + const body = (await res.json()) as JSONRPCErrorResponse + expect(body.error.code).toBe(A2A_ERROR_CODES.INVALID_PARAMS) + }) + it('get returns TASK_NOT_FOUND for an unregistered config', async () => { const { app } = buildHarness() const res = await postJsonRpc( diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 4221d01..fb37f96 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -85,7 +85,7 @@ describe('A2A payment ownership races', () => { verifyApiKey: async () => ({ consumerId: 'consumer', keyId: 'key', scopes: ['chat'] }), x402: { operatorAddress, chainId: 1, demoMode: true }, nonceStore: new MemoryNonceStore(), - a2a: { taskStore }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() app.route('/v1/agents', createAgentGateway(config)) @@ -139,7 +139,7 @@ describe('A2A payment ownership races', () => { paymentOperations: operations, }, nonceStore: new MemoryNonceStore(), - a2a: { taskStore }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() app.route('/v1/agents', createAgentGateway(config)) diff --git a/tests/a2a.test.ts b/tests/a2a.test.ts index a51c1da..0589ae9 100644 --- a/tests/a2a.test.ts +++ b/tests/a2a.test.ts @@ -430,6 +430,33 @@ describe('A2A — message/stream', () => { // ── tasks/get ───────────────────────────────────────────────────────────── describe('A2A — tasks/get', () => { + it('fails closed for production task access without an authorization hook', async () => { + const { app } = buildHarness({ + x402: { operatorAddress, chainId: 3799, verifySigner: async () => true, demoMode: false }, + }) + const sendRes = await postJsonRpc( + app, + 'test-agent', + { + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: textMessage('hi') }, + }, + apiKeyHeader(), + ) + const sent = (await sendRes.json()) as JSONRPCSuccessResponse + + const getRes = await postJsonRpc( + app, + 'test-agent', + { jsonrpc: '2.0', id: 2, method: 'tasks/get', params: { id: sent.result.id } }, + apiKeyHeader(), + ) + const body = (await getRes.json()) as JSONRPCErrorResponse + expect(body.error.code).toBe(A2A_ERROR_CODES.TASK_ACCESS_DENIED) + }) + it('returns the task created by a prior message/send', async () => { const { app } = buildHarness() const sendRes = await postJsonRpc( diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index 1515c46..18a2e5c 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -411,6 +411,66 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(settlements[0].requestId).toBe(usage[0].requestId) }) + it('does not signal a successful stream before settlement succeeds', async () => { + const { app } = buildHarness({ + settlePayment: async () => { throw new Error('settlement unavailable') }, + }) + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ nonce: '9001' }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + const body = await res.text() + expect(res.status).toBe(200) + expect(body).toContain('settlement unavailable') + expect(body).not.toContain('data: [DONE]') + }) + + it('records attribution before settlement so adapters can resolve the charge', async () => { + const order: string[] = [] + const { app } = buildHarness({ + recordUsage: async () => { order.push('record') }, + settlePayment: async () => { order.push('settle') }, + }) + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ nonce: '9002' }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + await res.text() + expect(order).toEqual(['record', 'settle']) + }) + + it('keeps legacy sandbox adapters working with visible-token estimates', async () => { + const { app, usage, settlements } = buildHarness({ + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'legacy' } } + }, + }), + }) + const res = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: 'Bearer sk_agent_legacy', + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + const streamed = await readSse(res) + expect(res.status).toBe(200) + expect(streamed.combinedText).toBe('legacy') + expect(streamed.done).toBe(true) + expect(usage[0]?.outputTokens).toBe(2) + expect(settlements).toHaveLength(1) + }) + it('threads a unique requestId per concurrent request — regression: two same-consumer requests get distinct ids', async () => { const { app, settlements, usage } = buildHarness({}, ['ok']) const requests = await Promise.all([ @@ -770,6 +830,22 @@ describe('createAgentGateway — production-config guard', () => { }, })).toThrow(/paymentProtocolVersion must be explicit/) }) + + it('refuses a custom A2A task store without atomic transitions', () => { + expect(() => createAgentGateway({ + resolveAgent: async () => null, + getSandbox: async () => ({ async *streamPrompt() { /* unused */ } }), + recordUsage: async () => { /* unused */ }, + x402: { operatorAddress, chainId: 3799, demoMode: true }, + a2a: { + taskStore: { + get: async () => undefined, + put: async () => undefined, + delete: async () => undefined, + }, + }, + })).toThrow(/atomic createIfAbsent and compareAndSet/) + }) }) describe('POST /:slug/chat/completions — error safety', () => { diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index 2d9158b..63f429a 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -147,6 +147,125 @@ describe('version 2 payment operations', () => { expect(settled[1].state).toBe('settled') }) + it('runs one reclaim side effect for concurrent expiry retries', async () => { + let now = 100 + let entered = 0 + let release!: () => void + const recoveryReady = new Promise((resolve) => { release = resolve }) + const operations = new MemoryPaymentOperations({ + now: () => now, + onReclaim: async () => { + entered += 1 + await recoveryReady + }, + }) + const owner = await operations.claimPayment(payload('1000', '16', '200'), context()) + now = 201 + const first = operations.reclaimPayment(owner.operationId) + while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) + const second = operations.reclaimPayment(owner.operationId) + release() + const recovered = await Promise.all([first, second]) + expect(entered).toBe(1) + expect(recovered[0].state).toBe('reclaimed') + expect(recovered[1].state).toBe('reclaimed') + }) + + it('runs one settlement-recovery side effect for concurrent crash retries', async () => { + let entered = 0 + let release!: () => void + const recoveryReady = new Promise((resolve) => { release = resolve }) + const operations = new MemoryPaymentOperations({ + onSettle: async () => { throw new Error('worker crashed during settlement') }, + onReclaim: async () => { + entered += 1 + await recoveryReady + }, + }) + const owner = await operations.claimPayment(payload('1000', '17'), context()) + const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage() } + await expect(operations.settlePayment(owner, input)).rejects.toThrow('worker crashed') + const first = operations.reclaimPayment(owner.operationId) + while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) + const second = operations.reclaimPayment(owner.operationId) + release() + const recovered = await Promise.all([first, second]) + expect(entered).toBe(1) + expect(recovered[0].state).toBe('settled') + expect(recovered[1].state).toBe('settled') + }) + + it('runs one settlement-recovery side effect for concurrent settlement retries', async () => { + let entered = 0 + let release!: () => void + const recoveryReady = new Promise((resolve) => { release = resolve }) + const operations = new MemoryPaymentOperations({ + onSettle: async () => { throw new Error('worker crashed during settlement') }, + onReclaim: async () => { + entered += 1 + await recoveryReady + }, + }) + const owner = await operations.claimPayment(payload('1000', '18'), context()) + const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage() } + await expect(operations.settlePayment(owner, input)).rejects.toThrow('worker crashed') + const first = operations.settlePayment(owner, input) + while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) + const second = operations.settlePayment(owner, input) + release() + const recovered = await Promise.all([first, second]) + expect(entered).toBe(1) + expect(recovered[0].state).toBe('settled') + expect(recovered[1].state).toBe('settled') + }) + + it('does not overlap settlement retry with reclaim recovery', async () => { + let entered = 0 + let release!: () => void + const recoveryReady = new Promise((resolve) => { release = resolve }) + const operations = new MemoryPaymentOperations({ + onSettle: async () => { throw new Error('worker crashed during settlement') }, + onReclaim: async () => { + entered += 1 + await recoveryReady + }, + }) + const owner = await operations.claimPayment(payload('1000', '19'), context()) + const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage() } + await expect(operations.settlePayment(owner, input)).rejects.toThrow('worker crashed') + const reclaim = operations.reclaimPayment(owner.operationId) + while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) + const settle = operations.settlePayment(owner, input) + release() + const recovered = await Promise.all([reclaim, settle]) + expect(entered).toBe(1) + expect(recovered[0].state).toBe('settled') + expect(recovered[1].state).toBe('settled') + }) + + it('does not overlap release retry with reclaim recovery', async () => { + let entered = 0 + let release!: () => void + const recoveryReady = new Promise((resolve) => { release = resolve }) + const operations = new MemoryPaymentOperations({ + onRelease: async () => { throw new Error('worker crashed during release') }, + onReclaim: async () => { + entered += 1 + await recoveryReady + }, + }) + const owner = await operations.claimPayment(payload('1000', '20'), context()) + await expect(operations.releasePayment(owner, 'sandbox failed')).rejects.toThrow('worker crashed') + const reclaim = operations.reclaimPayment(owner.operationId) + while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) + const retry = operations.releasePayment(owner, 'retry release') + release() + const recovered = await Promise.all([reclaim, retry]) + expect(entered).toBe(1) + expect(recovered[0].state).toBe('released') + expect(recovered[1].state).toBe('released') + }) + it('does not close a release when recovery has no refund proof', async () => { const operations = new MemoryPaymentOperations({ onRelease: async () => { throw new Error('acknowledgement lost') }, @@ -179,6 +298,7 @@ describe('bounded request pricing and sandbox receipts', () => { it('quotes conservative UTF-8 input and hidden provider costs', () => { expect(maximumBillableInputTokens({ ...agent, systemPrompt: '' }, '😀')).toBe(4) expect(requiredX402Amount(0.000001, 4, 2, 6, 3, 5, 0.00002)).toBe(20n) + expect(requiredX402Amount(0, 4, 2)).toBe(0n) }) it.each([ @@ -418,6 +538,42 @@ describe('bounded request pricing and sandbox receipts', () => { } await expect(run()).rejects.toThrow(/budget|provider|reasoning|tool/) }) + + it('delivers text before a later sandbox event arrives', async () => { + let release!: () => void + const laterEvent = new Promise((resolve) => { release = resolve }) + const box: SandboxBox = { + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'first' } } + await laterEvent + yield { + type: 'sandbox.usage', + data: { usage: settledUsage() }, + } + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => box, + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + } + const events = dispatchSandboxStreamRich( + agent, + 'hi', + 'consumer', + config, + undefined, + undefined, + 4, + )[Symbol.asyncIterator]() + await events.next() + const firstText = await events.next() + expect(firstText.value).toEqual({ kind: 'text', delta: 'first' }) + release() + await events.next() + await events.next() + }) }) function settledUsage() { diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index da37cbc..395f40d 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -203,8 +203,10 @@ describe('final payment boundary protocol guards', () => { const legacyResponse = await legacyResponsePromise await legacyResponse.text() - expect(modernResponse.status).toBe(200) - expect(legacyResponse.status).toBe(402) - expect(operations.get(`x402:${commitment}:5`)?.state).toBe('settled') + // The legacy gateway claims the shared nonce before its external callback. + // The modern gateway must lose without invoking a second payment owner. + expect(modernResponse.status).toBe(402) + expect(legacyResponse.status).toBe(200) + expect(operations.get(`x402:${commitment}:5`)).toBeUndefined() }) }) From 5c2e05b88395004cf07bc7243348728f40a08fa2 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 23:21:37 -0600 Subject: [PATCH 04/32] fix(payments): close ownership and recovery races --- src/a2a/handler.ts | 20 +++--- src/dispatch.ts | 16 ++--- src/nonce-store.ts | 28 ++++++++- src/payment-operations.ts | 7 ++- src/verify.ts | 6 +- tests/a2a-payment-races.test.ts | 70 +++++++++++++++++++++ tests/kv-stores.test.ts | 14 +++++ tests/payment-operations.test.ts | 26 ++++++++ tests/protocol-guards.test.ts | 103 +++++++++++++++++++++++++++++++ tests/verify.test.ts | 14 ++++- 10 files changed, 279 insertions(+), 25 deletions(-) diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 36a4ba9..4950cb9 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -830,12 +830,16 @@ async function guardMessageRequest( const authz = guard // Multi-turn continuation: if the caller addressed an existing task that is - // currently in `input-required`, append the new message and transition to - // `working`. Any other taskId (unknown OR pointing at a terminal/working + // currently in `input-required`, append the new message and reserve it as + // `submitted`. The handler transitions it to `working` only after payment + // succeeds, so cancellation during payment cannot start sandbox work. + // Any other taskId (unknown OR pointing at a terminal/working // task) means the caller is starting a fresh task and we mint a new id. if (typeof params.message.taskId === 'string') { const existing = await deps.taskStore.get(params.message.taskId) if (existing) { + const accessError = await authorizeTaskAccess(c, req, existing, deps) + if (accessError) return accessError if (existing.status.state !== 'input-required') { return c.json( fail( @@ -852,7 +856,7 @@ async function guardMessageRequest( } const continued: Task = { ...existing, - status: { state: 'working', timestamp: nowIso() }, + status: { state: 'submitted', timestamp: nowIso() }, history: [...(existing.history ?? []), appendedMessage], } if (!await compareAndSetTask(deps.taskStore, existing, continued)) { @@ -860,7 +864,7 @@ async function guardMessageRequest( fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${existing.id}' changed before continuation`), ) } - const claimError = await claimTaskPayment(c, req, continued, authz, deps) + const claimError = await claimTaskPayment(c, req, continued, authz, deps, existing) if (claimError) return claimError return { authz, task: continued } } @@ -898,16 +902,18 @@ async function claimTaskPayment( task: Task, authz: AuthorizedRequest, deps: A2AHandlerDeps, + paymentFailureTask?: Task, ): Promise { try { await claimPayment(authz, deps.config, deps.state) return undefined } catch { await releaseOwnedPayment(authz, deps, 'payment authorization failed') - const failed = withStatus(task, 'failed') + const failed = paymentFailureTask ?? withStatus(task, 'failed') try { - await deps.taskStore.put(failed) - await maybeDeliverPush(failed, deps) + if (await compareAndSetTask(deps.taskStore, task, failed) && isTerminal(failed.status.state)) { + await maybeDeliverPush(failed, deps) + } } catch (taskError) { console.error( `[a2a] failed to persist payment-failed task ${task.id}:`, diff --git a/src/dispatch.ts b/src/dispatch.ts index 86a3643..d09ddda 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -13,7 +13,7 @@ import type { Context } from 'hono' import { filterConsumerMessagesStrict, redactSystemPromptFromOutput } from './filter' import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' import { type RateLimitStore, checkRateLimit } from './rate-limit' -import type { NonceStore } from './nonce-store' +import { claimStoredNonce, type NonceStore } from './nonce-store' import type { PaymentOperation } from './payment-operations' import type { AgentMeta, @@ -513,7 +513,7 @@ export async function claimPayment( // the shared nonce while this callback is still running. Claim first so // an external reserve or charge cannot happen for a losing request. const legacyClaimed = config.x402.paymentProtocolVersion !== 2 && authz.paymentNonceKey - ? await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + ? await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) : undefined if (legacyClaimed === false) throw new Error('payment nonce was already consumed') const result = await config.x402.authorizePayment(authz.paymentPayload, context) @@ -528,13 +528,13 @@ export async function claimPayment( else if (config.x402.paymentProtocolVersion === 2) { throw new Error('version 2 payment authorization did not return an operation') } else if (authz.paymentNonceKey && legacyClaimed === undefined) { - const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) if (!claimed) throw new Error('payment nonce was already consumed') } } else if (config.x402.paymentOperations) { operation = await config.x402.paymentOperations.claimPayment(authz.paymentPayload, context) } else if (authz.paymentNonceKey) { - const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) if (!claimed) throw new Error('payment nonce was already consumed') } if (operation && operation.protocolVersion !== 2) { @@ -544,7 +544,7 @@ export async function claimPayment( throw new Error('durable payment operations are required to settle a claimed operation') } if (operation && authz.paymentNonceKey) { - const claimed = await claimNonce( + const claimed = await claimPaymentNonce( state.nonceStore, authz.paymentNonceKey, authz.paymentPayload, @@ -564,7 +564,7 @@ export async function claimPayment( } authz.paymentOperation = operation } else if (authz.paymentMethod === 'mpp' && authz.paymentNonceKey) { - const claimed = await claimNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) if (!claimed) throw new Error('payment nonce was already consumed') } @@ -627,7 +627,7 @@ export async function reclaimPayment( return config.x402.paymentOperations.reclaimPayment(operationId) } -async function claimNonce( +async function claimPaymentNonce( nonceStore: NonceStore, nonceKey: string, payload: Record, @@ -636,7 +636,7 @@ async function claimNonce( const expiry = Number(payload.expiry ?? Math.floor(Date.now() / 1000) + 3600) const ttl = Math.min(expiry - Math.floor(Date.now() / 1000), 3600) if (!Number.isFinite(ttl) || ttl <= 0) return false - return nonceStore.claim(nonceKey, Math.max(ttl, 60), ownerId) + return claimStoredNonce(nonceStore, nonceKey, Math.max(ttl, 60), ownerId) } /** diff --git a/src/nonce-store.ts b/src/nonce-store.ts index 326cae2..83d82dd 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -8,9 +8,10 @@ export interface NonceStore { hasSeen(nonce: string): Promise /** * Atomically claim a nonce. An owner id makes a retry by the same payment - * operation idempotent while a legacy claim still fails closed. + * operation idempotent while a legacy claim still fails closed. Version 1 + * stores can omit this method and retain their prior check-then-mark path. */ - claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise + claim?(nonce: string, ttlSeconds: number, ownerId?: string): Promise /** Mark nonce as used. TTL = how long to remember it (seconds). */ markSeen(nonce: string, ttlSeconds: number): Promise } @@ -74,7 +75,7 @@ export class MemoryNonceStore implements NonceStore { export interface KVNamespace { get(key: string, options?: { type?: 'text' | 'json' }): Promise put(key: string, value: string, options?: { expirationTtl?: number }): Promise - /** Required for version 2 gateway claims. A plain KV put is not atomic. */ + /** Optional atomic extension. Cloudflare KV does not provide this method. */ putIfAbsent?(key: string, value: string, options?: { expirationTtl?: number }): Promise delete(key: string): Promise } @@ -105,6 +106,11 @@ export class KvNonceStore implements NonceStore { } async claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise { + if (ownerId === undefined) { + if (await this.hasSeen(nonce)) return false + await this.markSeen(nonce, ttlSeconds) + return true + } if (!this.kv.putIfAbsent) { throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') } @@ -130,3 +136,19 @@ export class KvNonceStore implements NonceStore { return `${this.prefix}:${nonce}` } } + +/** Claim through a version 2 store or preserve the version 1 store contract. */ +export async function claimStoredNonce( + store: NonceStore, + nonce: string, + ttlSeconds: number, + ownerId?: string, +): Promise { + if (store.claim) return store.claim(nonce, ttlSeconds, ownerId) + if (ownerId !== undefined) { + throw new Error('NonceStore.claim is required for version 2 payment ownership') + } + if (await store.hasSeen(nonce)) return false + await store.markSeen(nonce, ttlSeconds) + return true +} diff --git a/src/payment-operations.ts b/src/payment-operations.ts index d35f919..993f6dd 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -83,6 +83,9 @@ export class MemoryPaymentOperations implements PaymentOperations { private readonly now: () => number constructor(private readonly options: MemoryPaymentOperationsOptions = {}) { + if (options.onClaim && !options.onReclaim) { + throw new Error('onReclaim is required when onClaim can reserve external funds') + } this.now = options.now ?? (() => Math.floor(Date.now() / 1000)) } @@ -185,9 +188,7 @@ export class MemoryPaymentOperations implements PaymentOperations { if (current.state === 'releasing') { const flight = this.releaseFlights.get(current.operationId) if (flight) return flight - const reclaimFlight = this.reclaimFlights.get(current.operationId) - if (reclaimFlight) return reclaimFlight - return this.runRelease(current, reason) + return this.reclaimPayment(current.operationId) } if (current.state !== 'claimed') throw new Error(`cannot release payment in state ${current.state}`) const releasing = { ...current, state: 'releasing' as const } diff --git a/src/verify.ts b/src/verify.ts index 17e7b6d..969c02e 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -1,5 +1,5 @@ import type { X402Config, MppConfig, ApiKeyInfo, GatewayConfig } from './types' -import type { NonceStore } from './nonce-store' +import { claimStoredNonce, type NonceStore } from './nonce-store' /** Return the canonical opaque nonce key used by the final payment claim. */ export function mppReplayNonceKey(authHeader: string): string | undefined { @@ -103,7 +103,7 @@ export async function verifyX402( if (nonceStore && markNonce) { // Mark seen with TTL matching the expiry window (max 1 hour) const ttl = Math.min(Number(expiry) - Math.floor(Date.now() / 1000), 3600) - const claimed = await nonceStore.claim(nonceKey, Math.max(ttl, 60)) + const claimed = await claimStoredNonce(nonceStore, nonceKey, Math.max(ttl, 60)) if (!claimed) return null } @@ -202,7 +202,7 @@ export async function verifyMpp( ? Math.floor(Date.now() / 1000) + 3600 : Number(payload.expiry) const ttl = Math.min(expiry - Math.floor(Date.now() / 1000), 3600) - const claimed = await nonceStore.claim(nonceKey!, Math.max(ttl, 60)) + const claimed = await claimStoredNonce(nonceStore, nonceKey!, Math.max(ttl, 60)) if (!claimed) return null } diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index fb37f96..73187ae 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -58,6 +58,76 @@ function usage() { } describe('A2A payment ownership races', () => { + it('does not execute a continuation canceled during payment authorization', async () => { + const taskStore = new InMemoryTaskStore() + await taskStore.put({ + kind: 'task', + id: 'task-payment-cancel', + contextId: 'ctx-race', + status: { state: 'input-required', timestamp: new Date().toISOString() }, + history: [message('initial', 'task-payment-cancel')], + }) + let authorizeEntered!: () => void + const authorizationStarted = new Promise((resolve) => { authorizeEntered = resolve }) + let finishAuthorization!: () => void + const authorizationReleased = new Promise((resolve) => { finishAuthorization = resolve }) + let runs = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => { + authorizeEntered() + await authorizationReleased + return true + }, + }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const continuation = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('76') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('continue', 'task-payment-cancel') }, + }), + }) + await authorizationStarted + + const canceled = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-payment-cancel' }, + }), + }) + expect(canceled.status).toBe(200) + finishAuthorization() + const response = await continuation + const body = await response.json() as { error?: { code?: number } } + + expect(body.error?.code).toBe(-32602) + expect(runs).toBe(0) + expect((await taskStore.get('task-payment-cancel'))?.status.state).toBe('canceled') + }) + it('allows only one concurrent continuation to claim and settle a task', async () => { const taskStore = new InMemoryTaskStore() await taskStore.put({ diff --git a/tests/kv-stores.test.ts b/tests/kv-stores.test.ts index 55b5825..32dbbf8 100644 --- a/tests/kv-stores.test.ts +++ b/tests/kv-stores.test.ts @@ -97,6 +97,20 @@ describe('KvNonceStore', () => { expect(results).toEqual([true, true]) expect(await isolateA.claim('shared-operation', 300, 'operation-2')).toBe(false) }) + + it('preserves legacy claims on the standard Cloudflare KV surface', async () => { + const backing = new StubKV() + const cloudflareKv: NonceKV = { + get: backing.get.bind(backing), + put: backing.put.bind(backing), + delete: backing.delete.bind(backing), + } + const store = new KvNonceStore(cloudflareKv) + + expect(await store.claim('legacy', 300)).toBe(true) + expect(await store.claim('legacy', 300)).toBe(false) + await expect(store.claim('version-2', 300, 'operation-1')).rejects.toThrow('atomic putIfAbsent') + }) }) describe('KvRateLimitStore', () => { diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index 63f429a..83d9e95 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -124,6 +124,32 @@ describe('version 2 payment operations', () => { expect(recovered.state).toBe('released') }) + it('reconciles an ambiguous release without repeating its side effect', async () => { + let releases = 0 + let recoveries = 0 + const operations = new MemoryPaymentOperations({ + onRelease: async () => { + releases += 1 + throw new Error('acknowledgement lost') + }, + onReclaim: async () => { recoveries += 1 }, + }) + const owner = await operations.claimPayment(payload('1000', '21'), context()) + + await expect(operations.releasePayment(owner, 'sandbox failed')).rejects.toThrow('acknowledgement lost') + const recovered = await operations.releasePayment(owner, 'retry release') + + expect(recovered.state).toBe('released') + expect(releases).toBe(1) + expect(recoveries).toBe(1) + }) + + it('requires a recovery callback when claim can reserve external funds', () => { + expect(() => new MemoryPaymentOperations({ + onClaim: async () => undefined, + })).toThrow('onReclaim is required') + }) + it('runs one settlement side effect for concurrent retries', async () => { let effects = 0 const operations = new MemoryPaymentOperations({ diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index 395f40d..d830bb6 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -138,6 +138,109 @@ describe('final payment boundary protocol guards', () => { expect(claims).toBe(0) }) + it('rejects a continuation before it can mutate another caller task', async () => { + const taskStore = new InMemoryTaskStore() + const paused: Task = { + kind: 'task', + id: 'victim-task', + contextId: 'victim-context', + status: { state: 'input-required', timestamp: new Date().toISOString() }, + history: [], + } + await taskStore.put(paused) + let sandboxRuns = 0 + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + sandboxRuns += 1 + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'work' } } + }, + }), + recordUsage: async () => undefined, + settlePayment: async () => undefined, + verifyApiKey: async () => ({ consumerId: 'intruder', keyId: 'intruder', scopes: ['chat'] }), + x402: { operatorAddress, chainId: 1, demoMode: true }, + a2a: { + taskStore, + authorizeTaskAccess: async (_task, context) => context.authorization === 'Bearer owner', + }, + })) + + const response = await app.request('/v1/agents/guards', { + method: 'POST', + headers: { 'Content-Type': 'application/json', Authorization: 'Bearer intruder' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: paused.id, + parts: [{ kind: 'text', text: 'continue as attacker' }], + }, + }, + }), + }) + + expect(response.status).toBe(403) + const body = await response.json() as { error?: { code: number } } + expect(body.error?.code).toBe(A2A_ERROR_CODES.TASK_ACCESS_DENIED) + expect(sandboxRuns).toBe(0) + expect(await taskStore.get(paused.id)).toEqual(paused) + }) + + it('keeps a paused task retryable when continuation payment fails', async () => { + const taskStore = new InMemoryTaskStore() + const paused: Task = { + kind: 'task', + id: 'paused-payment-task', + contextId: 'paused-context', + status: { state: 'input-required', timestamp: new Date().toISOString() }, + history: [], + } + await taskStore.put(paused) + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => box(), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => false, + }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + })) + + const response = await app.request('/v1/agents/guards', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('31') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: paused.id, + parts: [{ kind: 'text', text: 'retry' }], + }, + }, + }), + }) + + const body = await response.json() as { error?: { code: number } } + expect(body.error?.code).toBe(A2A_ERROR_CODES.INTERNAL_ERROR) + expect(await taskStore.get(paused.id)).toEqual(paused) + }) + it('shares one canonical nonce authority across mixed verifier versions', async () => { const store = new MemoryNonceStore() const config = { operatorAddress, chainId: 1, demoMode: true } diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 5742673..b681d54 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, beforeEach } from 'vitest' import { verifyX402, verifyMpp, defaultVerifyApiKey } from '../src/verify' -import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryNonceStore, type NonceStore } from '../src/nonce-store' import type { X402Config, MppConfig } from '../src/types' const operatorAddress = '0x1111111111111111111111111111111111111111' @@ -80,6 +80,18 @@ describe('verifyX402', () => { expect(second).toBeNull() }) + it('keeps version 1 custom nonce stores source-compatible', async () => { + const seen = new Set() + const nonceStore: NonceStore = { + hasSeen: async (nonce) => seen.has(nonce), + markSeen: async (nonce) => { seen.add(nonce) }, + } + const payload = buildSpendAuth({ nonce: '101' }) + + expect(await verifyX402(payload, baseConfig, nonceStore)).toBe('0xCommitmentAlice') + expect(await verifyX402(payload, baseConfig, nonceStore)).toBeNull() + }) + it('isolates nonces per commitment — regression: commitment-less nonce tracking lets Alice replay Bob\'s nonce', async () => { const nonceStore = new MemoryNonceStore() const aliceNonce = buildSpendAuth({ commitment: '0xAlice', nonce: '1' }) From 37faa4ea1cdf8c246693a6725f644732ed57a576 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 23:43:40 -0600 Subject: [PATCH 05/32] fix(payments): harden cancellation and ownership recovery --- src/a2a/handler.ts | 148 +++++++++++++++++++++++++------- src/dispatch.ts | 57 +++++++++++- src/payment-operations.ts | 3 + src/types.ts | 1 + tests/a2a-payment-races.test.ts | 63 ++++++++++++++ tests/protocol-guards.test.ts | 56 ++++++++++++ 6 files changed, 293 insertions(+), 35 deletions(-) diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 4950cb9..22dfcce 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -146,7 +146,7 @@ export function createA2AHandlers(deps: A2AHandlerDeps) { switch (parsed.method) { case 'message/send': - return handleMessageSend(c, slug, parsed, deps) + return handleMessageSend(c, slug, parsed, deps, cancels) case 'message/stream': return handleMessageStream(c, slug, parsed, deps, cancels) case 'tasks/get': @@ -180,15 +180,39 @@ async function handleMessageSend( slug: string, req: JSONRPCRequest, deps: A2AHandlerDeps, + cancels: CancelRegistry, ): Promise { const guard = await guardMessageRequest(c, slug, req, deps) if (guard instanceof Response) return guard const { authz, task } = guard + const controller = cancels.register(task.id) + try { + return await executeMessageSend(c, req, deps, authz, task, controller.signal) + } finally { + cancels.clear(task.id) + } +} + +async function executeMessageSend( + c: Context, + req: JSONRPCRequest, + deps: A2AHandlerDeps, + authz: AuthorizedRequest, + task: Task, + signal: AbortSignal, +): Promise { const workingTask: Task = task.status.state === 'working' ? task : { ...task, status: { state: 'working', timestamp: nowIso() } } if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { await releaseOwnedPayment(authz, deps, 'A2A task changed before execution started') + if (signal.aborted) { + const canceled = await deps.taskStore.get(task.id) + if (canceled?.status.state === 'canceled') { + await maybeDeliverPush(canceled, deps) + return c.json(ok(req.id, canceled)) + } + } return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) } @@ -203,7 +227,7 @@ async function handleMessageSend( authz.userMessage, authz.consumerId, deps.config, - undefined, + signal, task.id, authz.maxOutputTokens, )) { @@ -249,12 +273,36 @@ async function handleMessageSend( ) } + if (signal.aborted) { + const canceled = await completeCanceledTask( + authz, + workingTask, + responseText, + usage, + workObserved, + deps, + ) + return c.json(ok(req.id, canceled)) + } + // Settle for the work done so far before short-circuiting on input-required. // The user has been charged for the partial response, which is the right // commercial behavior — the sandbox produced tokens. try { if (!usage) throw new Error('sandbox did not provide a usage receipt') if (!await claimTaskFinalization(deps.taskStore, workingTask)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask?.status.state === 'canceled') { + const canceled = await completeCanceledTask( + authz, + currentTask, + responseText, + usage, + workObserved, + deps, + ) + return c.json(ok(req.id, canceled)) + } throw new Error('A2A task changed before payment settlement') } await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) @@ -386,23 +434,37 @@ async function handleMessageStream( // exists; otherwise retain ownership when output or hidden work was // observed because releasing would make paid work free. if (controller.signal.aborted) { - if (usage) { - try { - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) - } catch (settlementError) { - console.error( - `[a2a] canceled task settlement retained for ${authz.requestId}:`, - settlementError instanceof Error ? settlementError.message : String(settlementError), - ) - } - } else { - await releaseOrRetainPayment(authz, deps, 'a2a task canceled', workObserved || usage !== undefined) - } - const canceled = withStatus(task, 'canceled', undefined, [ - responseTextToArtifact(responseText, `${task.id}-artifact-0`), - ]) - try { - await deps.taskStore.put(canceled) + const canceled = await completeCanceledTask( + authz, + workingTask, + responseText, + usage, + workObserved, + deps, + ) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: canceled.status, + final: true, + }) + return + } + + // Settle once for whatever the sandbox produced (full or partial). + if (!usage) throw new Error('sandbox did not provide a usage receipt') + if (!cancels.beginFinalization(task.id) || !await claimTaskFinalization(deps.taskStore, workingTask)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask?.status.state === 'canceled') { + const canceled = await completeCanceledTask( + authz, + currentTask, + responseText, + usage, + workObserved, + deps, + ) send({ kind: 'status-update', taskId: task.id, @@ -410,19 +472,8 @@ async function handleMessageStream( status: canceled.status, final: true, }) - await maybeDeliverPush(canceled, deps) - } catch (taskError) { - console.error( - `[a2a] failed to persist canceled task ${task.id}:`, - taskError instanceof Error ? taskError.message : String(taskError), - ) + return } - return - } - - // Settle once for whatever the sandbox produced (full or partial). - if (!usage) throw new Error('sandbox did not provide a usage receipt') - if (!cancels.beginFinalization(task.id) || !await claimTaskFinalization(deps.taskStore, workingTask)) { await releaseOrRetainPayment( authz, deps, @@ -955,6 +1006,41 @@ async function releaseOrRetainPayment( } } +async function completeCanceledTask( + authz: AuthorizedRequest, + task: Task, + responseText: string, + usage: SandboxUsageReceipt | undefined, + workObserved: boolean, + deps: A2AHandlerDeps, +): Promise { + if (usage) { + try { + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) + } catch (settlementError) { + console.error( + `[a2a] canceled task settlement retained for ${authz.requestId}:`, + settlementError instanceof Error ? settlementError.message : String(settlementError), + ) + } + } else { + await releaseOrRetainPayment(authz, deps, 'a2a task canceled', workObserved) + } + const currentTask = await deps.taskStore.get(task.id) + const canceledBase = currentTask?.status.state === 'canceled' + ? currentTask + : withStatus(task, 'canceled') + const canceled: Task = responseText + ? { + ...canceledBase, + artifacts: [responseTextToArtifact(responseText, `${task.id}-artifact-0`)], + } + : canceledBase + await deps.taskStore.put(canceled) + await maybeDeliverPush(canceled, deps) + return canceled +} + // ── Helpers ─────────────────────────────────────────────────────────────── const FINALIZING_METADATA_KEY = 'gatewayFinalizing' diff --git a/src/dispatch.ts b/src/dispatch.ts index d09ddda..4f65ddc 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -22,6 +22,7 @@ import type { GatewayConfig, PaymentMethod, SandboxExecutionBudget, + SandboxStreamEvent, SandboxUsageReceipt, } from './types' import { @@ -65,6 +66,7 @@ export interface AuthorizedRequest { paymentPayload: Record | null paymentNonceKey?: string paymentOperation?: PaymentOperation + paymentOperationAcquired?: boolean } function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { @@ -543,12 +545,18 @@ export async function claimPayment( if (operation && !config.x402.paymentOperations) { throw new Error('durable payment operations are required to settle a claimed operation') } + if (operation) { + authz.paymentOperationAcquired = operation.acquiredByRequestId === context.requestId + if (!authz.paymentOperationAcquired) { + throw new Error('payment operation was already claimed') + } + } if (operation && authz.paymentNonceKey) { const claimed = await claimPaymentNonce( state.nonceStore, authz.paymentNonceKey, authz.paymentPayload, - operation.operationId, + `${operation.operationId}:${context.requestId}`, ) if (!claimed) { try { @@ -596,7 +604,7 @@ export async function releasePayment( config: GatewayConfig, reason: string, ): Promise { - if (!authz.paymentOperation || !config.x402.paymentOperations) return + if (!authz.paymentOperation || authz.paymentOperationAcquired !== true || !config.x402.paymentOperations) return await config.x402.paymentOperations.releasePayment(authz.paymentOperation, reason) } @@ -735,9 +743,17 @@ export async function* dispatchSandboxStreamRich( systemPrompt: agent.systemPrompt, maxOutputTokens: outputLimit, executionBudget, + signal, }) - for await (const event of promptStream) { - if (signal?.aborted) return + const iterator = promptStream[Symbol.asyncIterator]() + while (true) { + const next = await readSandboxEvent(iterator, signal) + if (next === ABORTED_SANDBOX_READ) { + closeSandboxIterator(iterator) + return + } + if (next.done) break + const event = next.value if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) if (event.data?.reasoning?.tokens !== undefined) { observedReasoningTokens += nonNegativeSafeInteger(event.data.reasoning.tokens, 'reasoning tokens') @@ -805,6 +821,39 @@ export async function* dispatchSandboxStreamRich( yield { kind: 'usage', usage } } +const ABORTED_SANDBOX_READ = Symbol('aborted-sandbox-read') + +async function readSandboxEvent( + iterator: AsyncIterator, + signal?: AbortSignal, +): Promise | typeof ABORTED_SANDBOX_READ> { + if (!signal) return iterator.next() + if (signal.aborted) return ABORTED_SANDBOX_READ + return new Promise((resolve, reject) => { + const onAbort = () => resolve(ABORTED_SANDBOX_READ) + signal.addEventListener('abort', onAbort, { once: true }) + iterator.next().then( + (result) => { + signal.removeEventListener('abort', onAbort) + resolve(result) + }, + (error: unknown) => { + signal.removeEventListener('abort', onAbort) + reject(error) + }, + ) + }) +} + +function closeSandboxIterator(iterator: AsyncIterator): void { + try { + const closing = iterator.return?.() + if (closing) void Promise.resolve(closing).catch(() => undefined) + } catch { + // The request is already canceled. Adapter cleanup must not delay it. + } +} + /** * Record usage event + settle payment + invoke the observer. Both wire * formats call this once their stream has drained, so settlement happens diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 993f6dd..7add441 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -17,6 +17,8 @@ export type PaymentOperationState = export interface PaymentOperation { protocolVersion: typeof PAYMENT_PROTOCOL_VERSION operationId: string + /** Request that atomically created this operation. Idempotent reads retain the original value. */ + acquiredByRequestId: string nonceKey: string authorizationId: string reservedAmount: bigint @@ -112,6 +114,7 @@ export class MemoryPaymentOperations implements PaymentOperations { const operation: PaymentOperation = { protocolVersion: PAYMENT_PROTOCOL_VERSION, operationId, + acquiredByRequestId: context.requestId, nonceKey, authorizationId: typeof payload.authHash === 'string' ? payload.authHash : operationId, reservedAmount, diff --git a/src/types.ts b/src/types.ts index ce7e7aa..b5a4dc7 100644 --- a/src/types.ts +++ b/src/types.ts @@ -245,6 +245,7 @@ export interface SandboxBox { systemPrompt?: string maxOutputTokens?: number executionBudget?: SandboxExecutionBudget + signal?: AbortSignal }, ): AsyncIterable } diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 73187ae..b959820 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -183,6 +183,69 @@ describe('A2A payment ownership races', () => { expect(finalTask?.history).toHaveLength(2) }) + it('returns a canceled synchronous task without losing payment ownership', async () => { + const taskStore = new InMemoryTaskStore() + const operations = new MemoryPaymentOperations() + let outputSeen!: () => void + const outputReady = new Promise((resolve) => { outputSeen = resolve }) + let releaseSandbox!: () => void + const sandboxReleased = new Promise((resolve) => { releaseSandbox = resolve }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + outputSeen() + await sandboxReleased + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('75') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-sync-cancel') }, + }), + }) + await outputReady + + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-sync-cancel' }, + }), + }) + expect(cancel.status).toBe(200) + releaseSandbox() + const response = await send + const body = await response.json() as { result?: { status?: { state?: string } }; error?: unknown } + + expect(body.error).toBeUndefined() + expect(body.result?.status?.state).toBe('canceled') + expect((await taskStore.get('task-sync-cancel'))?.status.state).toBe('canceled') + expect(operations.get(`x402:${commitment}:75`)?.state).toBe('claimed') + }) + it('retains ownership when cancellation follows delivered output without a receipt', async () => { const taskStore = new InMemoryTaskStore() const operations = new MemoryPaymentOperations() diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index d830bb6..a22c5c1 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -312,4 +312,60 @@ describe('final payment boundary protocol guards', () => { expect(legacyResponse.status).toBe(200) expect(operations.get(`x402:${commitment}:5`)).toBeUndefined() }) + + it('lets only one separately configured version 2 gateway own a payment', async () => { + const nonceStore = new MemoryNonceStore() + let claimsReady = 0 + let releaseClaims!: () => void + const claimsReleased = new Promise((resolve) => { releaseClaims = resolve }) + let runs = 0 + let settlements = 0 + const stores = [ + new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 } }), + new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 } }), + ] + const apps = stores.map((operations) => { + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield* box().streamPrompt('run') + }, + }), + recordUsage: async () => undefined, + nonceStore, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + authorizePayment: async (payload, context) => { + const operation = await operations.claimPayment(payload, context) + claimsReady += 1 + if (claimsReady === 2) releaseClaims() + await claimsReleased + return operation + }, + }, + })) + return app + }) + + const responses = await Promise.all(apps.map((app) => app.request('/v1/agents/guards/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('32') }, + body: JSON.stringify({ max_tokens: 1, messages: [{ role: 'user', content: 'run once' }] }), + }))) + await Promise.all(responses.map((response) => response.text())) + + expect(responses.map((response) => response.status).sort()).toEqual([200, 402]) + expect(runs).toBe(1) + expect(settlements).toBe(1) + expect(stores + .map((store) => store.get(`x402:${commitment}:32`)?.state) + .sort()).toEqual(['released', 'settled']) + }) }) From 5f47d85c70b6709ef530cb8568883829acfa6906 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 23:52:17 -0600 Subject: [PATCH 06/32] fix(release): include observer entry and settle cancel races --- src/a2a/handler.ts | 6 +++++- tsup.config.ts | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 22dfcce..cab70a4 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -454,7 +454,10 @@ async function handleMessageStream( // Settle once for whatever the sandbox produced (full or partial). if (!usage) throw new Error('sandbox did not provide a usage receipt') - if (!cancels.beginFinalization(task.id) || !await claimTaskFinalization(deps.taskStore, workingTask)) { + // Let the durable task-store CAS decide the cancellation race. Mark + // the local registry only after that CAS wins, so cancel can replace a + // still-pending finalization instead of being rejected prematurely. + if (!await claimTaskFinalization(deps.taskStore, workingTask)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask?.status.state === 'canceled') { const canceled = await completeCanceledTask( @@ -482,6 +485,7 @@ async function handleMessageStream( ) return } + cancels.beginFinalization(task.id) await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) if (inputRequiredSeen) { diff --git a/tsup.config.ts b/tsup.config.ts index a4a5072..e996db4 100644 --- a/tsup.config.ts +++ b/tsup.config.ts @@ -9,6 +9,7 @@ export default defineConfig({ 'src/publish.ts', 'src/rate-limit.ts', 'src/nonce-store.ts', + 'src/observer.ts', ], format: ['esm'], dts: true, From 2335d1532baadf7ee4dd713d3d3e3460f308034c Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 23:53:21 -0600 Subject: [PATCH 07/32] test(payments): cover final ownership races --- tests/a2a-payment-races.test.ts | 235 +++++++++++++++++++++++++++++++- tests/protocol-guards.test.ts | 68 +++++++++ 2 files changed, 302 insertions(+), 1 deletion(-) diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index b959820..c6c4c0d 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -1,7 +1,7 @@ import { Hono } from 'hono' import { describe, expect, it } from 'vitest' -import { InMemoryTaskStore } from '../src/a2a/task-store' +import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { createAgentGateway } from '../src/middleware' import { MemoryNonceStore } from '../src/nonce-store' import { MemoryPaymentOperations } from '../src/payment-operations' @@ -246,6 +246,239 @@ describe('A2A payment ownership races', () => { expect(operations.get(`x402:${commitment}:75`)?.state).toBe('claimed') }) + it('settles when cancellation wins the final task update', async () => { + const innerStore = new InMemoryTaskStore() + let finalizationSeen!: () => void + const finalizationReady = new Promise((resolve) => { finalizationSeen = resolve }) + let releaseFinalization!: () => void + const finalizationReleased = new Promise((resolve) => { releaseFinalization = resolve }) + const taskStore: TaskStore = { + get: (id) => innerStore.get(id), + put: (task) => innerStore.put(task), + createIfAbsent: (task) => innerStore.createIfAbsent(task), + delete: (id) => innerStore.delete(id), + async compareAndSet(expected, next) { + if (next.metadata?.gatewayFinalizing === true) { + finalizationSeen() + await finalizationReleased + } + return innerStore.compareAndSet(expected, next) + }, + } + let settlements = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { settlements += 1 }, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('78') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-finalization-cancel') }, + }), + }) + await finalizationReady + + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-finalization-cancel' }, + }), + }) + expect(cancel.status).toBe(200) + releaseFinalization() + const response = await send + const body = await response.json() as { result?: { status?: { state?: string } }; error?: unknown } + + expect(body.error).toBeUndefined() + expect(body.result?.status?.state).toBe('canceled') + expect(settlements).toBe(1) + expect(operations.get(`x402:${commitment}:78`)?.state).toBe('settled') + }) + + it('settles a stream when cancellation wins the final task update', async () => { + const innerStore = new InMemoryTaskStore() + let cancellationStored!: () => void + const cancellationReady = new Promise((resolve) => { cancellationStored = resolve }) + let releaseCancellation!: () => void + const cancellationReleased = new Promise((resolve) => { releaseCancellation = resolve }) + const taskStore: TaskStore = { + get: (id) => innerStore.get(id), + put: (task) => innerStore.put(task), + createIfAbsent: (task) => innerStore.createIfAbsent(task), + delete: (id) => innerStore.delete(id), + async compareAndSet(expected, next) { + const transitioned = await innerStore.compareAndSet(expected, next) + if (transitioned && next.status.state === 'canceled') { + cancellationStored() + await cancellationReleased + } + return transitioned + }, + } + let sandboxDrained!: () => void + const sandboxReady = new Promise((resolve) => { sandboxDrained = resolve }) + let releaseSandbox!: () => void + const sandboxReleased = new Promise((resolve) => { releaseSandbox = resolve }) + let settlements = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { settlements += 1 }, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + sandboxDrained() + await sandboxReleased + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const stream = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('80') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/stream', + params: { message: message('run', 'task-stream-finalization-cancel') }, + }), + }) + const reader = stream.body!.getReader() + await sandboxReady + + const cancelPromise = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-stream-finalization-cancel' }, + }), + }) + await cancellationReady + releaseSandbox() + let wire = '' + const decoder = new TextDecoder() + while (true) { + const chunk = await reader.read() + if (chunk.done) break + wire += decoder.decode(chunk.value, { stream: true }) + } + releaseCancellation() + const cancel = await cancelPromise + + expect(cancel.status).toBe(200) + expect(wire).toContain('"state":"canceled"') + expect(settlements).toBe(1) + expect(operations.get(`x402:${commitment}:80`)?.state).toBe('settled') + }) + + it('returns after cancel when a sandbox stops emitting events', async () => { + const taskStore = new InMemoryTaskStore() + const operations = new MemoryPaymentOperations() + let outputSeen!: () => void + const outputReady = new Promise((resolve) => { outputSeen = resolve }) + const never = new Promise(() => undefined) + let sandboxSignal: AbortSignal | undefined + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt(_message, opts) { + sandboxSignal = opts?.signal + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + outputSeen() + await never + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('79') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-silent-cancel') }, + }), + }) + await outputReady + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-silent-cancel' }, + }), + }) + expect(cancel.status).toBe(200) + const response = await Promise.race([ + send, + new Promise((_resolve, reject) => setTimeout(() => reject(new Error('send did not cancel')), 100)), + ]) + const body = await response.json() as { result?: { status?: { state?: string } }; error?: unknown } + + expect(body.error).toBeUndefined() + expect(body.result?.status?.state).toBe('canceled') + expect(sandboxSignal?.aborted).toBe(true) + expect(operations.get(`x402:${commitment}:79`)?.state).toBe('claimed') + }) + it('retains ownership when cancellation follows delivered output without a receipt', async () => { const taskStore = new InMemoryTaskStore() const operations = new MemoryPaymentOperations() diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index a22c5c1..2af416d 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -368,4 +368,72 @@ describe('final payment boundary protocol guards', () => { .map((store) => store.get(`x402:${commitment}:32`)?.state) .sort()).toEqual(['released', 'settled']) }) + + it('does not let an idempotent retry release the payment owner', async () => { + const nonceStore = new MemoryNonceStore() + let settlements = 0 + let releases = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { settlements += 1 }, + onRelease: async () => { releases += 1 }, + }) + let operationPromise: ReturnType | undefined + let authorizations = 0 + let releaseAuthorizations!: () => void + const authorizationsReleased = new Promise((resolve) => { releaseAuthorizations = resolve }) + let runStarted!: () => void + const sandboxStarted = new Promise((resolve) => { runStarted = resolve }) + let finishRun!: () => void + const runReleased = new Promise((resolve) => { finishRun = resolve }) + let runs = 0 + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + runStarted() + await runReleased + yield* box().streamPrompt('run') + }, + }), + recordUsage: async () => undefined, + nonceStore, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + authorizePayment: async (payload, context) => { + operationPromise ??= operations.claimPayment(payload, context) + const operation = await operationPromise + authorizations += 1 + if (authorizations === 2) releaseAuthorizations() + await authorizationsReleased + return operation + }, + }, + })) + const request = () => app.request('/v1/agents/guards/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('33') }, + body: JSON.stringify({ max_tokens: 1, messages: [{ role: 'user', content: 'run once' }] }), + }) + const requests = [request(), request()] + await sandboxStarted + const loser = await Promise.race(requests) + expect(loser.status).toBe(402) + expect(operations.get(`x402:${commitment}:33`)?.state).toBe('claimed') + expect(releases).toBe(0) + + finishRun() + const responses = await Promise.all(requests) + await Promise.all(responses.map((response) => response.text())) + expect(responses.map((response) => response.status).sort()).toEqual([200, 402]) + expect(runs).toBe(1) + expect(settlements).toBe(1) + expect(releases).toBe(0) + expect(operations.get(`x402:${commitment}:33`)?.state).toBe('settled') + }) }) From 1bef5f0441785e04ed386792fef8343d8989b821 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Thu, 13 Aug 2026 23:59:45 -0600 Subject: [PATCH 08/32] fix(api): type payment protocol marker --- src/middleware.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/middleware.ts b/src/middleware.ts index 5262c78..a453ec6 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -251,6 +251,11 @@ export function createAgentGateway(config: GatewayConfig) { // when an old binary ignores the version 2 operation contract. Object.assign(createAgentGateway, { paymentProtocolVersion: 2 as const }) +/** Public package-boundary version marker for durable payment operations. */ +export namespace createAgentGateway { + export const paymentProtocolVersion = 2 as const +} + /** * Drain the sandbox stream into an OpenAI-shaped SSE response, settle the * payment, fire observer hooks. Identical pre-refactor behavior, just lifted From 2492f1e1c0016d6948e397fb9787f9df59934133 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 00:06:39 -0600 Subject: [PATCH 09/32] fix(payments): retain unreceipted execution --- README.md | 4 ++- src/a2a/handler.ts | 9 ++++++ src/dispatch.ts | 28 ++++++++++++++---- src/middleware.ts | 5 ++++ src/payment-operations.ts | 49 ++++++++++++++++++++++++++++++-- tests/a2a-payment-races.test.ts | 6 ++-- tests/payment-operations.test.ts | 24 +++++++++++++++- tests/protocol-guards.test.ts | 7 ++--- 8 files changed, 116 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 5cdbad7..804ff0c 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,9 @@ Set `x402.demoMode: true` only for local development and tests; that explicit mo Keep `verifySigner` free of side effects. Use `authorizePayment` to reserve or claim funds after rate limits, content checks, and product authorization succeed. For production version 2, set `x402.paymentProtocolVersion: 2`, provide `paymentOperations`, and return its operation from `authorizePayment`. -The operation store owns claim, partial settle, release, and expiry reclaim. +The operation store owns claim, execution start, receipt retention, partial settle, release, and expiry reclaim. +An executing or retained operation cannot expire into a refund. +A retained operation can settle later when recovery obtains its usage receipt. Keep version 1 explicitly configured while old and new gateways coexist; shared nonce storage must reject a version 1 claim owned by a version 2 operation. Before it calls the verifier, the gateway requires the signed amount to cover filtered input plus the requested output limit. The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index cab70a4..4781644 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -17,6 +17,7 @@ import { type AuthorizedRequest, type GatewayState, authenticateAndGuard, + beginPaymentExecution, claimPayment, dispatchSandboxStreamRich, releasePayment, @@ -230,6 +231,10 @@ async function executeMessageSend( signal, task.id, authz.maxOutputTokens, + async () => { + await beginPaymentExecution(authz, deps.config) + if (authz.paymentOperation) workObserved = true + }, )) { if (event.kind === 'text') { responseText += event.delta @@ -403,6 +408,10 @@ async function handleMessageStream( controller.signal, task.id, authz.maxOutputTokens, + async () => { + await beginPaymentExecution(authz, deps.config) + if (authz.paymentOperation) workObserved = true + }, )) { if (event.kind === 'text') { responseText += event.delta diff --git a/src/dispatch.ts b/src/dispatch.ts index 4f65ddc..d68533d 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -13,7 +13,7 @@ import type { Context } from 'hono' import { filterConsumerMessagesStrict, redactSystemPromptFromOutput } from './filter' import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' import { type RateLimitStore, checkRateLimit } from './rate-limit' -import { claimStoredNonce, type NonceStore } from './nonce-store' +import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' import type { PaymentOperation } from './payment-operations' import type { AgentMeta, @@ -608,6 +608,15 @@ export async function releasePayment( await config.x402.paymentOperations.releasePayment(authz.paymentOperation, reason) } +/** Mark a durable reservation active immediately before sandbox execution. */ +export async function beginPaymentExecution( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + if (!authz.paymentOperation || authz.paymentOperationAcquired !== true || !config.x402.paymentOperations) return + authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) +} + /** * Release only when no sandbox work was observed. Once output or a receipt * exists, retain the owner for settlement or background recovery. @@ -619,6 +628,9 @@ export async function releasePaymentAfterFailure( workObserved: boolean, ): Promise { if (workObserved) { + if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { + authz.paymentOperation = await config.x402.paymentOperations.retainPayment(authz.paymentOperation, reason) + } console.error( `[agent-gateway] retaining payment ownership after sandbox work for ${authz.requestId}: ${reason}`, ) @@ -641,10 +653,12 @@ async function claimPaymentNonce( payload: Record, ownerId?: string, ): Promise { - const expiry = Number(payload.expiry ?? Math.floor(Date.now() / 1000) + 3600) - const ttl = Math.min(expiry - Math.floor(Date.now() / 1000), 3600) - if (!Number.isFinite(ttl) || ttl <= 0) return false - return claimStoredNonce(nonceStore, nonceKey, Math.max(ttl, 60), ownerId) + const expiry = payload.expiry === undefined + ? BigInt(Math.floor(Date.now() / 1000) + 3600) + : BigInt(String(payload.expiry)) + const ttl = nonceTtlSeconds(expiry) + if (ttl === undefined) return false + return claimStoredNonce(nonceStore, nonceKey, ttl, ownerId) } /** @@ -708,6 +722,7 @@ export async function* dispatchSandboxStreamRich( signal?: AbortSignal, sessionId?: string, maxOutputTokens?: number, + onExecutionStart?: () => Promise, ): AsyncIterable { const box = await config.getSandbox(agent) const outputLimit = maxOutputTokens ?? config.defaultOutputTokens ?? 1024 @@ -738,6 +753,9 @@ export async function* dispatchSandboxStreamRich( (config.executionBudget?.maxReasoningTokens ?? outputLimit) + (config.executionBudget?.maxToolTokens ?? outputLimit)) * agent.pricePerTokenUsd, } + if (signal?.aborted) return + await onExecutionStart?.() + if (signal?.aborted) return const promptStream = box.streamPrompt(userMessage, { sessionId: sessionId ?? `consumer:${consumerId}`, systemPrompt: agent.systemPrompt, diff --git a/src/middleware.ts b/src/middleware.ts index a453ec6..dc743c0 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -6,6 +6,7 @@ import { type AuthorizedRequest, type GatewayState, authenticateAndGuard, + beginPaymentExecution, claimPayment, dispatchSandboxStreamRich, releasePayment, @@ -310,6 +311,10 @@ function streamChatCompletions( undefined, undefined, maxOutputTokens, + async () => { + await beginPaymentExecution(authz, config) + if (authz.paymentOperation) workObserved = true + }, )) { if (event.kind === 'text') { sendChunk(event.delta) diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 7add441..97012c9 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -6,6 +6,8 @@ export const PAYMENT_PROTOCOL_VERSION = 2 as const export type PaymentOperationState = | 'claiming' | 'claimed' + | 'executing' + | 'retained' | 'settling' | 'settled' | 'releasing' @@ -19,6 +21,8 @@ export interface PaymentOperation { operationId: string /** Request that atomically created this operation. Idempotent reads retain the original value. */ acquiredByRequestId: string + executionStartedAt?: number + retentionReason?: string nonceKey: string authorizationId: string reservedAmount: bigint @@ -59,6 +63,10 @@ export interface PaymentOperations { payload: Record, context: PaymentAuthorizationContext, ): Promise + /** Prevent expiry reclaim while the sandbox can consume provider resources. */ + beginPaymentExecution(operation: PaymentOperation): Promise + /** Preserve funds after work when the final usage receipt is still missing. */ + retainPayment(operation: PaymentOperation, reason: string): Promise settlePayment( operation: PaymentOperation, input: PaymentSettlementInput, @@ -163,7 +171,7 @@ export class MemoryPaymentOperations implements PaymentOperations { this.settleFlights.delete(current.operationId) } } - } else if (current.state !== 'claimed') { + } else if (current.state !== 'claimed' && current.state !== 'executing' && current.state !== 'retained') { throw new Error(`cannot settle payment in state ${current.state}`) } validateSettlement(current, input) @@ -185,6 +193,38 @@ export class MemoryPaymentOperations implements PaymentOperations { } } + async beginPaymentExecution(operation: PaymentOperation): Promise { + const current = this.requireCurrent(operation) + if (current.state === 'executing') return current + if (current.state !== 'claimed') { + throw new Error(`cannot begin payment execution in state ${current.state}`) + } + const executing = { + ...current, + state: 'executing' as const, + executionStartedAt: this.now(), + } + this.operations.set(current.operationId, executing) + return executing + } + + async retainPayment(operation: PaymentOperation, reason: string): Promise { + const current = this.requireCurrent(operation) + if (current.state === 'retained' || current.state === 'settling' || current.state === 'settled') { + return current + } + if (current.state !== 'claimed' && current.state !== 'executing') { + throw new Error(`cannot retain payment in state ${current.state}`) + } + const retained = { + ...current, + state: 'retained' as const, + retentionReason: reason, + } + this.operations.set(current.operationId, retained) + return retained + } + async releasePayment(operation: PaymentOperation, reason: string): Promise { const current = this.requireCurrent(operation) if (current.state === 'released') return current @@ -193,7 +233,9 @@ export class MemoryPaymentOperations implements PaymentOperations { if (flight) return flight return this.reclaimPayment(current.operationId) } - if (current.state !== 'claimed') throw new Error(`cannot release payment in state ${current.state}`) + if (current.state !== 'claimed' && current.state !== 'executing') { + throw new Error(`cannot release payment in state ${current.state}`) + } const releasing = { ...current, state: 'releasing' as const } this.operations.set(current.operationId, releasing) return this.runRelease(releasing, reason) @@ -203,6 +245,9 @@ export class MemoryPaymentOperations implements PaymentOperations { const current = this.operations.get(operationId) if (!current) throw new Error('payment operation was not found') if (current.state === 'reclaimed') return current + if (current.state === 'executing' || current.state === 'retained') { + throw new Error(`cannot reclaim payment in state ${current.state} without a usage receipt`) + } if (current.state === 'settling') { const flight = this.settleFlights.get(operationId) if (flight) return flight diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index c6c4c0d..2fafbb7 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -243,7 +243,7 @@ describe('A2A payment ownership races', () => { expect(body.error).toBeUndefined() expect(body.result?.status?.state).toBe('canceled') expect((await taskStore.get('task-sync-cancel'))?.status.state).toBe('canceled') - expect(operations.get(`x402:${commitment}:75`)?.state).toBe('claimed') + expect(operations.get(`x402:${commitment}:75`)?.state).toBe('retained') }) it('settles when cancellation wins the final task update', async () => { @@ -476,7 +476,7 @@ describe('A2A payment ownership races', () => { expect(body.error).toBeUndefined() expect(body.result?.status?.state).toBe('canceled') expect(sandboxSignal?.aborted).toBe(true) - expect(operations.get(`x402:${commitment}:79`)?.state).toBe('claimed') + expect(operations.get(`x402:${commitment}:79`)?.state).toBe('retained') }) it('retains ownership when cancellation follows delivered output without a receipt', async () => { @@ -542,7 +542,7 @@ describe('A2A payment ownership races', () => { // Drain the stream so its recovery path runs. } - expect(operations.get(`x402:${commitment}:77`)?.state).toBe('claimed') + expect(operations.get(`x402:${commitment}:77`)?.state).toBe('retained') const canceled = await taskStore.get('task-cancel') expect(canceled?.status.state).toBe('canceled') }) diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index 83d9e95..df6f467 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -107,6 +107,28 @@ describe('version 2 payment operations', () => { expect(recoveredCalls).toBe(1) }) + it('does not refund active or retained work after authorization expiry', async () => { + let now = 100 + const operations = new MemoryPaymentOperations({ now: () => now }) + const owner = await operations.claimPayment(payload('1000', '22', '200'), context()) + const executing = await operations.beginPaymentExecution(owner) + expect(executing.state).toBe('executing') + now = 201 + + await expect(operations.reclaimPayment(owner.operationId)).rejects.toThrow('executing') + const retained = await operations.retainPayment(executing, 'usage receipt pending') + expect(retained.state).toBe('retained') + expect(retained.retentionReason).toBe('usage receipt pending') + await expect(operations.reclaimPayment(owner.operationId)).rejects.toThrow('retained') + + const settled = await operations.settlePayment(retained, { + amount: 200n, + totalCostUsd: 0.2, + usage: settledUsage(), + }) + expect(settled.state).toBe('settled') + }) + it('recovers a release after a worker crash between state and side effect', async () => { let fail = true const operations = new MemoryPaymentOperations({ @@ -461,7 +483,7 @@ describe('bounded request pricing and sandbox receipts', () => { expect(wire).toContain('sandbox exceeded max output tokens') expect(wire).not.toContain('0123456789abcdefghijklmnop') expect(releases).toBe(0) - expect(operations.get(`x402:${'0x' + 'ab'.repeat(32)}:15`)?.state).toBe('claimed') + expect(operations.get(`x402:${'0x' + 'ab'.repeat(32)}:15`)?.state).toBe('retained') }) it('retains hidden usage when a final provider receipt omits it', async () => { diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index 2af416d..c7ddff6 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -422,13 +422,12 @@ describe('final payment boundary protocol guards', () => { }) const requests = [request(), request()] await sandboxStarted - const loser = await Promise.race(requests) - expect(loser.status).toBe(402) - expect(operations.get(`x402:${commitment}:33`)?.state).toBe('claimed') + const responses = await Promise.all(requests) + expect(responses.map((response) => response.status).sort()).toEqual([200, 402]) + expect(operations.get(`x402:${commitment}:33`)?.state).toBe('executing') expect(releases).toBe(0) finishRun() - const responses = await Promise.all(requests) await Promise.all(responses.map((response) => response.text())) expect(responses.map((response) => response.status).sort()).toEqual([200, 402]) expect(runs).toBe(1) From e4bd05c6dd3c8c550fc5cdd98de61c0814505fed Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 00:12:28 -0600 Subject: [PATCH 10/32] fix(payments): retain nonce through signed expiry --- src/nonce-store.ts | 15 +++++++++++++++ src/verify.ts | 17 +++++++++-------- tests/nonce-store.test.ts | 11 ++++++++++- tests/verify.test.ts | 33 +++++++++++++++++++++++++++++++++ 4 files changed, 67 insertions(+), 9 deletions(-) diff --git a/src/nonce-store.ts b/src/nonce-store.ts index 83d82dd..c21809e 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -16,6 +16,21 @@ export interface NonceStore { markSeen(nonce: string, ttlSeconds: number): Promise } +/** + * Return the seconds for which a signed nonce must remain stored. + * + * The signed expiry is the replay boundary. A fixed one-hour cap would allow + * a still-valid authorization to replay after the nonce entry expires. + */ +export function nonceTtlSeconds( + expiry: bigint, + nowSeconds = Math.floor(Date.now() / 1000), +): number | undefined { + const remaining = expiry - BigInt(nowSeconds) + if (remaining <= 0n || remaining > BigInt(Number.MAX_SAFE_INTEGER)) return undefined + return Math.max(Number(remaining), 60) +} + // --------------------------------------------------------------------------- // In-memory implementation — single-worker, ephemeral // --------------------------------------------------------------------------- diff --git a/src/verify.ts b/src/verify.ts index 969c02e..63002b4 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -1,5 +1,5 @@ import type { X402Config, MppConfig, ApiKeyInfo, GatewayConfig } from './types' -import { claimStoredNonce, type NonceStore } from './nonce-store' +import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' /** Return the canonical opaque nonce key used by the final payment claim. */ export function mppReplayNonceKey(authHeader: string): string | undefined { @@ -101,9 +101,9 @@ export async function verifyX402( // Check and mark only after the signature is accepted. Otherwise an // invalid request can burn a valid payer nonce and deny the real request. if (nonceStore && markNonce) { - // Mark seen with TTL matching the expiry window (max 1 hour) - const ttl = Math.min(Number(expiry) - Math.floor(Date.now() / 1000), 3600) - const claimed = await claimStoredNonce(nonceStore, nonceKey, Math.max(ttl, 60)) + const ttl = nonceTtlSeconds(expiry) + if (ttl === undefined) return null + const claimed = await claimStoredNonce(nonceStore, nonceKey, ttl) if (!claimed) return null } @@ -199,10 +199,11 @@ export async function verifyMpp( if (nonceStore && payload.nonce !== undefined && markNonce) { const expiry = payload.expiry === undefined - ? Math.floor(Date.now() / 1000) + 3600 - : Number(payload.expiry) - const ttl = Math.min(expiry - Math.floor(Date.now() / 1000), 3600) - const claimed = await claimStoredNonce(nonceStore, nonceKey!, Math.max(ttl, 60)) + ? BigInt(Math.floor(Date.now() / 1000) + 3600) + : BigInt(String(payload.expiry)) + const ttl = nonceTtlSeconds(expiry) + if (ttl === undefined) return null + const claimed = await claimStoredNonce(nonceStore, nonceKey!, ttl) if (!claimed) return null } diff --git a/tests/nonce-store.test.ts b/tests/nonce-store.test.ts index 7e8193a..8161666 100644 --- a/tests/nonce-store.test.ts +++ b/tests/nonce-store.test.ts @@ -1,5 +1,14 @@ import { describe, it, expect, vi, afterEach } from 'vitest' -import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryNonceStore, nonceTtlSeconds } from '../src/nonce-store' + +describe('nonceTtlSeconds', () => { + it('covers the complete signed validity window', () => { + expect(nonceTtlSeconds(101n, 100)).toBe(60) + expect(nonceTtlSeconds(7_300n, 100)).toBe(7_200) + expect(nonceTtlSeconds(100n, 100)).toBeUndefined() + expect(nonceTtlSeconds(BigInt(Number.MAX_SAFE_INTEGER) + 101n, 100)).toBeUndefined() + }) +}) describe('MemoryNonceStore', () => { afterEach(() => vi.useRealTimers()) diff --git a/tests/verify.test.ts b/tests/verify.test.ts index b681d54..1a4952c 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -80,6 +80,39 @@ describe('verifyX402', () => { expect(second).toBeNull() }) + it('rejects replay for the full signed lifetime beyond one hour', async () => { + const nonceStore = new MemoryNonceStore() + const payload = buildSpendAuth({ + nonce: '102', + expiry: String(Math.floor(Date.now() / 1000) + 7_200), + }) + + expect(await verifyX402(payload, baseConfig, nonceStore)).toBe('0xCommitmentAlice') + expect(await verifyX402(payload, baseConfig, nonceStore)).toBeNull() + }) + + it('retains a nonce until its signed expiry, beyond the old one-hour cap', async () => { + const ttls: number[] = [] + const nonceStore: NonceStore = { + hasSeen: async () => false, + claim: async (_nonce, ttlSeconds) => { + ttls.push(ttlSeconds) + return true + }, + markSeen: async () => undefined, + } + const now = Math.floor(Date.now() / 1000) + const result = await verifyX402( + buildSpendAuth({ nonce: '102', expiry: String(now + 7200) }), + baseConfig, + nonceStore, + ) + + expect(result).toBe('0xCommitmentAlice') + expect(ttls[0]).toBeGreaterThan(3600) + expect(ttls[0]).toBeLessThanOrEqual(7200) + }) + it('keeps version 1 custom nonce stores source-compatible', async () => { const seen = new Set() const nonceStore: NonceStore = { From a32c1f6b9b4f8e16e8b683caa35c88e381d8b83f Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 00:13:34 -0600 Subject: [PATCH 11/32] fix(gateway): close mixed-auth recovery gaps --- README.md | 4 +- docs/a2a-long-horizon.md | 5 +- src/a2a/handler.ts | 2 + src/dispatch.ts | 161 ++++++++++++++++--------------- src/middleware.ts | 41 ++++---- src/types.ts | 12 ++- tests/middleware.test.ts | 91 ++++++++++++++++- tests/payment-operations.test.ts | 49 ++++++++++ 8 files changed, 262 insertions(+), 103 deletions(-) diff --git a/README.md b/README.md index 804ff0c..5b82ee8 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,9 @@ Before it calls the verifier, the gateway requires the signed amount to cover fi The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. Sandbox adapters should emit a complete `sandbox.usage` receipt. -Version 2 payment operations reject missing receipts; legacy adapters use visible-token estimates only. +Requests with a version 2 payment operation reject missing receipts; API-key and MPP requests keep the legacy visible-token estimate path. +Older custom A2A task stores remain source-compatible through a process-safe fallback. +Use an atomic task store for multi-worker production deployments. MPP is method-specific. Configure `mpp.verifySigner` for production MPP credentials; it receives the decoded JSON payload when available plus the original decoded credential, and returns the authenticated consumer ID or `null`. diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 3a8c112..1e3882d 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -15,8 +15,9 @@ All four are gated on configuration — they cost nothing for agents that don't By default `GatewayConfig.a2a.taskStore` is in-memory: fast, zero-config, fine for tests and single-machine deployments. Production deployments swap in `SqlTaskStore` against any SQL store — D1, postgres, sqlite, libSQL, Turso — via a `SqlAdapter` shim. -Custom production task stores must implement both atomic methods, `createIfAbsent` and `compareAndSet`. -The gateway rejects a custom store without those methods because a read-then-write fallback can run paid work twice across workers. +Custom production task stores should implement both atomic methods, `createIfAbsent` and `compareAndSet`. +Older stores remain compatible through a read-then-write fallback, which is process-safe only and can run paid work twice across workers. +Use `SqlTaskStore` or another atomic adapter for multi-worker production deployments. Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. The hook receives the task and request headers so the application can enforce task ownership. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 4781644..280a17e 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -235,6 +235,7 @@ async function executeMessageSend( await beginPaymentExecution(authz, deps.config) if (authz.paymentOperation) workObserved = true }, + authz.paymentOperation !== undefined, )) { if (event.kind === 'text') { responseText += event.delta @@ -412,6 +413,7 @@ async function handleMessageStream( await beginPaymentExecution(authz, deps.config) if (authz.paymentOperation) workObserved = true }, + authz.paymentOperation !== undefined, )) { if (event.kind === 'text') { responseText += event.delta diff --git a/src/dispatch.ts b/src/dispatch.ts index d68533d..4c1f52e 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -723,6 +723,7 @@ export async function* dispatchSandboxStreamRich( sessionId?: string, maxOutputTokens?: number, onExecutionStart?: () => Promise, + requiresReceipt = config.x402.paymentOperations !== undefined, ): AsyncIterable { const box = await config.getSandbox(agent) const outputLimit = maxOutputTokens ?? config.defaultOutputTokens ?? 1024 @@ -764,79 +765,80 @@ export async function* dispatchSandboxStreamRich( signal, }) const iterator = promptStream[Symbol.asyncIterator]() - while (true) { - const next = await readSandboxEvent(iterator, signal) - if (next === ABORTED_SANDBOX_READ) { - closeSandboxIterator(iterator) - return - } - if (next.done) break - const event = next.value - if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) - if (event.data?.reasoning?.tokens !== undefined) { - observedReasoningTokens += nonNegativeSafeInteger(event.data.reasoning.tokens, 'reasoning tokens') - yield { kind: 'activity' } - } - if (event.data?.tool) { - observedToolCalls += 1 - observedToolTokens += - nonNegativeSafeInteger(event.data.tool.inputTokens ?? 0, 'tool input tokens') + - nonNegativeSafeInteger(event.data.tool.outputTokens ?? 0, 'tool output tokens') - yield { kind: 'activity' } - } - enforceUsageBudget(withObservedUsage( - usageParts, - observedReasoningTokens, - observedToolTokens, - observedToolCalls, - ), executionBudget) - if ( - event.type === 'message.part.updated' && - event.data?.part?.type === 'text' && - event.data.delta - ) { - const remainingBytes = maxOutputBytes - outputBytes - if (remainingBytes <= 0) throw new Error('sandbox exceeded max output tokens') - const bounded = truncateUtf8(event.data.delta, remainingBytes, encoder) - if (bounded.truncated) { + try { + while (true) { + const next = await readSandboxEvent(iterator, signal) + if (next === ABORTED_SANDBOX_READ) return + if (next.done) break + const event = next.value + if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) + if (event.data?.reasoning?.tokens !== undefined) { + observedReasoningTokens += nonNegativeSafeInteger(event.data.reasoning.tokens, 'reasoning tokens') yield { kind: 'activity' } - throw new Error('sandbox exceeded max output tokens') } - outputBytes += bounded.bytes - legacyOutputText += bounded.text - yield { kind: 'activity' } - yield { kind: 'text', delta: redactSystemPromptFromOutput(bounded.text, agent.systemPrompt) } - continue - } - if (event.type === 'input-required' || event.data?.inputRequired) { - const usage = completeUsage( + if (event.data?.tool) { + observedToolCalls += 1 + observedToolTokens += + nonNegativeSafeInteger(event.data.tool.inputTokens ?? 0, 'tool input tokens') + + nonNegativeSafeInteger(event.data.tool.outputTokens ?? 0, 'tool output tokens') + yield { kind: 'activity' } + } + enforceUsageBudget(withObservedUsage( usageParts, observedReasoningTokens, observedToolTokens, observedToolCalls, - userMessage, - legacyOutputText, - executionBudget, - config.x402.paymentOperations !== undefined, - ) - yield { kind: 'input-required', prompt: event.data?.inputRequired?.prompt } - // Terminal for the sandbox stream — sandbox SHOULD stop emitting until - // the gateway dispatches a continuation message with the new user input. - yield { kind: 'usage', usage } - return + ), executionBudget) + if ( + event.type === 'message.part.updated' && + event.data?.part?.type === 'text' && + event.data.delta + ) { + const remainingBytes = maxOutputBytes - outputBytes + if (remainingBytes <= 0) throw new Error('sandbox exceeded max output tokens') + const bounded = truncateUtf8(event.data.delta, remainingBytes, encoder) + if (bounded.truncated) { + yield { kind: 'activity' } + throw new Error('sandbox exceeded max output tokens') + } + outputBytes += bounded.bytes + legacyOutputText += bounded.text + yield { kind: 'activity' } + yield { kind: 'text', delta: redactSystemPromptFromOutput(bounded.text, agent.systemPrompt) } + continue + } + if (event.type === 'input-required' || event.data?.inputRequired) { + const usage = completeUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + userMessage, + legacyOutputText, + executionBudget, + requiresReceipt, + ) + yield { kind: 'input-required', prompt: event.data?.inputRequired?.prompt } + // Terminal for the sandbox stream — sandbox SHOULD stop emitting until + // the gateway dispatches a continuation message with the new user input. + yield { kind: 'usage', usage } + return + } } + const usage = completeUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + userMessage, + legacyOutputText, + executionBudget, + requiresReceipt, + ) + yield { kind: 'usage', usage } + } finally { + closeSandboxIterator(iterator) } - const usage = completeUsage( - usageParts, - observedReasoningTokens, - observedToolTokens, - observedToolCalls, - userMessage, - legacyOutputText, - executionBudget, - config.x402.paymentOperations !== undefined, - ) - yield { kind: 'usage', usage } } const ABORTED_SANDBOX_READ = Symbol('aborted-sandbox-read') @@ -913,9 +915,6 @@ export async function settleAndRecord( startMs: authz.startMs, } try { - // Record attribution before settlement. Marketplace adapters use this - // row to map the request to its agent and consumer before deducting funds. - await config.recordUsage(usageEvent) if (authz.paymentOperation && config.x402.paymentOperations) { const amount = actualX402Amount( agent.pricePerTokenUsd, @@ -930,15 +929,23 @@ export async function settleAndRecord( authz.paymentOperation, { amount, totalCostUsd: totalCost, usage }, ) - } else if (config.settlePayment) { - await config.settlePayment( - { - method: authz.paymentMethod, - consumerId: authz.consumerId, - requestId: authz.requestId, - }, - totalCost, - ) + // Durable settlement happens first. If attribution storage is + // unavailable, recovery must never refund delivered work. + await config.recordUsage(usageEvent) + } else { + // Legacy adapters retain attribution-before-charge because their + // settlement callback may resolve that usage row. + await config.recordUsage(usageEvent) + if (config.settlePayment) { + await config.settlePayment( + { + method: authz.paymentMethod, + consumerId: authz.consumerId, + requestId: authz.requestId, + }, + totalCost, + ) + } } await obs?.onRequestComplete?.(ctx, usageEvent) } catch (err) { diff --git a/src/middleware.ts b/src/middleware.ts index dc743c0..6d358e0 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -232,14 +232,6 @@ export function createAgentGateway(config: GatewayConfig) { // settleAndRecord, so every security and billing guarantee applies uniformly // regardless of which protocol the caller used. const taskStore = config.a2a?.taskStore ?? new InMemoryTaskStore() - if ( - config.a2a?.taskStore && - (!taskStore.createIfAbsent || !taskStore.compareAndSet) - ) { - throw new Error( - 'createAgentGateway: custom A2A taskStore must implement atomic createIfAbsent and compareAndSet', - ) - } const pushStore = config.a2a?.pushStore const a2a = createA2AHandlers({ config, state, taskStore, pushStore }) gw.get('/:slug/.well-known/agent.json', a2a.handleAgentCard) @@ -281,6 +273,11 @@ function streamChatCompletions( let outputText = '' let usage: import('./types').SandboxUsageReceipt | undefined let workObserved = false + const requestSignal = c.req.raw.signal + const abortController = new AbortController() + const abortFromRequest = () => abortController.abort() + if (requestSignal.aborted) abortFromRequest() + else requestSignal.addEventListener('abort', abortFromRequest, { once: true }) const ctx: RequestContext = { requestId, agentSlug: agent.slug, @@ -291,6 +288,7 @@ function streamChatCompletions( async start(controller) { const encoder = new TextEncoder() const sendChunk = (delta: string) => { + if (controller.desiredSize === null) return outputText += delta const chunk: ChatCompletionChunk = { id: `chatcmpl-${Date.now()}`, @@ -308,13 +306,14 @@ function streamChatCompletions( userMessage, consumerId, config, - undefined, + abortController.signal, undefined, maxOutputTokens, async () => { await beginPaymentExecution(authz, config) if (authz.paymentOperation) workObserved = true }, + authz.paymentOperation !== undefined, )) { if (event.kind === 'text') { sendChunk(event.delta) @@ -341,8 +340,10 @@ function streamChatCompletions( model: agent.slug, choices: [{ index: 0, delta: {}, finish_reason: 'stop' }], } - controller.enqueue(encoder.encode(`data: ${JSON.stringify(done)}\n\n`)) - controller.enqueue(encoder.encode('data: [DONE]\n\n')) + if (controller.desiredSize !== null) { + controller.enqueue(encoder.encode(`data: ${JSON.stringify(done)}\n\n`)) + controller.enqueue(encoder.encode('data: [DONE]\n\n')) + } } catch (err) { const rawMessage = err instanceof Error ? err.message : String(err) // Never expose stack traces / absolute paths from sandbox internals. @@ -366,15 +367,21 @@ function streamChatCompletions( releaseError instanceof Error ? releaseError.message : String(releaseError), ) } - controller.enqueue( - encoder.encode( - `data: ${JSON.stringify({ error: { message: safeMessage, type: 'server_error' } })}\n\n`, - ), - ) + if (!abortController.signal.aborted && controller.desiredSize !== null) { + controller.enqueue( + encoder.encode( + `data: ${JSON.stringify({ error: { message: safeMessage, type: 'server_error' } })}\n\n`, + ), + ) + } } finally { - controller.close() + requestSignal.removeEventListener('abort', abortFromRequest) + if (controller.desiredSize !== null) controller.close() } }, + cancel() { + abortController.abort() + }, }) return new Response(stream, { diff --git a/src/types.ts b/src/types.ts index b5a4dc7..8af750a 100644 --- a/src/types.ts +++ b/src/types.ts @@ -191,10 +191,14 @@ export interface GatewayUsageEvent { paymentMethod: PaymentMethod inputTokens: number outputTokens: number - reasoningTokens: number - toolTokens: number - toolCallCount: number - providerCostUsd: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + reasoningTokens?: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + toolTokens?: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + toolCallCount?: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + providerCostUsd?: number totalCostUsd: number ownerEarnedUsd: number platformFeeUsd: number diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index 18a2e5c..1f51220 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -429,6 +429,41 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(body).not.toContain('data: [DONE]') }) + it('aborts and closes the sandbox iterator when an HTTP reader cancels', async () => { + let sandboxSignal: AbortSignal | undefined + let finishCleanup!: () => void + const cleanup = new Promise((resolve) => { finishCleanup = resolve }) + const { app } = buildHarness({ + getSandbox: async () => ({ + async *streamPrompt(_message: string, opts?: { signal?: AbortSignal }) { + sandboxSignal = opts?.signal + try { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'partial' } } + await new Promise((resolve) => { + opts?.signal?.addEventListener('abort', () => resolve(), { once: true }) + }) + } finally { + finishCleanup() + } + }, + }), + }) + const response = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: 'Bearer sk_agent_cancel', + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + const reader = response.body!.getReader() + await reader.read() + await reader.cancel() + await cleanup + + expect(sandboxSignal?.aborted).toBe(true) + }) + it('records attribution before settlement so adapters can resolve the charge', async () => { const order: string[] = [] const { app } = buildHarness({ @@ -471,6 +506,58 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(settlements).toHaveLength(1) }) + it('requires complete receipts only for requests with durable payment ownership', async () => { + const operations = new MemoryPaymentOperations() + const { app } = buildHarness({ + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'legacy' } } + }, + }), + verifyApiKey: async () => ({ consumerId: 'api-consumer', keyId: 'api-key', scopes: ['chat'] }), + mpp: { + realm: 'agents.tangle.tools', + method: 'stripe', + verifySigner: async () => 'mpp:consumer', + }, + x402: { + operatorAddress, + chainId: 3799, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + }) + const apiKeyResponse = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: 'Bearer sk_agent_legacy', + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + const mppCredential = Buffer.from(JSON.stringify({})).toString('base64url') + const mppResponse = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${mppCredential}`, + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + + const [apiKeyStream, mppStream] = await Promise.all([ + readSse(apiKeyResponse), + readSse(mppResponse), + ]) + expect(apiKeyResponse.status).toBe(200) + expect(mppResponse.status).toBe(200) + expect(apiKeyStream.combinedText).toBe('legacy') + expect(mppStream.combinedText).toBe('legacy') + expect(apiKeyStream.done).toBe(true) + expect(mppStream.done).toBe(true) + }) + it('threads a unique requestId per concurrent request — regression: two same-consumer requests get distinct ids', async () => { const { app, settlements, usage } = buildHarness({}, ['ok']) const requests = await Promise.all([ @@ -831,7 +918,7 @@ describe('createAgentGateway — production-config guard', () => { })).toThrow(/paymentProtocolVersion must be explicit/) }) - it('refuses a custom A2A task store without atomic transitions', () => { + it('keeps older custom A2A task stores source-compatible', () => { expect(() => createAgentGateway({ resolveAgent: async () => null, getSandbox: async () => ({ async *streamPrompt() { /* unused */ } }), @@ -844,7 +931,7 @@ describe('createAgentGateway — production-config guard', () => { delete: async () => undefined, }, }, - })).toThrow(/atomic createIfAbsent and compareAndSet/) + })).not.toThrow() }) }) diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index df6f467..bdb3fa6 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -60,6 +60,55 @@ function context(requiredAmount = 500n): PaymentAuthorizationContext { } describe('version 2 payment operations', () => { + it('settles durable work before a usage-store failure can trigger recovery', async () => { + const order: string[] = [] + let operationId: string | undefined + const operations = new MemoryPaymentOperations({ + onSettle: async (operation) => { + order.push('settle') + operationId = operation.operationId + }, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'sandbox.usage', data: { usage: settledUsage() } } + }, + }), + recordUsage: async () => { + order.push('record') + throw new Error('usage database unavailable') + }, + x402: { + operatorAddress: '0x1', + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const response = await app.request('/v1/agents/payment-tests/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': JSON.stringify({ + ...payload('1000000000', '20'), + operator: '0x1', + }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + await response.text() + + expect(order).toEqual(['settle', 'record']) + expect(operationId).toBeDefined() + expect(operations.get(operationId!)?.state).toBe('settled') + }) + it('has one atomic owner and refunds the unused reservation', async () => { const operations = new MemoryPaymentOperations() const results = await Promise.allSettled([ From 8451fb71a37e7eadf8ffb1ca12679c78fe415b25 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 00:47:39 -0600 Subject: [PATCH 12/32] ci: verify gateway pull requests --- .github/workflows/publish.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 2fb6f5b..00e7ca2 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -1,6 +1,7 @@ name: Publish on: + pull_request: push: tags: - 'v*' From 7f597b399dbf844107991517a12340f754085896 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 17:03:55 -0600 Subject: [PATCH 13/32] fix(payments): close replay and cancellation gaps --- src/dispatch.ts | 105 +++++++++++++++++++++-- src/middleware.ts | 4 +- src/nonce-store.ts | 14 +-- src/payment-operations.ts | 4 +- src/verify.ts | 66 ++++++++------ tests/a2a-payment-races.test.ts | 2 + tests/middleware.test.ts | 101 +++++++++++++++++++++- tests/nonce-store.test.ts | 33 ++++++- tests/payment-operations.test.ts | 142 +++++++++++++++++++++++++++++-- tests/protocol-guards.test.ts | 11 ++- tests/verify.test.ts | 29 +++++++ 11 files changed, 454 insertions(+), 57 deletions(-) diff --git a/src/dispatch.ts b/src/dispatch.ts index 4c1f52e..66a29cd 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -29,6 +29,7 @@ import { defaultVerifyApiKey, isApiKeyAuthEnabled, isMppAuthEnabled, + mppPaymentPayload, mppReplayNonceKey, verifyMpp, verifyX402, @@ -313,6 +314,7 @@ export async function authenticateAndGuard( } consumerId = signer paymentMethod = 'mpp' + x402Payload = mppPaymentPayload(authHeader) ?? null paymentNonceKey = mppReplayNonceKey(authHeader) } else if (authHeader.startsWith('Bearer ')) { const verify = config.verifyApiKey ?? (config.x402.demoMode ? defaultVerifyApiKey : null) @@ -571,9 +573,47 @@ export async function claimPayment( } } authz.paymentOperation = operation - } else if (authz.paymentMethod === 'mpp' && authz.paymentNonceKey) { - const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) - if (!claimed) throw new Error('payment nonce was already consumed') + } else if (authz.paymentMethod === 'mpp') { + const durablePayload = durableMppPaymentPayload(authz.paymentPayload) + if (durablePayload && config.x402.paymentOperations) { + const context = { + requestId: authz.requestId, + agentId: authz.agent.id, + requiredAmount: authz.requiredPaymentAmount, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + } + const operation = await config.x402.paymentOperations.claimPayment(durablePayload, context) + if (operation.protocolVersion !== 2) throw new Error('payment operation protocol version mismatch') + if (operation.acquiredByRequestId !== context.requestId) { + throw new Error('payment operation was already claimed') + } + authz.paymentPayload = durablePayload + authz.paymentOperation = operation + authz.paymentOperationAcquired = true + if (authz.paymentNonceKey) { + const claimed = await claimPaymentNonce( + state.nonceStore, + authz.paymentNonceKey, + durablePayload, + `${operation.operationId}:${context.requestId}`, + ) + if (!claimed) { + try { + await config.x402.paymentOperations.releasePayment(operation, 'shared payment nonce was already owned') + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } + throw new Error('payment nonce was already consumed') + } + } + } else if (authz.paymentNonceKey) { + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) + if (!claimed) throw new Error('payment nonce was already consumed') + } } try { @@ -661,6 +701,31 @@ async function claimPaymentNonce( return claimStoredNonce(nonceStore, nonceKey, ttl, ownerId) } +function durableMppPaymentPayload( + payload: Record | null, +): Record | undefined { + if (!payload) return undefined + const commitment = payload.commitment ?? payload.from + const amount = payload.amount ?? payload.value + const nonce = payload.nonce + if (typeof commitment !== 'string' || commitment.length === 0) return undefined + if (amount === undefined || nonce === undefined) return undefined + const amountText = String(amount) + const nonceText = String(nonce) + if (!/^\d+$/.test(amountText) || !/^\d+$/.test(nonceText)) return undefined + const expiryText = payload.expiry === undefined + ? String(Math.floor(Date.now() / 1000) + 3600) + : String(payload.expiry) + if (!/^\d+$/.test(expiryText)) return undefined + return { + ...payload, + commitment, + amount: amountText, + nonce: nonceText, + expiry: expiryText, + } +} + /** * Yield the inner sandbox's response as text deltas, applying the * system-prompt redaction filter on each delta so leakage of the agent's @@ -724,8 +789,11 @@ export async function* dispatchSandboxStreamRich( maxOutputTokens?: number, onExecutionStart?: () => Promise, requiresReceipt = config.x402.paymentOperations !== undefined, + onSandboxStart?: () => void, ): AsyncIterable { + if (signal?.aborted) return const box = await config.getSandbox(agent) + if (signal?.aborted) return const outputLimit = maxOutputTokens ?? config.defaultOutputTokens ?? 1024 if (!Number.isSafeInteger(outputLimit) || outputLimit <= 0) { throw new Error('max output tokens must be a positive safe integer') @@ -764,6 +832,7 @@ export async function* dispatchSandboxStreamRich( executionBudget, signal, }) + onSandboxStart?.() const iterator = promptStream[Symbol.asyncIterator]() try { while (true) { @@ -837,7 +906,7 @@ export async function* dispatchSandboxStreamRich( ) yield { kind: 'usage', usage } } finally { - closeSandboxIterator(iterator) + await closeSandboxIterator(iterator) } } @@ -850,7 +919,10 @@ async function readSandboxEvent( if (!signal) return iterator.next() if (signal.aborted) return ABORTED_SANDBOX_READ return new Promise((resolve, reject) => { - const onAbort = () => resolve(ABORTED_SANDBOX_READ) + const onAbort = () => { + signal.removeEventListener('abort', onAbort) + resolve(ABORTED_SANDBOX_READ) + } signal.addEventListener('abort', onAbort, { once: true }) iterator.next().then( (result) => { @@ -865,12 +937,27 @@ async function readSandboxEvent( }) } -function closeSandboxIterator(iterator: AsyncIterator): void { +const SANDBOX_CLEANUP_TIMEOUT_MS = 50 + +async function closeSandboxIterator(iterator: AsyncIterator): Promise { + let closing: PromiseLike | undefined try { - const closing = iterator.return?.() - if (closing) void Promise.resolve(closing).catch(() => undefined) + const result = iterator.return?.() + if (result) closing = Promise.resolve(result) } catch { - // The request is already canceled. Adapter cleanup must not delay it. + return + } + if (!closing) return + let timeout: ReturnType | undefined + try { + await Promise.race([ + Promise.resolve(closing).catch(() => undefined), + new Promise((resolve) => { + timeout = setTimeout(resolve, SANDBOX_CLEANUP_TIMEOUT_MS) + }), + ]) + } finally { + if (timeout !== undefined) clearTimeout(timeout) } } diff --git a/src/middleware.ts b/src/middleware.ts index 6d358e0..70656f5 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -311,9 +311,11 @@ function streamChatCompletions( maxOutputTokens, async () => { await beginPaymentExecution(authz, config) - if (authz.paymentOperation) workObserved = true }, authz.paymentOperation !== undefined, + () => { + if (authz.paymentOperation) workObserved = true + }, )) { if (event.kind === 'text') { sendChunk(event.delta) diff --git a/src/nonce-store.ts b/src/nonce-store.ts index c21809e..7628af6 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -121,19 +121,21 @@ export class KvNonceStore implements NonceStore { } async claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise { - if (ownerId === undefined) { + if (!this.kv.putIfAbsent) { + if (ownerId !== undefined) { + throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') + } if (await this.hasSeen(nonce)) return false await this.markSeen(nonce, ttlSeconds) return true } - if (!this.kv.putIfAbsent) { - throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') - } const ttl = Math.max(ttlSeconds, 60) const key = this.key(nonce) const value = ownerId ?? '1' - const existing = await this.kv.get(key) - if (existing !== null) return ownerId !== undefined && existing === ownerId + if (ownerId !== undefined) { + const existing = await this.kv.get(key) + if (existing !== null) return existing === ownerId + } const inserted = await this.kv.putIfAbsent(key, value, { expirationTtl: ttl }) if (inserted || ownerId === undefined) return inserted // Another isolate may have won between get and putIfAbsent. Re-read so a diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 97012c9..9a52305 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -93,8 +93,8 @@ export class MemoryPaymentOperations implements PaymentOperations { private readonly now: () => number constructor(private readonly options: MemoryPaymentOperationsOptions = {}) { - if (options.onClaim && !options.onReclaim) { - throw new Error('onReclaim is required when onClaim can reserve external funds') + if ((options.onClaim || options.onSettle || options.onRelease) && !options.onReclaim) { + throw new Error('onReclaim is required when payment callbacks can lose acknowledgement') } this.now = options.now ?? (() => Math.floor(Date.now() / 1000)) } diff --git a/src/verify.ts b/src/verify.ts index 63002b4..c0854e1 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -3,17 +3,42 @@ import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-stor /** Return the canonical opaque nonce key used by the final payment claim. */ export function mppReplayNonceKey(authHeader: string): string | undefined { + const decoded = decodeMppCredential(authHeader) + return decoded ? canonicalMppNonceKey(decoded.method, decoded.payload) : undefined +} + +/** Return the decoded MPP payload for a durable payment claim. */ +export function mppPaymentPayload(authHeader: string): Record | undefined { + return decodeMppCredential(authHeader)?.payload +} + +interface DecodedMppCredential { + method: string + credential: string + payload: Record +} + +function decodeMppCredential(authHeader: string): DecodedMppCredential | undefined { const match = authHeader.match(/^Payment\s+(\S+)\s+(\S+)$/i) if (!match) return undefined - const [, method, credentialB64] = match + const [, rawMethod, credentialB64] = match + if (!/^[A-Za-z0-9_-]+$/.test(credentialB64)) return undefined try { const decoded = Buffer.from(credentialB64, 'base64url').toString('utf-8') - const credential = JSON.parse(decoded) as Record - const nested = credential.payload - const payload = nested && typeof nested === 'object' && !Array.isArray(nested) - ? nested as Record - : credential - return canonicalMppNonceKey(method, payload) + let payload: Record = {} + try { + const credential = JSON.parse(decoded) as unknown + if (credential && typeof credential === 'object' && !Array.isArray(credential)) { + const record = credential as Record + const nested = record.payload + payload = nested && typeof nested === 'object' && !Array.isArray(nested) + ? nested as Record + : record + } + } catch { + // Method-specific verifiers may accept a non-JSON credential format. + } + return { method: rawMethod.toLowerCase(), credential: decoded, payload } } catch { return undefined } @@ -42,7 +67,7 @@ export function isApiKeyAuthEnabled( export function isMppAuthEnabled( config: Pick, ): boolean { - const method = config.mpp?.method ?? 'blueprintevm' + const method = (config.mpp?.method ?? 'blueprintevm').toLowerCase() return Boolean( config.mpp && (config.mpp.verifySigner !== undefined || @@ -84,7 +109,7 @@ export async function verifyX402( // Reject payments that cannot cover the request's maximum charge. The // check runs before the host verifier because that callback can reserve or // settle funds as part of its production verification path. - if (amount < minimumAmount || minimumAmount <= 0n) return null + if (amount <= 0n || minimumAmount < 0n || amount < minimumAmount) return null const nonceKey = `${String(raw.commitment).toLowerCase()}:${nonce.toString()}` if (nonceStore && await nonceStore.hasSeen(nonceKey)) return null @@ -136,25 +161,14 @@ export async function verifyMpp( const match = authHeader.match(/^Payment\s+(\S+)\s+(\S+)$/i) if (!match) return null - const [, method, credentialB64] = match - if (config.method && method !== config.method) return null + const [, rawMethod] = match + const method = rawMethod.toLowerCase() + if (config.method && method !== config.method.toLowerCase()) return null try { - if (!/^[A-Za-z0-9_-]+$/.test(credentialB64)) return null - const decoded = Buffer.from(credentialB64, 'base64url').toString('utf-8') - let payload: Record = {} - try { - const credential = JSON.parse(decoded) as unknown - if (credential && typeof credential === 'object' && !Array.isArray(credential)) { - const nested = (credential as Record).payload - payload = - nested && typeof nested === 'object' && !Array.isArray(nested) - ? (nested as Record) - : (credential as Record) - } - } catch { - // Method-specific verifiers may accept a non-JSON credential format. - } + const decodedCredential = decodeMppCredential(authHeader) + if (!decodedCredential) return null + const { credential: decoded, payload } = decodedCredential // Validate common EVM fields before a production verifier can reserve or // settle funds. BlueprinTEVM carries x402-equivalent token amounts and diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 2fafbb7..385e981 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -268,6 +268,7 @@ describe('A2A payment ownership races', () => { let settlements = 0 const operations = new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 }, + onReclaim: async () => undefined, }) const config: GatewayConfig = { resolveAgent: async () => agent, @@ -350,6 +351,7 @@ describe('A2A payment ownership races', () => { let settlements = 0 const operations = new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 }, + onReclaim: async () => undefined, }) const config: GatewayConfig = { resolveAgent: async () => agent, diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index 1f51220..0f2e38f 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -16,7 +16,7 @@ import type { GatewayUsageEvent, ApiKeyInfo, } from '../src/types' -import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryNonceStore, type NonceStore } from '../src/nonce-store' import { MemoryRateLimitStore } from '../src/rate-limit' import { MemoryPaymentOperations } from '../src/payment-operations' @@ -464,6 +464,54 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(sandboxSignal?.aborted).toBe(true) }) + it('releases a durable payment when cancellation wins after authorization but before sandbox start', async () => { + const controller = new AbortController() + const operations = new MemoryPaymentOperations() + const originalBegin = operations.beginPaymentExecution.bind(operations) + operations.beginPaymentExecution = async (operation) => { + const executing = await originalBegin(operation) + controller.abort() + return executing + } + let sandboxCalls = 0 + const { app } = buildHarness({ + getSandbox: async () => ({ + async *streamPrompt() { + sandboxCalls += 1 + yield { type: 'sandbox.usage', data: { usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0, + budgetEnforced: true, + } } } + }, + }), + x402: { + operatorAddress, + chainId: 3799, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + }) + const response = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + signal: controller.signal, + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': buildSpendAuth({ nonce: '9003' }), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + await response.text() + + expect(sandboxCalls).toBe(0) + expect(operations.get('x402:0xcommitmentalice:9003')?.state).toBe('released') + }) + it('records attribution before settlement so adapters can resolve the charge', async () => { const order: string[] = [] const { app } = buildHarness({ @@ -506,6 +554,57 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(settlements).toHaveLength(1) }) + it('durably claims an x402-compatible MPP receipt and rejects its replay', async () => { + const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) + const nonceStore: NonceStore = { + hasSeen: async () => false, + claim: async () => true, + markSeen: async () => undefined, + } + const credential = Buffer.from(JSON.stringify({ + payload: { + commitment: '0xCommitmentAlice', + signature: '0xSignatureBytes', + operator: operatorAddress, + amount: fundedRequestAmount, + nonce: '901', + expiry: String(Math.floor(Date.now() / 1000) + 600), + }, + })).toString('base64url') + const { app } = buildHarness({ + nonceStore, + mpp: { + realm: 'agents.tangle.tools', + method: 'blueprintevm', + verifySigner: async () => 'mpp:consumer', + }, + x402: { + operatorAddress, + chainId: 3799, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + }) + const request = () => app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment blueprintevm ${credential}`, + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + + const first = await request() + await first.text() + const second = await request() + await second.text() + + expect(first.status).toBe(200) + expect(second.status).toBe(402) + expect(operations.get('x402:0xcommitmentalice:901')?.state).toBe('settled') + }) + it('requires complete receipts only for requests with durable payment ownership', async () => { const operations = new MemoryPaymentOperations() const { app } = buildHarness({ diff --git a/tests/nonce-store.test.ts b/tests/nonce-store.test.ts index 8161666..7862853 100644 --- a/tests/nonce-store.test.ts +++ b/tests/nonce-store.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, afterEach } from 'vitest' -import { MemoryNonceStore, nonceTtlSeconds } from '../src/nonce-store' +import { KvNonceStore, MemoryNonceStore, nonceTtlSeconds, type KVNamespace } from '../src/nonce-store' describe('nonceTtlSeconds', () => { it('covers the complete signed validity window', () => { @@ -65,3 +65,34 @@ describe('MemoryNonceStore', () => { } }) }) + +describe('KvNonceStore', () => { + it('allows only one mixed legacy and version 2 claim', async () => { + const values = new Map() + const kv: KVNamespace = { + async get(key) { + return values.get(key) ?? null + }, + async put(key, value) { + values.set(key, value) + }, + async putIfAbsent(key, value) { + if (values.has(key)) return false + values.set(key, value) + return true + }, + async delete(key) { + values.delete(key) + }, + } + const store = new KvNonceStore(kv) + + const results = await Promise.all([ + store.claim('mixed-version', 60), + store.claim('mixed-version', 60, 'operation-1'), + ]) + + expect(results.filter(Boolean)).toHaveLength(1) + expect(await store.claim('mixed-version', 60, 'operation-2')).toBe(false) + }) +}) diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index bdb3fa6..0afcfb1 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { Hono } from 'hono' import { dispatchSandboxStreamRich, @@ -68,6 +68,7 @@ describe('version 2 payment operations', () => { order.push('settle') operationId = operation.operationId }, + onReclaim: async () => undefined, }) const config: GatewayConfig = { resolveAgent: async () => agent, @@ -221,6 +222,13 @@ describe('version 2 payment operations', () => { })).toThrow('onReclaim is required') }) + it.each([ + ['settlement', { onSettle: async () => undefined }], + ['release', { onRelease: async () => undefined }], + ])('requires recovery for an external %s acknowledgement', (_name, options) => { + expect(() => new MemoryPaymentOperations(options)).toThrow('onReclaim is required') + }) + it('runs one settlement side effect for concurrent retries', async () => { let effects = 0 const operations = new MemoryPaymentOperations({ @@ -228,6 +236,7 @@ describe('version 2 payment operations', () => { effects += 1 await new Promise((resolve) => setTimeout(resolve, 1)) }, + onReclaim: async () => undefined, }) const owner = await operations.claimPayment(payload('1000', '13'), context()) const input = { @@ -363,14 +372,10 @@ describe('version 2 payment operations', () => { expect(recovered[1].state).toBe('released') }) - it('does not close a release when recovery has no refund proof', async () => { - const operations = new MemoryPaymentOperations({ + it('rejects a release callback without a recovery proof', () => { + expect(() => new MemoryPaymentOperations({ onRelease: async () => { throw new Error('acknowledgement lost') }, - }) - const owner = await operations.claimPayment(payload('1000', '14'), context()) - await expect(operations.releasePayment(owner, 'sandbox failed')).rejects.toThrow() - await expect(operations.reclaimPayment(owner.operationId)).rejects.toThrow('recovery is not configured') - expect(operations.get(owner.operationId)?.state).toBe('releasing') + })).toThrow('onReclaim is required') }) it('reclaims a claim that crashed after durable ownership but before completion', async () => { @@ -491,7 +496,10 @@ describe('bounded request pricing and sandbox receipts', () => { it('bounds output before delivery and retains payment ownership after an over-limit stream', async () => { let releases = 0 - const operations = new MemoryPaymentOperations({ onRelease: async () => { releases += 1 } }) + const operations = new MemoryPaymentOperations({ + onRelease: async () => { releases += 1 }, + onReclaim: async () => undefined, + }) const box: SandboxBox = { async *streamPrompt() { yield { @@ -671,6 +679,122 @@ describe('bounded request pricing and sandbox receipts', () => { await events.next() await events.next() }) + + it('waits for bounded iterator cleanup after cancellation', async () => { + let signalNext!: () => void + const nextStarted = new Promise((resolve) => { signalNext = resolve }) + let signalCleanup!: () => void + const cleanupStarted = new Promise((resolve) => { signalCleanup = resolve }) + let finishCleanup!: () => void + const cleanupFinished = new Promise>((resolve) => { + finishCleanup = () => resolve({ done: true, value: undefined }) + }) + const iterator: AsyncIterator = { + next: async () => { + signalNext() + return new Promise>(() => undefined) + }, + return: async () => { + signalCleanup() + return cleanupFinished + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + streamPrompt: () => ({ [Symbol.asyncIterator]: () => iterator }), + }), + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + } + const controller = new AbortController() + let completed = false + const running = (async () => { + for await (const _event of dispatchSandboxStreamRich( + agent, + 'hi', + 'consumer', + config, + controller.signal, + undefined, + 4, + )) { + // The iterator never emits an event before cancellation. + } + completed = true + })() + + await nextStarted + controller.abort() + await cleanupStarted + await Promise.resolve() + expect(completed).toBe(false) + + finishCleanup() + await running + expect(completed).toBe(true) + }) + + it('bounds cleanup when an iterator ignores cancellation', async () => { + vi.useFakeTimers() + try { + let signalNext!: () => void + const nextStarted = new Promise((resolve) => { signalNext = resolve }) + let signalCleanup!: () => void + const cleanupStarted = new Promise((resolve) => { signalCleanup = resolve }) + let cleanupCalls = 0 + const iterator: AsyncIterator = { + next: async () => { + signalNext() + return new Promise>(() => undefined) + }, + return: async () => { + cleanupCalls += 1 + signalCleanup() + return new Promise>(() => undefined) + }, + } + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + streamPrompt: () => ({ [Symbol.asyncIterator]: () => iterator }), + }), + recordUsage: async () => undefined, + x402: { operatorAddress: '0x1', chainId: 1, demoMode: true }, + } + const controller = new AbortController() + let completed = false + const running = (async () => { + for await (const _event of dispatchSandboxStreamRich( + agent, + 'hi', + 'consumer', + config, + controller.signal, + undefined, + 4, + )) { + // The iterator never emits an event before cancellation. + } + completed = true + })() + + await nextStarted + controller.abort() + await cleanupStarted + expect(cleanupCalls).toBe(1) + expect(completed).toBe(false) + + vi.advanceTimersByTime(49) + await Promise.resolve() + expect(completed).toBe(false) + vi.advanceTimersByTime(1) + await running + expect(completed).toBe(true) + } finally { + vi.useRealTimers() + } + }) }) function settledUsage() { diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index c7ddff6..f4f3469 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -321,8 +321,14 @@ describe('final payment boundary protocol guards', () => { let runs = 0 let settlements = 0 const stores = [ - new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 } }), - new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 } }), + new MemoryPaymentOperations({ + onSettle: async () => { settlements += 1 }, + onReclaim: async () => undefined, + }), + new MemoryPaymentOperations({ + onSettle: async () => { settlements += 1 }, + onReclaim: async () => undefined, + }), ] const apps = stores.map((operations) => { const app = new Hono() @@ -376,6 +382,7 @@ describe('final payment boundary protocol guards', () => { const operations = new MemoryPaymentOperations({ onSettle: async () => { settlements += 1 }, onRelease: async () => { releases += 1 }, + onReclaim: async () => undefined, }) let operationPromise: ReturnType | undefined let authorizations = 0 diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 1a4952c..92b3180 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -64,6 +64,10 @@ describe('verifyX402', () => { expect(calls).toBe(0) }) + it('accepts a positive authorization when the requested agent price is zero', async () => { + expect(await verifyX402(buildSpendAuth({ amount: '1' }), baseConfig, undefined, 0n)).toBe('0xCommitmentAlice') + }) + it('rejects expired payments — regression: forever-valid sigs enable drained-wallet attacks', async () => { const expired = buildSpendAuth({ expiry: String(Math.floor(Date.now() / 1000) - 10) }) expect(await verifyX402(expired, baseConfig)).toBeNull() @@ -256,6 +260,31 @@ describe('verifyMpp', () => { expect(calls).toBe(0) }) + it('normalizes MPP method casing before applying the configured payment ceiling', async () => { + const underfunded = buildCredential({ + commitment: '0xAlice', + operator: operatorAddress, + amount: '19999', + nonce: '9', + expiry: String(Math.floor(Date.now() / 1000) + 600), + }).replace('Payment blueprintevm ', 'Payment BLUEPRINTEVM ') + + expect(await verifyMpp( + underfunded, + { realm: 'agents.tangle.tools' }, + baseConfig, + undefined, + 20000n, + )).toBeNull() + expect(await verifyMpp( + underfunded, + mppConfig, + baseConfig, + undefined, + 1n, + )).toBe('0xAlice') + }) + it('rejects MPP in production when no method verifier is configured', async () => { const header = buildCredential({ commitment: '0xAlice', operator: operatorAddress }) const productionX402: X402Config = { ...baseConfig, demoMode: false } From f51b44a6fb518909ed433f70c5652b22d70080d9 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 17:05:12 -0600 Subject: [PATCH 14/32] fix(a2a): enforce atomic task ownership and recovery --- docs/a2a-long-horizon.md | 11 +- src/a2a/handler.ts | 556 ++++++++++++++++++++++++++++---- src/a2a/task-store-sql.ts | 5 +- src/a2a/task-store.ts | 4 +- tests/a2a-atomicity.test.ts | 433 +++++++++++++++++++++++++ tests/a2a-payment-races.test.ts | 2 +- 6 files changed, 933 insertions(+), 78 deletions(-) create mode 100644 tests/a2a-atomicity.test.ts diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 1e3882d..19d49bc 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -15,10 +15,17 @@ All four are gated on configuration — they cost nothing for agents that don't By default `GatewayConfig.a2a.taskStore` is in-memory: fast, zero-config, fine for tests and single-machine deployments. Production deployments swap in `SqlTaskStore` against any SQL store — D1, postgres, sqlite, libSQL, Turso — via a `SqlAdapter` shim. -Custom production task stores should implement both atomic methods, `createIfAbsent` and `compareAndSet`. -Older stores remain compatible through a read-then-write fallback, which is process-safe only and can run paid work twice across workers. +Custom production task stores must implement both atomic methods, `createIfAbsent` and `compareAndSet`. +The gateway rejects a production task store that lacks either method before it serves A2A requests. +Explicit `x402.demoMode` keeps a read-then-write fallback for local tests only. Use `SqlTaskStore` or another atomic adapter for multi-worker production deployments. +Before payment settlement, the gateway stores a recovery record in task metadata. +The record contains the payment operation, usage receipt, output artifact, and a five-minute lease. +After a restart, the first task read after lease expiry resumes settlement with the same operation. +If settlement acknowledgement fails, the gateway keeps the operation and retries after the lease expires. +If the record is malformed or has no recoverable operation, the gateway expires the task as failed. + Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. The hook receives the task and request headers so the application can enforce task ownership. Explicit `x402.demoMode` permits these methods without the hook for local tests only. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 280a17e..eb6071b 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -24,8 +24,13 @@ import { releasePaymentAfterFailure, settleAndRecord, } from '../dispatch' -import type { GatewayConfig } from '../types' -import type { SandboxUsageReceipt } from '../types' +import type { PaymentOperation } from '../payment-operations' +import type { + GatewayConfig, + PaymentMethod, + SandboxExecutionBudget, + SandboxUsageReceipt, +} from '../types' import { buildAgentCard } from './agent-card' import { fail, ok, parseEnvelope } from './jsonrpc' import { @@ -37,6 +42,7 @@ import type { TaskStore } from './task-store' import { extractTextFromMessage, responseTextToArtifact } from './translate' import { A2A_ERROR_CODES, + type Artifact, type JSONRPCRequest, type Message, type MessageSendParams, @@ -55,6 +61,41 @@ export interface A2AHandlerDeps { pushStore?: PushNotificationStore } +interface SerializedPaymentOperation { + protocolVersion: 2 + operationId: string + acquiredByRequestId: string + executionStartedAt?: number + retentionReason?: string + nonceKey: string + authorizationId: string + reservedAmount: string + settledAmount: string + refundAmount: string + expiresAt: number + state: PaymentOperation['state'] +} + +interface FinalizationRecord { + version: 1 + lease: { id: string; expiresAt: number } + agentSlug: string + requestId: string + consumerId: string + paymentMethod: PaymentMethod + startMs: number + operationId: string | null + paymentOperation: SerializedPaymentOperation | null + receipt: SandboxUsageReceipt + artifact: Artifact | null + inputRequired: boolean + inputRequiredPrompt?: string + maxOutputTokens: number + executionBudget: SandboxExecutionBudget + recoveryAttempts?: number + recoveryError?: string +} + /** Terminal task states — fire-once push delivery occurs on these transitions. */ const TERMINAL_STATES: ReadonlySet = new Set([ 'completed', @@ -106,6 +147,10 @@ class CancelRegistry { } export function createA2AHandlers(deps: A2AHandlerDeps) { + const runtimeDeps: A2AHandlerDeps = { + ...deps, + taskStore: normalizeTaskStore(deps.taskStore, deps.config.x402.demoMode === true), + } const cancels = new CancelRegistry() // GET /:slug/.well-known/agent.json @@ -147,23 +192,23 @@ export function createA2AHandlers(deps: A2AHandlerDeps) { switch (parsed.method) { case 'message/send': - return handleMessageSend(c, slug, parsed, deps, cancels) + return handleMessageSend(c, slug, parsed, runtimeDeps, cancels) case 'message/stream': - return handleMessageStream(c, slug, parsed, deps, cancels) + return handleMessageStream(c, slug, parsed, runtimeDeps, cancels) case 'tasks/get': - return handleTasksGet(c, parsed, deps) + return handleTasksGet(c, parsed, runtimeDeps) case 'tasks/cancel': - return handleTasksCancel(c, parsed, deps, cancels) + return handleTasksCancel(c, parsed, runtimeDeps, cancels) case 'tasks/resubscribe': - return handleTasksResubscribe(c, parsed, deps) + return handleTasksResubscribe(c, parsed, runtimeDeps) case 'tasks/pushNotificationConfig/set': - return handlePushSet(c, parsed, deps) + return handlePushSet(c, parsed, runtimeDeps) case 'tasks/pushNotificationConfig/get': - return handlePushGet(c, parsed, deps) + return handlePushGet(c, parsed, runtimeDeps) case 'tasks/pushNotificationConfig/list': - return handlePushList(c, parsed, deps) + return handlePushList(c, parsed, runtimeDeps) case 'tasks/pushNotificationConfig/delete': - return handlePushDelete(c, parsed, deps) + return handlePushDelete(c, parsed, runtimeDeps) default: return c.json( fail(parsed.id, A2A_ERROR_CODES.METHOD_NOT_FOUND, `unknown method '${parsed.method}'`), @@ -222,6 +267,7 @@ async function executeMessageSend( let workObserved = false let inputRequiredPrompt: string | undefined let inputRequiredSeen = false + let finalizationLeaseId: string | undefined try { for await (const event of dispatchSandboxStreamRich( authz.agent, @@ -296,7 +342,18 @@ async function executeMessageSend( // commercial behavior — the sandbox produced tokens. try { if (!usage) throw new Error('sandbox did not provide a usage receipt') - if (!await claimTaskFinalization(deps.taskStore, workingTask)) { + const finalizationArtifact = responseText + ? responseTextToArtifact(responseText, `${task.id}-artifact-0`) + : task.artifacts?.[0] ?? null + const finalization = buildFinalizationRecord( + authz, + usage, + finalizationArtifact, + inputRequiredSeen, + inputRequiredPrompt, + ) + const finalizingTask = withFinalizationRecord(workingTask, finalization) + if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask?.status.state === 'canceled') { const canceled = await completeCanceledTask( @@ -311,7 +368,31 @@ async function executeMessageSend( } throw new Error('A2A task changed before payment settlement') } + finalizationLeaseId = finalization.lease.id await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) + const settledBase = clearFinalizationMarker(finalizingTask) + const result = inputRequiredSeen + ? withStatus( + settledBase, + 'input-required', + inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, + responseText + ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] + : task.artifacts, + ) + : withStatus(settledBase, 'completed', undefined, [ + responseTextToArtifact(responseText, `${task.id}-artifact-0`), + ]) + if (!await compareAndSetTask(deps.taskStore, finalizingTask, result)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask && (isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required')) { + return c.json(ok(req.id, currentTask)) + } + throw new Error('A2A task changed after payment settlement') + } + if (inputRequiredSeen) return c.json(ok(req.id, result)) + await maybeDeliverPush(result, deps) + return c.json(ok(req.id, result)) } catch (err) { await releaseOrRetainPayment( authz, @@ -319,6 +400,15 @@ async function executeMessageSend( err instanceof Error ? err.message : String(err), workObserved || usage !== undefined, ) + if (finalizationLeaseId) { + await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalizationLeaseId, + asError(err), + ) + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) + } const currentTask = await deps.taskStore.get(task.id) const failed = currentTask && isTerminal(currentTask.status.state) ? currentTask @@ -334,27 +424,6 @@ async function executeMessageSend( } return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) } - - if (inputRequiredSeen) { - const paused = withStatus( - workingTask, - 'input-required', - inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, - responseText - ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] - : task.artifacts, - ) - await deps.taskStore.put(paused) - // input-required is non-terminal — do NOT deliver push notifications. - return c.json(ok(req.id, paused)) - } - - const completed = withStatus(workingTask, 'completed', undefined, [ - responseTextToArtifact(responseText, `${task.id}-artifact-0`), - ]) - await deps.taskStore.put(completed) - await maybeDeliverPush(completed, deps) - return c.json(ok(req.id, completed)) } // ── message/stream (SSE) ────────────────────────────────────────────────── @@ -395,6 +464,7 @@ async function handleMessageStream( : { ...task, status: workingStatus.status } let inputRequiredPrompt: string | undefined let inputRequiredSeen = false + let finalizationLeaseId: string | undefined try { if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { throw new Error('A2A task changed before execution started') @@ -465,10 +535,19 @@ async function handleMessageStream( // Settle once for whatever the sandbox produced (full or partial). if (!usage) throw new Error('sandbox did not provide a usage receipt') + const finalizationArtifact = responseTextToArtifact(responseText, `${task.id}-artifact-0`) + const finalization = buildFinalizationRecord( + authz, + usage, + finalizationArtifact, + inputRequiredSeen, + inputRequiredPrompt, + ) // Let the durable task-store CAS decide the cancellation race. Mark // the local registry only after that CAS wins, so cancel can replace a // still-pending finalization instead of being rejected prematurely. - if (!await claimTaskFinalization(deps.taskStore, workingTask)) { + const finalizingTask = withFinalizationRecord(workingTask, finalization) + if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask?.status.state === 'canceled') { const canceled = await completeCanceledTask( @@ -496,19 +575,32 @@ async function handleMessageStream( ) return } + finalizationLeaseId = finalization.lease.id cancels.beginFinalization(task.id) await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) if (inputRequiredSeen) { const paused = withStatus( - task, + clearFinalizationMarker(finalizingTask), 'input-required', inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, responseText ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] : task.artifacts, ) - await deps.taskStore.put(paused) + if (!await compareAndSetTask(deps.taskStore, finalizingTask, paused)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: currentTask.status, + final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', + }) + } + return + } send({ kind: 'status-update', taskId: task.id, @@ -520,7 +612,23 @@ async function handleMessageStream( return } - // Final: artifact lastChunk + completed status. + // Final: persist the terminal task before emitting terminal events. + const completed = withStatus(clearFinalizationMarker(finalizingTask), 'completed', undefined, [ + responseTextToArtifact(responseText, `${task.id}-artifact-0`), + ]) + if (!await compareAndSetTask(deps.taskStore, finalizingTask, completed)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: currentTask.status, + final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', + }) + } + return + } send({ kind: 'artifact-update', taskId: task.id, @@ -533,10 +641,6 @@ async function handleMessageStream( append: true, lastChunk: true, }) - const completed = withStatus(task, 'completed', undefined, [ - responseTextToArtifact(responseText, `${task.id}-artifact-0`), - ]) - await deps.taskStore.put(completed) send({ kind: 'status-update', taskId: task.id, @@ -552,6 +656,24 @@ async function handleMessageStream( err instanceof Error ? err.message : String(err), workObserved || usage !== undefined, ) + if (finalizationLeaseId) { + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalizationLeaseId, + asError(err), + ) + if (retained) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: retained.status, + final: false, + }) + } + return + } const currentTask = await deps.taskStore.get(task.id) const failed = currentTask && isTerminal(currentTask.status.state) ? currentTask @@ -619,14 +741,19 @@ async function handleTasksGet( if (!params || typeof params.id !== 'string') { return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) } - const task = await deps.taskStore.get(params.id) - if (!task) { + const storedTask = await deps.taskStore.get(params.id) + if (!storedTask) { return c.json( fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), ) } - const accessError = await authorizeTaskAccess(c, req, task, deps) + const accessError = await authorizeTaskAccess(c, req, storedTask, deps) if (accessError) return accessError + const task = await recoverFinalizationIfNeeded( + storedTask, + deps, + c.req.param('slug') ?? '', + ) return c.json(ok(req.id, task)) } @@ -640,14 +767,19 @@ async function handleTasksCancel( if (!params || typeof params.id !== 'string') { return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) } - const task = await deps.taskStore.get(params.id) - if (!task) { + const storedTask = await deps.taskStore.get(params.id) + if (!storedTask) { return c.json( fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), ) } - const accessError = await authorizeTaskAccess(c, req, task, deps) + const accessError = await authorizeTaskAccess(c, req, storedTask, deps) if (accessError) return accessError + const task = await recoverFinalizationIfNeeded( + storedTask, + deps, + c.req.param('slug') ?? '', + ) if (isTerminal(task.status.state)) { return c.json( fail( @@ -715,14 +847,19 @@ async function handleTasksResubscribe( if (!params || typeof params.id !== 'string') { return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) } - const task = await deps.taskStore.get(params.id) - if (!task) { + const storedTask = await deps.taskStore.get(params.id) + if (!storedTask) { return c.json( fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), ) } - const accessError = await authorizeTaskAccess(c, req, task, deps) + const accessError = await authorizeTaskAccess(c, req, storedTask, deps) if (accessError) return accessError + const task = await recoverFinalizationIfNeeded( + storedTask, + deps, + c.req.param('slug') ?? '', + ) const final = isTerminal(task.status.state) || task.status.state === 'input-required' const event: TaskStatusUpdateEvent = { kind: 'status-update', @@ -902,10 +1039,11 @@ async function guardMessageRequest( // Any other taskId (unknown OR pointing at a terminal/working // task) means the caller is starting a fresh task and we mint a new id. if (typeof params.message.taskId === 'string') { - const existing = await deps.taskStore.get(params.message.taskId) - if (existing) { - const accessError = await authorizeTaskAccess(c, req, existing, deps) + const storedExisting = await deps.taskStore.get(params.message.taskId) + if (storedExisting) { + const accessError = await authorizeTaskAccess(c, req, storedExisting, deps) if (accessError) return accessError + const existing = await recoverFinalizationIfNeeded(storedExisting, deps, slug) if (existing.status.state !== 'input-required') { return c.json( fail( @@ -1021,6 +1159,28 @@ async function releaseOrRetainPayment( } } +/** Keep an owned finalization record durable when settlement acknowledgement is lost. */ +async function retainFinalizationForRecovery( + taskStore: TaskStore, + taskId: string, + leaseId: string, + error: Error, +): Promise { + const current = await taskStore.get(taskId) + if (!current) return undefined + const record = readFinalizationRecord(current) + if (!record || record.lease.id !== leaseId) return undefined + const retry: FinalizationRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, + recoveryAttempts: (record.recoveryAttempts ?? 0) + 1, + recoveryError: error.message, + } + const next = withFinalizationRecord(current, retry) + if (await compareAndSetTask(taskStore, current, next)) return next + return await taskStore.get(taskId) +} + async function completeCanceledTask( authz: AuthorizedRequest, task: Task, @@ -1059,6 +1219,7 @@ async function completeCanceledTask( // ── Helpers ─────────────────────────────────────────────────────────────── const FINALIZING_METADATA_KEY = 'gatewayFinalizing' +const FINALIZATION_LEASE_MS = 5 * 60 * 1000 async function authorizeTaskAccess( c: Context, @@ -1102,32 +1263,283 @@ function isHttpsUrl(value: string): boolean { } async function createTask(taskStore: TaskStore, task: Task): Promise { - if (taskStore.createIfAbsent) return taskStore.createIfAbsent(task) - if (await taskStore.get(task.id)) return false - await taskStore.put(task) - return true + if (!taskStore.createIfAbsent) { + throw new Error('A2A task store does not provide createIfAbsent') + } + return taskStore.createIfAbsent(task) } async function compareAndSetTask(taskStore: TaskStore, expected: Task, next: Task): Promise { - if (taskStore.compareAndSet) return taskStore.compareAndSet(expected, next) - // Older adapters remain source-compatible, but their fallback is only - // process-safe. Durable adapters must implement compareAndSet. - const current = await taskStore.get(expected.id) - if (!current || JSON.stringify(current) !== JSON.stringify(expected)) return false - await taskStore.put(next) - return true + if (!taskStore.compareAndSet) { + throw new Error('A2A task store does not provide compareAndSet') + } + return taskStore.compareAndSet(expected, next) +} + +function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): TaskStore { + if (typeof taskStore.createIfAbsent === 'function' && typeof taskStore.compareAndSet === 'function') { + return taskStore + } + if (!allowUnsafeFallback) { + throw new Error( + 'A2A production task store must implement createIfAbsent and compareAndSet', + ) + } + return { + get: (id) => taskStore.get(id), + put: (task) => taskStore.put(task), + delete: (id) => taskStore.delete(id), + async createIfAbsent(task) { + if (await taskStore.get(task.id)) return false + await taskStore.put(task) + return true + }, + async compareAndSet(expected, next) { + const current = await taskStore.get(expected.id) + if (!current || JSON.stringify(current) !== JSON.stringify(expected)) return false + await taskStore.put(next) + return true + }, + } +} + +function buildFinalizationRecord( + authz: AuthorizedRequest, + receipt: SandboxUsageReceipt, + artifact: Artifact | null, + inputRequired: boolean, + inputRequiredPrompt: string | undefined, +): FinalizationRecord { + const operation = authz.paymentOperation + return { + version: 1, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, + agentSlug: authz.agent.slug, + requestId: authz.requestId, + consumerId: authz.consumerId, + paymentMethod: authz.paymentMethod, + startMs: authz.startMs, + operationId: operation?.operationId ?? null, + paymentOperation: operation ? serializePaymentOperation(operation) : null, + receipt, + artifact, + inputRequired, + ...(inputRequiredPrompt ? { inputRequiredPrompt } : {}), + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + } +} + +function serializePaymentOperation(operation: PaymentOperation): SerializedPaymentOperation { + return { + protocolVersion: 2, + operationId: operation.operationId, + acquiredByRequestId: operation.acquiredByRequestId, + ...(operation.executionStartedAt !== undefined + ? { executionStartedAt: operation.executionStartedAt } + : {}), + ...(operation.retentionReason ? { retentionReason: operation.retentionReason } : {}), + nonceKey: operation.nonceKey, + authorizationId: operation.authorizationId, + reservedAmount: operation.reservedAmount.toString(), + settledAmount: operation.settledAmount.toString(), + refundAmount: operation.refundAmount.toString(), + expiresAt: operation.expiresAt, + state: operation.state, + } +} + +function deserializePaymentOperation(value: SerializedPaymentOperation): PaymentOperation { + if (value.protocolVersion !== 2) throw new Error('unsupported A2A payment operation version') + if (!value.operationId || !value.acquiredByRequestId || !value.nonceKey || !value.authorizationId) { + throw new Error('incomplete A2A payment operation recovery record') + } + return { + protocolVersion: 2, + operationId: value.operationId, + acquiredByRequestId: value.acquiredByRequestId, + ...(value.executionStartedAt !== undefined ? { executionStartedAt: value.executionStartedAt } : {}), + ...(value.retentionReason ? { retentionReason: value.retentionReason } : {}), + nonceKey: value.nonceKey, + authorizationId: value.authorizationId, + reservedAmount: BigInt(value.reservedAmount), + settledAmount: BigInt(value.settledAmount), + refundAmount: BigInt(value.refundAmount), + expiresAt: value.expiresAt, + state: value.state, + } } -async function claimTaskFinalization(taskStore: TaskStore, task: Task): Promise { - if (isTaskFinalizing(task)) return false - return compareAndSetTask(taskStore, task, { +function withFinalizationRecord(task: Task, record: FinalizationRecord): Task { + return { ...task, - metadata: { ...(task.metadata ?? {}), [FINALIZING_METADATA_KEY]: true }, - }) + metadata: { ...(task.metadata ?? {}), [FINALIZING_METADATA_KEY]: record }, + } +} + +function readFinalizationRecord(task: Task): FinalizationRecord | undefined { + const raw = task.metadata?.[FINALIZING_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const record = raw as Partial + if ( + record.version !== 1 || + !record.lease || + typeof record.lease.id !== 'string' || + typeof record.lease.expiresAt !== 'number' + ) { + return undefined + } + return record as FinalizationRecord } function isTaskFinalizing(task: Task): boolean { - return task.metadata?.[FINALIZING_METADATA_KEY] === true + const marker = task.metadata?.[FINALIZING_METADATA_KEY] + return marker === true || (typeof marker === 'object' && marker !== null) +} + +async function recoverFinalizationIfNeeded( + task: Task, + deps: A2AHandlerDeps, + requestedAgentSlug: string, +): Promise { + if (!isTaskFinalizing(task)) return task + const record = readFinalizationRecord(task) + if (!record) { + return expireFinalization( + task, + deps, + null, + new Error('A2A finalization record is missing'), + ) + } + if (record.lease.expiresAt > Date.now()) return task + + const renewed: FinalizationRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, + } + const leasedTask = withFinalizationRecord(task, renewed) + if (!await compareAndSetTask(deps.taskStore, task, leasedTask)) { + return await deps.taskStore.get(task.id) ?? task + } + + try { + const agentSlug = renewed.agentSlug || requestedAgentSlug + const agent = await deps.config.resolveAgent(agentSlug) + if (!agent || !agent.enabled) throw new Error('A2A recovery agent is unavailable') + + let paymentOperation: PaymentOperation | undefined + if (renewed.operationId || renewed.paymentOperation) { + if (!renewed.operationId || !renewed.paymentOperation) { + throw new Error('A2A payment operation recovery record is incomplete') + } + if (renewed.operationId !== renewed.paymentOperation.operationId) { + throw new Error('A2A payment operation recovery id does not match') + } + if (!deps.config.x402.paymentOperations) { + throw new Error('A2A payment operation recovery is not configured') + } + paymentOperation = deserializePaymentOperation(renewed.paymentOperation) + } else if (!deps.config.x402.demoMode) { + throw new Error('A2A production recovery requires a durable payment operation') + } + + const authz: AuthorizedRequest = { + agent, + consumerId: renewed.consumerId, + paymentMethod: renewed.paymentMethod, + keyInfo: null, + userMessage: '[recovered A2A task]', + rateLimitRemaining: undefined, + requestId: renewed.requestId, + startMs: renewed.startMs, + maxOutputTokens: renewed.maxOutputTokens, + executionBudget: renewed.executionBudget, + requiredPaymentAmount: 0n, + paymentPayload: null, + ...(paymentOperation + ? { paymentOperation, paymentOperationAcquired: true } + : {}), + } + await settleAndRecord(agent, authz, renewed.receipt, deps.config, deps.state.obs) + + const recovered = finalizationResultTask(leasedTask, renewed) + if (!await compareAndSetTask(deps.taskStore, leasedTask, recovered)) { + return await deps.taskStore.get(task.id) ?? recovered + } + await maybeDeliverPush(recovered, deps) + return recovered + } catch (error) { + const recoveryError = error instanceof Error ? error : new Error(String(error)) + console.error( + `[a2a] finalization recovery failed for ${task.id}:`, + recoveryError.message, + ) + if (renewed.operationId && renewed.paymentOperation) { + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + renewed.lease.id, + recoveryError, + ) + if (retained) return retained + } + return expireFinalization(leasedTask, deps, renewed, recoveryError) + } +} + +function finalizationResultTask(task: Task, record: FinalizationRecord): Task { + const cleanTask = clearFinalizationMarker(task) + if (record.inputRequired) { + return withStatus( + cleanTask, + 'input-required', + record.inputRequiredPrompt ? agentMessage(cleanTask, record.inputRequiredPrompt) : undefined, + record.artifact ? [record.artifact] : cleanTask.artifacts, + ) + } + return withStatus( + cleanTask, + 'completed', + undefined, + record.artifact ? [record.artifact] : cleanTask.artifacts, + ) +} + +async function expireFinalization( + task: Task, + deps: A2AHandlerDeps, + record: FinalizationRecord | null, + error: Error, +): Promise { + const cleanTask = clearFinalizationMarker(task) + const failed: Task = { + ...withStatus(cleanTask, 'failed'), + metadata: { + ...(cleanTask.metadata ?? {}), + gatewayFinalizationRecovery: { + operationId: record?.operationId ?? null, + error: error.message, + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await maybeDeliverPush(failed, deps) + return failed + } + return await deps.taskStore.get(task.id) ?? failed +} + +function clearFinalizationMarker(task: Task): Task { + if (!task.metadata || !(FINALIZING_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[FINALIZING_METADATA_KEY] + return Object.keys(metadata).length > 0 + ? { ...task, metadata } + : (() => { + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata + })() } function isTerminal(state: Task['status']['state']): boolean { @@ -1139,6 +1551,10 @@ function isTerminal(state: Task['status']['state']): boolean { ) } +function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)) +} + function nowIso(): string { return new Date().toISOString() } diff --git a/src/a2a/task-store-sql.ts b/src/a2a/task-store-sql.ts index 857749a..58a3baf 100644 --- a/src/a2a/task-store-sql.ts +++ b/src/a2a/task-store-sql.ts @@ -9,9 +9,8 @@ * Schema is one table: tasks keyed by id with the full JSON payload, plus a * secondary index on `context_id` so `tasks/resubscribe` and conversational * lookups by context are O(log n). TTL is enforced at read time the same way - * `InMemoryTaskStore` does — the gateway is single-writer per task id so a - * stale row is invisible to callers regardless of when the row is physically - * deleted. + * `InMemoryTaskStore` does — `createIfAbsent` and `compareAndSet` make task + * ownership safe when multiple gateway workers share the database. * * Why not bake in a specific driver? Hono workers run on Cloudflare (D1), * Node (pg / sqlite), Bun, Deno. Burning a hard dependency on one client diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index f320491..85fa1ff 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -10,9 +10,9 @@ import type { Task } from './types' export interface TaskStore { get(id: string): Promise put(task: Task): Promise - /** Insert only when the task id is absent. Required for explicit A2A ids. */ + /** Insert only when the task id is absent. Required outside explicit demo mode. */ createIfAbsent?(task: Task): Promise - /** Replace only when the stored task still equals `expected`. */ + /** Replace only when the stored task still equals `expected`. Required for production races. */ compareAndSet?(expected: Task, next: Task): Promise delete(id: string): Promise } diff --git a/tests/a2a-atomicity.test.ts b/tests/a2a-atomicity.test.ts new file mode 100644 index 0000000..e988a7c --- /dev/null +++ b/tests/a2a-atomicity.test.ts @@ -0,0 +1,433 @@ +import { Hono } from 'hono' +import { describe, expect, it } from 'vitest' + +import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' +import { createAgentGateway } from '../src/middleware' +import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryPaymentOperations, type PaymentOperation } from '../src/payment-operations' +import type { AgentMeta, GatewayConfig, SandboxBox, SandboxUsageReceipt } from '../src/types' +import type { Artifact, Task } from '../src/a2a/types' + +const operatorAddress = '0x1111111111111111111111111111111111111111' +const commitment = `0x${'cd'.repeat(32)}` + +const agent: AgentMeta = { + id: 'agent-a2a-atomicity', + ownerId: 'owner', + slug: 'a2a-atomicity', + systemPrompt: '', + pricePerTokenUsd: 0.000001, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +const receipt: SandboxUsageReceipt = { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.000002, + budgetEnforced: true, +} + +const artifact: Artifact = { + artifactId: 'task-restart-artifact', + name: 'response', + parts: [{ kind: 'text', text: 'recovered output' }], +} + +function paymentHeader(nonce: string): string { + return JSON.stringify({ + commitment, + signature: '0xsig', + operator: operatorAddress, + amount: '1000000000', + nonce, + expiry: String(Math.floor(Date.now() / 1000) + 300), + }) +} + +function requestBody(taskId: string, text: string) { + return { + jsonrpc: '2.0', + id: text, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId, + contextId: 'ctx-atomicity', + messageId: `message-${text}`, + parts: [{ kind: 'text', text }], + }, + }, + } +} + +class TwoWorkerCreateBarrier implements TaskStore { + private readonly inner = new InMemoryTaskStore() + private entered = 0 + private release!: () => void + private readonly bothEntered = new Promise((resolve) => { this.release = resolve }) + + get(id: string): Promise { + return this.inner.get(id) + } + + put(task: Task): Promise { + return this.inner.put(task) + } + + delete(id: string): Promise { + return this.inner.delete(id) + } + + async createIfAbsent(task: Task): Promise { + this.entered += 1 + if (this.entered === 2) this.release() + await this.bothEntered + return this.inner.createIfAbsent(task) + } + + compareAndSet(expected: Task, next: Task): Promise { + return this.inner.compareAndSet(expected, next) + } +} + +function atomicityConfig( + taskStore: TaskStore, + counters: { runs: number; records: number; settlements: number }, +): GatewayConfig { + const sandbox: SandboxBox = { + async *streamPrompt() { + counters.runs += 1 + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'done' } } + yield { type: 'sandbox.usage', data: { usage: receipt } } + }, + } + return { + resolveAgent: async (slug) => (slug === agent.slug ? agent : null), + getSandbox: async () => sandbox, + recordUsage: async () => { counters.records += 1 }, + settlePayment: async () => { counters.settlements += 1 }, + x402: { + operatorAddress, + chainId: 1, + paymentProtocolVersion: 1, + verifySigner: async () => true, + }, + nonceStore: new MemoryNonceStore(), + a2a: { + taskStore, + authorizeTaskAccess: async () => true, + }, + } +} + +function recoveredOperation(operation: PaymentOperation) { + return { + protocolVersion: operation.protocolVersion, + operationId: operation.operationId, + acquiredByRequestId: operation.acquiredByRequestId, + ...(operation.executionStartedAt !== undefined + ? { executionStartedAt: operation.executionStartedAt } + : {}), + ...(operation.retentionReason ? { retentionReason: operation.retentionReason } : {}), + nonceKey: operation.nonceKey, + authorizationId: operation.authorizationId, + reservedAmount: operation.reservedAmount.toString(), + settledAmount: operation.settledAmount.toString(), + refundAmount: operation.refundAmount.toString(), + expiresAt: operation.expiresAt, + state: operation.state, + } +} + +describe('A2A task atomicity and restart recovery', () => { + it('rejects a non-atomic task store outside explicit demo mode', () => { + const legacyStore: TaskStore = { + get: async () => undefined, + put: async () => undefined, + delete: async () => undefined, + } + const config = atomicityConfig(legacyStore, { runs: 0, records: 0, settlements: 0 }) + + expect(() => createAgentGateway(config)).toThrow( + /A2A production task store must implement createIfAbsent and compareAndSet/, + ) + }) + + it('lets exactly one of two workers create, execute, and settle a task', async () => { + const taskStore = new TwoWorkerCreateBarrier() + const counters = { runs: 0, records: 0, settlements: 0 } + const first = new Hono() + const second = new Hono() + first.route('/v1/agents', createAgentGateway(atomicityConfig(taskStore, counters))) + second.route('/v1/agents', createAgentGateway(atomicityConfig(taskStore, counters))) + + const request = (app: Hono, nonce: string, text: string) => app.request( + `/v1/agents/${agent.slug}`, + { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader(nonce), + }, + body: JSON.stringify(requestBody('task-two-workers', text)), + }, + ) + + const [firstResponse, secondResponse] = await Promise.all([ + request(first, '101', 'one'), + request(second, '102', 'two'), + ]) + const bodies = await Promise.all([ + firstResponse.json(), + secondResponse.json(), + ]) as Array<{ result?: { status?: { state?: string } }; error?: { code?: number } }> + expect(bodies.filter((body) => body.result?.status?.state === 'completed')).toHaveLength(1) + expect(bodies.filter((body) => body.error?.code === -32602)).toHaveLength(1) + expect(counters.runs).toBe(1) + expect(counters.records).toBe(1) + expect(counters.settlements).toBe(1) + }) + + it('replays a crashed finalization after restart and clears the lease marker', async () => { + const taskStore = new InMemoryTaskStore() + const counters = { records: 0, settlements: 0 } + const operations = new MemoryPaymentOperations({ + onSettle: async () => { counters.settlements += 1 }, + onReclaim: async () => undefined, + }) + const paymentOperation = await operations.claimPayment( + { + commitment, + nonce: '103', + amount: '1000000000', + expiry: String(Math.floor(Date.now() / 1000) + 300), + }, + { + requestId: 'crashed-request', + agentId: agent.id, + requiredAmount: 1n, + maxOutputTokens: 1024, + executionBudget: { + maxInputTokens: 1024, + maxOutputTokens: 1024, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 1, + }, + }, + ) + const executing = await operations.beginPaymentExecution(paymentOperation) + const executionBudget = { + maxInputTokens: 1024, + maxOutputTokens: 1024, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 1, + } + const task: Task = { + kind: 'task', + id: 'task-restart', + contextId: 'ctx-restart', + status: { state: 'working', timestamp: new Date().toISOString() }, + artifacts: [artifact], + metadata: { + gatewayFinalizing: { + version: 1, + lease: { id: 'crashed-lease', expiresAt: Date.now() - 1 }, + agentSlug: agent.slug, + requestId: 'crashed-request', + consumerId: 'consumer-restart', + paymentMethod: 'x402', + startMs: Date.now() - 100, + operationId: executing.operationId, + paymentOperation: recoveredOperation(executing), + receipt, + artifact, + inputRequired: false, + maxOutputTokens: 1024, + executionBudget, + }, + }, + } + await taskStore.put(task) + + const config: GatewayConfig = { + resolveAgent: async (slug) => (slug === agent.slug ? agent : null), + getSandbox: async () => ({ async *streamPrompt() { throw new Error('restart recovery must not execute sandbox') } }), + recordUsage: async () => { counters.records += 1 }, + x402: { + operatorAddress, + chainId: 1, + paymentProtocolVersion: 2, + verifySigner: async () => true, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { + taskStore, + authorizeTaskAccess: async () => true, + }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + + const response = await app.request(`/v1/agents/${agent.slug}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/get', + params: { id: 'task-restart' }, + }), + }) + const body = await response.json() as { + result?: Task + error?: unknown + } + + expect(body.error).toBeUndefined() + expect(body.result?.status.state).toBe('completed') + expect(body.result?.artifacts).toEqual([artifact]) + expect((await taskStore.get('task-restart'))?.metadata?.gatewayFinalizing).toBeUndefined() + expect(operations.get(executing.operationId)?.state).toBe('settled') + expect(counters.records).toBe(1) + expect(counters.settlements).toBe(1) + }) + + it('keeps a settling payment recoverable after a lost settlement acknowledgement', async () => { + let settlementAttempts = 0 + let recoveryAttempts = 0 + let records = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { + settlementAttempts += 1 + throw new Error('settlement acknowledgement lost') + }, + onReclaim: async () => { recoveryAttempts += 1 }, + }) + const claimed = await operations.claimPayment( + { + commitment, + nonce: '104', + amount: '1000000000', + expiry: String(Math.floor(Date.now() / 1000) + 300), + }, + { + requestId: 'retry-request', + agentId: agent.id, + requiredAmount: 1n, + maxOutputTokens: 1024, + executionBudget: { + maxInputTokens: 1024, + maxOutputTokens: 1024, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 1, + }, + }, + ) + const executing = await operations.beginPaymentExecution(claimed) + const executionBudget = { + maxInputTokens: 1024, + maxOutputTokens: 1024, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 1, + } + const taskStore = new InMemoryTaskStore() + await taskStore.put({ + kind: 'task', + id: 'task-retry', + contextId: 'ctx-retry', + status: { state: 'working', timestamp: new Date().toISOString() }, + metadata: { + gatewayFinalizing: { + version: 1, + lease: { id: 'retry-lease', expiresAt: Date.now() - 1 }, + agentSlug: agent.slug, + requestId: 'retry-request', + consumerId: 'consumer-retry', + paymentMethod: 'x402', + startMs: Date.now() - 100, + operationId: executing.operationId, + paymentOperation: recoveredOperation(executing), + receipt, + artifact, + inputRequired: false, + maxOutputTokens: 1024, + executionBudget, + }, + }, + }) + const config: GatewayConfig = { + resolveAgent: async (slug) => (slug === agent.slug ? agent : null), + getSandbox: async () => ({ async *streamPrompt() { throw new Error('recovery must not execute sandbox') } }), + recordUsage: async () => { records += 1 }, + x402: { + operatorAddress, + chainId: 1, + paymentProtocolVersion: 2, + verifySigner: async () => true, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const getTask = () => app.request(`/v1/agents/${agent.slug}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/get', + params: { id: 'task-retry' }, + }), + }) + + const firstBody = await (await getTask()).json() as { result?: Task } + expect(firstBody.result?.status.state).toBe('working') + expect(operations.get(executing.operationId)?.state).toBe('settling') + expect(settlementAttempts).toBe(1) + const retained = await taskStore.get('task-retry') + const marker = retained?.metadata?.gatewayFinalizing as { + operationId: string + recoveryAttempts?: number + lease: { id: string; expiresAt: number } + } + expect(marker.operationId).toBe(executing.operationId) + expect(marker.recoveryAttempts).toBe(1) + + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayFinalizing: { + ...marker, + lease: { ...marker.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + const secondBody = await (await getTask()).json() as { result?: Task } + expect(secondBody.result?.status.state).toBe('completed') + expect((await taskStore.get('task-retry'))?.metadata?.gatewayFinalizing).toBeUndefined() + expect(operations.get(executing.operationId)?.state).toBe('settled') + expect(recoveryAttempts).toBe(1) + expect(records).toBe(1) + }) +}) diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 385e981..02c9203 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -258,7 +258,7 @@ describe('A2A payment ownership races', () => { createIfAbsent: (task) => innerStore.createIfAbsent(task), delete: (id) => innerStore.delete(id), async compareAndSet(expected, next) { - if (next.metadata?.gatewayFinalizing === true) { + if (next.metadata?.gatewayFinalizing) { finalizationSeen() await finalizationReleased } From b128b0776ff1540bd0eb9e6c78f49cbdf77eddd3 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 17:47:47 -0600 Subject: [PATCH 15/32] fix(a2a): recover canceled settlement failures --- docs/a2a-long-horizon.md | 2 + src/a2a/handler.ts | 89 +++++++++++++- tests/a2a-payment-races.test.ts | 206 +++++++++++++++++++++++++++++++- 3 files changed, 289 insertions(+), 8 deletions(-) diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 19d49bc..9e576f5 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -24,6 +24,8 @@ Before payment settlement, the gateway stores a recovery record in task metadata The record contains the payment operation, usage receipt, output artifact, and a five-minute lease. After a restart, the first task read after lease expiry resumes settlement with the same operation. If settlement acknowledgement fails, the gateway keeps the operation and retries after the lease expires. +Cancellation stores the recovery record before it completes the canceled task. +The canceled task therefore keeps a retryable lease when settlement acknowledgement fails. If the record is malformed or has no recoverable operation, the gateway expires the task as failed. Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index eb6071b..591018f 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -76,6 +76,8 @@ interface SerializedPaymentOperation { state: PaymentOperation['state'] } +type FinalizationState = 'completed' | 'input-required' | 'canceled' + interface FinalizationRecord { version: 1 lease: { id: string; expiresAt: number } @@ -90,6 +92,7 @@ interface FinalizationRecord { artifact: Artifact | null inputRequired: boolean inputRequiredPrompt?: string + finalState?: FinalizationState maxOutputTokens: number executionBudget: SandboxExecutionBudget recoveryAttempts?: number @@ -279,9 +282,11 @@ async function executeMessageSend( authz.maxOutputTokens, async () => { await beginPaymentExecution(authz, deps.config) - if (authz.paymentOperation) workObserved = true }, authz.paymentOperation !== undefined, + () => { + workObserved = true + }, )) { if (event.kind === 'text') { responseText += event.delta @@ -481,9 +486,11 @@ async function handleMessageStream( authz.maxOutputTokens, async () => { await beginPaymentExecution(authz, deps.config) - if (authz.paymentOperation) workObserved = true }, authz.paymentOperation !== undefined, + () => { + workObserved = true + }, )) { if (event.kind === 'text') { responseText += event.delta @@ -1190,21 +1197,74 @@ async function completeCanceledTask( deps: A2AHandlerDeps, ): Promise { if (usage) { + const current = await deps.taskStore.get(task.id) ?? task + const finalization = buildFinalizationRecord( + authz, + usage, + responseText + ? responseTextToArtifact(responseText, `${task.id}-artifact-0`) + : current.artifacts?.[0] ?? null, + false, + undefined, + 'canceled', + ) + let finalizingTask: Task | undefined + let candidate = current + for (let attempt = 0; attempt < 8; attempt += 1) { + if (isTaskFinalizing(candidate)) return candidate + if (isTerminal(candidate.status.state) && candidate.status.state !== 'canceled') return candidate + // Store the lease before settlement can move the payment to settling. + const next = withFinalizationRecord(candidate, finalization) + if (await compareAndSetTask(deps.taskStore, candidate, next)) { + finalizingTask = next + break + } + const latest = await deps.taskStore.get(task.id) + if (!latest) break + if (isTerminal(latest.status.state) && latest.status.state !== 'canceled') return latest + candidate = latest + } + if (!finalizingTask) { + throw new Error(`A2A task '${task.id}' changed before cancellation settlement`) + } + try { await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) } catch (settlementError) { + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalization.lease.id, + asError(settlementError), + ) + const recoveryTask = retained ?? finalizingTask console.error( `[a2a] canceled task settlement retained for ${authz.requestId}:`, settlementError instanceof Error ? settlementError.message : String(settlementError), ) + await maybeDeliverPush(recoveryTask, deps) + return recoveryTask } - } else { - await releaseOrRetainPayment(authz, deps, 'a2a task canceled', workObserved) + + const canceled = withStatus( + clearFinalizationMarker(finalizingTask), + 'canceled', + undefined, + responseText + ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] + : finalizingTask.artifacts, + ) + if (!await compareAndSetTask(deps.taskStore, finalizingTask, canceled)) { + return await deps.taskStore.get(task.id) ?? canceled + } + await maybeDeliverPush(canceled, deps) + return canceled } + await releaseOrRetainPayment(authz, deps, 'a2a task canceled', workObserved) const currentTask = await deps.taskStore.get(task.id) const canceledBase = currentTask?.status.state === 'canceled' ? currentTask - : withStatus(task, 'canceled') + : withStatus(currentTask ?? task, 'canceled') const canceled: Task = responseText ? { ...canceledBase, @@ -1309,6 +1369,7 @@ function buildFinalizationRecord( artifact: Artifact | null, inputRequired: boolean, inputRequiredPrompt: string | undefined, + finalState: FinalizationState = inputRequired ? 'input-required' : 'completed', ): FinalizationRecord { const operation = authz.paymentOperation return { @@ -1325,6 +1386,7 @@ function buildFinalizationRecord( artifact, inputRequired, ...(inputRequiredPrompt ? { inputRequiredPrompt } : {}), + finalState, maxOutputTokens: authz.maxOutputTokens, executionBudget: authz.executionBudget, } @@ -1490,7 +1552,22 @@ async function recoverFinalizationIfNeeded( function finalizationResultTask(task: Task, record: FinalizationRecord): Task { const cleanTask = clearFinalizationMarker(task) - if (record.inputRequired) { + const finalState = record.finalState ?? ( + task.status.state === 'canceled' + ? 'canceled' + : record.inputRequired + ? 'input-required' + : 'completed' + ) + if (finalState === 'canceled') { + return withStatus( + cleanTask, + 'canceled', + undefined, + record.artifact ? [record.artifact] : cleanTask.artifacts, + ) + } + if (finalState === 'input-required') { return withStatus( cleanTask, 'input-required', diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 02c9203..a8b0f5b 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -4,7 +4,7 @@ import { describe, expect, it } from 'vitest' import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { createAgentGateway } from '../src/middleware' import { MemoryNonceStore } from '../src/nonce-store' -import { MemoryPaymentOperations } from '../src/payment-operations' +import { MemoryPaymentOperations, type PaymentOperation } from '../src/payment-operations' import type { AgentMeta, GatewayConfig, SandboxBox } from '../src/types' const operatorAddress = '0x1111111111111111111111111111111111111111' @@ -128,6 +128,82 @@ describe('A2A payment ownership races', () => { expect((await taskStore.get('task-payment-cancel'))?.status.state).toBe('canceled') }) + it('releases payment when cancellation interrupts execution before sandbox start', async () => { + const taskStore = new InMemoryTaskStore() + let executionStarted!: () => void + const executionReady = new Promise((resolve) => { executionStarted = resolve }) + let releaseExecution!: () => void + const executionReleased = new Promise((resolve) => { releaseExecution = resolve }) + let sandboxStarted = false + + class BlockingExecutionOperations extends MemoryPaymentOperations { + override async beginPaymentExecution(operation: PaymentOperation): Promise { + const executing = await super.beginPaymentExecution(operation) + executionStarted() + await executionReleased + return executing + } + } + + const operations = new BlockingExecutionOperations() + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + streamPrompt() { + sandboxStarted = true + return (async function* () { + yield { type: 'sandbox.usage', data: { usage: usage() } } + })() + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('82') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-before-sandbox') }, + }), + }) + await executionReady + + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-before-sandbox' }, + }), + }) + expect(cancel.status).toBe(200) + releaseExecution() + + const response = await send + const body = await response.json() as { result?: { status?: { state?: string } }; error?: unknown } + + expect(body.error).toBeUndefined() + expect(body.result?.status?.state).toBe('canceled') + expect(sandboxStarted).toBe(false) + expect(operations.get(`x402:${commitment}:82`)?.state).toBe('released') + }) + it('allows only one concurrent continuation to claim and settle a task', async () => { const taskStore = new InMemoryTaskStore() await taskStore.put({ @@ -330,6 +406,7 @@ describe('A2A payment ownership races', () => { const cancellationReady = new Promise((resolve) => { cancellationStored = resolve }) let releaseCancellation!: () => void const cancellationReleased = new Promise((resolve) => { releaseCancellation = resolve }) + let cancellationBlocked = false const taskStore: TaskStore = { get: (id) => innerStore.get(id), put: (task) => innerStore.put(task), @@ -337,7 +414,8 @@ describe('A2A payment ownership races', () => { delete: (id) => innerStore.delete(id), async compareAndSet(expected, next) { const transitioned = await innerStore.compareAndSet(expected, next) - if (transitioned && next.status.state === 'canceled') { + if (transitioned && next.status.state === 'canceled' && !cancellationBlocked) { + cancellationBlocked = true cancellationStored() await cancellationReleased } @@ -548,4 +626,128 @@ describe('A2A payment ownership races', () => { const canceled = await taskStore.get('task-cancel') expect(canceled?.status.state).toBe('canceled') }) + + it('keeps a canceled task recoverable when settlement acknowledgement is lost', async () => { + const innerStore = new InMemoryTaskStore() + let finalizationSeen!: () => void + const finalizationReady = new Promise((resolve) => { finalizationSeen = resolve }) + let releaseFinalization!: () => void + const finalizationReleased = new Promise((resolve) => { releaseFinalization = resolve }) + let firstFinalization = true + const taskStore: TaskStore = { + get: (id) => innerStore.get(id), + put: (task) => innerStore.put(task), + createIfAbsent: (task) => innerStore.createIfAbsent(task), + delete: (id) => innerStore.delete(id), + async compareAndSet(expected, next) { + if (firstFinalization && next.metadata?.gatewayFinalizing) { + firstFinalization = false + finalizationSeen() + await finalizationReleased + } + return innerStore.compareAndSet(expected, next) + }, + } + let settlementAttempts = 0 + let recoveryAttempts = 0 + let records = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { + settlementAttempts += 1 + throw new Error('settlement acknowledgement lost') + }, + onReclaim: async () => { recoveryAttempts += 1 }, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => { records += 1 }, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('83') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-canceled-settlement-recovery') }, + }), + }) + await finalizationReady + + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-canceled-settlement-recovery' }, + }), + }) + expect(cancel.status).toBe(200) + releaseFinalization() + + const response = await send + const body = await response.json() as { result?: { status?: { state?: string } }; error?: unknown } + expect(body.error).toBeUndefined() + expect(body.result?.status?.state).toBe('canceled') + expect(settlementAttempts).toBe(1) + expect(operations.get(`x402:${commitment}:83`)?.state).toBe('settling') + + const retained = await taskStore.get('task-canceled-settlement-recovery') + const marker = retained?.metadata?.gatewayFinalizing as { + lease: { id: string; expiresAt: number } + recoveryAttempts?: number + } + expect(retained?.status.state).toBe('canceled') + expect(marker.lease.expiresAt).toBeGreaterThan(Date.now()) + expect(marker.recoveryAttempts).toBe(1) + + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayFinalizing: { + ...marker, + lease: { ...marker.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + const recovered = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 3, + method: 'tasks/get', + params: { id: 'task-canceled-settlement-recovery' }, + }), + }) + const recoveredBody = await recovered.json() as { + result?: { status?: { state?: string } } + } + expect(recoveredBody.result?.status?.state).toBe('canceled') + expect((await taskStore.get('task-canceled-settlement-recovery'))?.metadata?.gatewayFinalizing).toBeUndefined() + expect(operations.get(`x402:${commitment}:83`)?.state).toBe('settled') + expect(recoveryAttempts).toBe(1) + expect(records).toBe(1) + }) }) From 5310b5341322efea1033eab34a068412818012bb Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 17:48:00 -0600 Subject: [PATCH 16/32] fix(payments): require atomic replay claims --- src/dispatch.ts | 37 +++++++++++++++-------------- src/nonce-store.ts | 48 +++++++++++--------------------------- src/verify.ts | 28 ++++++++++++++-------- tests/kv-stores.test.ts | 46 ++++++++++++++++++------------------ tests/middleware.test.ts | 49 +++++++++++++++++++++++++++++++++++++-- tests/nonce-store.test.ts | 43 +++++++++++++++++++++------------- tests/verify.test.ts | 17 +++++--------- 7 files changed, 154 insertions(+), 114 deletions(-) diff --git a/src/dispatch.ts b/src/dispatch.ts index 66a29cd..8b14f89 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -574,6 +574,9 @@ export async function claimPayment( } authz.paymentOperation = operation } else if (authz.paymentMethod === 'mpp') { + if (!authz.paymentNonceKey) { + throw new Error('MPP payment has no replay identity') + } const durablePayload = durableMppPaymentPayload(authz.paymentPayload) if (durablePayload && config.x402.paymentOperations) { const context = { @@ -591,26 +594,24 @@ export async function claimPayment( authz.paymentPayload = durablePayload authz.paymentOperation = operation authz.paymentOperationAcquired = true - if (authz.paymentNonceKey) { - const claimed = await claimPaymentNonce( - state.nonceStore, - authz.paymentNonceKey, - durablePayload, - `${operation.operationId}:${context.requestId}`, - ) - if (!claimed) { - try { - await config.x402.paymentOperations.releasePayment(operation, 'shared payment nonce was already owned') - } catch (releaseError) { - console.error( - `[agent-gateway] payment release failed for ${authz.requestId}:`, - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) - } - throw new Error('payment nonce was already consumed') + const claimed = await claimPaymentNonce( + state.nonceStore, + authz.paymentNonceKey, + durablePayload, + `${operation.operationId}:${context.requestId}`, + ) + if (!claimed) { + try { + await config.x402.paymentOperations.releasePayment(operation, 'shared payment nonce was already owned') + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) } + throw new Error('payment nonce was already consumed') } - } else if (authz.paymentNonceKey) { + } else { const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) if (!claimed) throw new Error('payment nonce was already consumed') } diff --git a/src/nonce-store.ts b/src/nonce-store.ts index 7628af6..415595d 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -4,16 +4,14 @@ */ export interface NonceStore { - /** Check if nonce has been seen. Returns true if already used (reject). */ - hasSeen(nonce: string): Promise + /** Optional observation. This method never grants ownership. */ + hasSeen?(nonce: string): Promise /** * Atomically claim a nonce. An owner id makes a retry by the same payment - * operation idempotent while a legacy claim still fails closed. Version 1 - * stores can omit this method and retain their prior check-then-mark path. + * operation idempotent. Calls without an owner id use first-writer-wins + * semantics for legacy payment authorization. */ - claim?(nonce: string, ttlSeconds: number, ownerId?: string): Promise - /** Mark nonce as used. TTL = how long to remember it (seconds). */ - markSeen(nonce: string, ttlSeconds: number): Promise + claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise } /** @@ -62,11 +60,6 @@ export class MemoryNonceStore implements NonceStore { return true } - async markSeen(nonce: string, ttlSeconds: number): Promise { - this.seen.set(nonce, { expiresAt: Date.now() + ttlSeconds * 1000 }) - this.evictExpired() - } - private evictExpired() { const now = Date.now() // Evict at most every 60 seconds to avoid O(n) on every request @@ -90,7 +83,7 @@ export class MemoryNonceStore implements NonceStore { export interface KVNamespace { get(key: string, options?: { type?: 'text' | 'json' }): Promise put(key: string, value: string, options?: { expirationTtl?: number }): Promise - /** Optional atomic extension. Cloudflare KV does not provide this method. */ + /** Linearizable create-if-absent extension. Cloudflare KV does not provide it. */ putIfAbsent?(key: string, value: string, options?: { expirationTtl?: number }): Promise delete(key: string): Promise } @@ -102,7 +95,7 @@ export interface KVNamespace { * Cloudflare routes requests across multiple isolates. Without shared state, * an attacker could retry a replayed nonce against a different isolate and * have it accepted. Version 2 requires an atomic binding for payment claims. - * This store keeps the nonce only for legacy replay protection. + * This store requires an atomic binding for every payment claim. * * Usage: * const nonceStore = new KvNonceStore(env.NONCE_KV, 'x402') @@ -116,18 +109,12 @@ export class KvNonceStore implements NonceStore { ) {} async hasSeen(nonce: string): Promise { - const value = await this.kv.get(this.key(nonce)) - return value !== null + return (await this.kv.get(this.key(nonce))) !== null } async claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise { if (!this.kv.putIfAbsent) { - if (ownerId !== undefined) { - throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') - } - if (await this.hasSeen(nonce)) return false - await this.markSeen(nonce, ttlSeconds) - return true + throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') } const ttl = Math.max(ttlSeconds, 60) const key = this.key(nonce) @@ -143,29 +130,20 @@ export class KvNonceStore implements NonceStore { return (await this.kv.get(key)) === ownerId } - async markSeen(nonce: string, ttlSeconds: number): Promise { - // KV minimum TTL is 60 seconds - const ttl = Math.max(ttlSeconds, 60) - await this.kv.put(this.key(nonce), '1', { expirationTtl: ttl }) - } - private key(nonce: string): string { return `${this.prefix}:${nonce}` } } -/** Claim through a version 2 store or preserve the version 1 store contract. */ +/** Claim through the single atomic nonce-ownership contract. */ export async function claimStoredNonce( store: NonceStore, nonce: string, ttlSeconds: number, ownerId?: string, ): Promise { - if (store.claim) return store.claim(nonce, ttlSeconds, ownerId) - if (ownerId !== undefined) { - throw new Error('NonceStore.claim is required for version 2 payment ownership') + if (typeof store.claim !== 'function') { + throw new Error('NonceStore.claim is required for payment replay protection') } - if (await store.hasSeen(nonce)) return false - await store.markSeen(nonce, ttlSeconds) - return true + return store.claim(nonce, ttlSeconds, ownerId) } diff --git a/src/verify.ts b/src/verify.ts index c0854e1..4d77863 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -4,7 +4,7 @@ import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-stor /** Return the canonical opaque nonce key used by the final payment claim. */ export function mppReplayNonceKey(authHeader: string): string | undefined { const decoded = decodeMppCredential(authHeader) - return decoded ? canonicalMppNonceKey(decoded.method, decoded.payload) : undefined + return decoded ? canonicalMppNonceKey(decoded.method, decoded.payload, decoded.credential) : undefined } /** Return the decoded MPP payload for a durable payment claim. */ @@ -44,8 +44,16 @@ function decodeMppCredential(authHeader: string): DecodedMppCredential | undefin } } -function canonicalMppNonceKey(method: string, payload: Record): string | undefined { - if (payload.nonce === undefined) return undefined +function canonicalMppNonceKey( + method: string, + payload: Record, + credential: string, +): string { + if (payload.nonce === undefined) { + // Generic MPP methods may not expose a numeric nonce. The signed receipt + // itself is still the replay identity and must be claimed exactly once. + return `mpp:${method.toLowerCase()}:receipt:${Buffer.from(credential).toString('base64url')}` + } const nonce = BigInt(String(payload.nonce)).toString() const commitment = payload.commitment // BlueprinTEVM carries the same SpendAuth identity as x402. Keep one @@ -112,7 +120,7 @@ export async function verifyX402( if (amount <= 0n || minimumAmount < 0n || amount < minimumAmount) return null const nonceKey = `${String(raw.commitment).toLowerCase()}:${nonce.toString()}` - if (nonceStore && await nonceStore.hasSeen(nonceKey)) return null + if (nonceStore?.hasSeen && await nonceStore.hasSeen(nonceKey)) return null if (config.verifySigner) { const verified = await config.verifySigner(raw, { @@ -123,8 +131,8 @@ export async function verifyX402( return null } - // Check and mark only after the signature is accepted. Otherwise an - // invalid request can burn a valid payer nonce and deny the real request. + // Claim only after the signature is accepted. Otherwise invalid traffic + // can burn a valid payer nonce and deny the real request. if (nonceStore && markNonce) { const ttl = nonceTtlSeconds(expiry) if (ttl === undefined) return null @@ -191,8 +199,8 @@ export async function verifyMpp( return null } - const nonceKey = nonceStore ? canonicalMppNonceKey(method, payload) ?? null : null - if (nonceKey && await nonceStore!.hasSeen(nonceKey)) return null + const nonceKey = canonicalMppNonceKey(method, payload, decoded) + if (nonceStore?.hasSeen && await nonceStore.hasSeen(nonceKey)) return null let consumerId: string | null = null if (config.verifySigner) { @@ -211,13 +219,13 @@ export async function verifyMpp( } if (!consumerId) return null - if (nonceStore && payload.nonce !== undefined && markNonce) { + if (nonceStore && markNonce) { const expiry = payload.expiry === undefined ? BigInt(Math.floor(Date.now() / 1000) + 3600) : BigInt(String(payload.expiry)) const ttl = nonceTtlSeconds(expiry) if (ttl === undefined) return null - const claimed = await claimStoredNonce(nonceStore, nonceKey!, ttl) + const claimed = await claimStoredNonce(nonceStore, nonceKey, ttl) if (!claimed) return null } diff --git a/tests/kv-stores.test.ts b/tests/kv-stores.test.ts index 32dbbf8..08e83d5 100644 --- a/tests/kv-stores.test.ts +++ b/tests/kv-stores.test.ts @@ -24,8 +24,10 @@ class StubKV implements NonceKV, RlKV { } async putIfAbsent(key: string, value: string, options?: { expirationTtl?: number }): Promise { - if (await this.get(key) !== null) return false - await this.put(key, value, options) + const entry = this.store.get(key) + if (entry && entry.expiresAt >= this.now()) return false + const ttl = options?.expirationTtl ?? 86400 + this.store.set(key, { value, expiresAt: this.now() + ttl * 1000 }) return true } @@ -37,42 +39,43 @@ class StubKV implements NonceKV, RlKV { describe('KvNonceStore', () => { afterEach(() => vi.useRealTimers()) - it('returns false on unseen nonce', async () => { + it('claims an unseen nonce and rejects its replay', async () => { const store = new KvNonceStore(new StubKV()) - expect(await store.hasSeen('fresh')).toBe(false) + expect(await store.claim('fresh', 60)).toBe(true) + expect(await store.claim('fresh', 60)).toBe(false) }) - it('returns true after markSeen — regression: missed replay would let attackers reuse payments', async () => { + it('retains an atomic claim — regression: missed replay would let attackers reuse payments', async () => { const store = new KvNonceStore(new StubKV()) - await store.markSeen('replay', 3600) - expect(await store.hasSeen('replay')).toBe(true) + expect(await store.claim('replay', 3600)).toBe(true) + expect(await store.claim('replay', 3600)).toBe(false) }) it('enforces 60s minimum TTL — regression: KV rejects shorter TTLs so shorter expiries silently drop', async () => { const kv = new StubKV() - const putSpy = vi.spyOn(kv, 'put') + const putIfAbsentSpy = vi.spyOn(kv, 'putIfAbsent') const store = new KvNonceStore(kv) - await store.markSeen('n1', 10) // request 10 seconds - expect(putSpy).toHaveBeenCalledWith(expect.stringContaining('nonce:'), '1', { expirationTtl: 60 }) + await store.claim('n1', 10) // request 10 seconds + expect(putIfAbsentSpy).toHaveBeenCalledWith(expect.stringContaining('nonce:'), '1', { expirationTtl: 60 }) }) it('honors TTL expiry via KV eviction semantics', async () => { vi.useFakeTimers() vi.setSystemTime(new Date('2026-01-01T00:00:00Z')) const store = new KvNonceStore(new StubKV()) - await store.markSeen('fleeting', 60) - expect(await store.hasSeen('fleeting')).toBe(true) + expect(await store.claim('fleeting', 60)).toBe(true) + expect(await store.claim('fleeting', 60)).toBe(false) vi.advanceTimersByTime(61_000) - expect(await store.hasSeen('fleeting')).toBe(false) + expect(await store.claim('fleeting', 60)).toBe(true) }) it('namespaces by prefix — regression: collisions across shared KVs', async () => { const kv = new StubKV() const x402 = new KvNonceStore(kv, 'x402') const mpp = new KvNonceStore(kv, 'mpp') - await x402.markSeen('42', 300) - expect(await x402.hasSeen('42')).toBe(true) - expect(await mpp.hasSeen('42')).toBe(false) + expect(await x402.claim('42', 300)).toBe(true) + expect(await x402.claim('42', 300)).toBe(false) + expect(await mpp.claim('42', 300)).toBe(true) }) it('integrates with verifyX402 across distributed isolates — regression: same nonce accepted on a second isolate', async () => { @@ -81,9 +84,9 @@ describe('KvNonceStore', () => { const isolateA = new KvNonceStore(sharedKv) const isolateB = new KvNonceStore(sharedKv) - await isolateA.markSeen('shared-nonce', 300) - // Isolate B sees the same nonce as used — no cross-isolate bypass - expect(await isolateB.hasSeen('shared-nonce')).toBe(true) + expect(await isolateA.claim('shared-nonce', 300)).toBe(true) + // Isolate B sees the same nonce as used — no cross-isolate bypass. + expect(await isolateB.claim('shared-nonce', 300)).toBe(false) }) it('keeps same-owner claims idempotent when isolates race on the atomic insert', async () => { @@ -98,7 +101,7 @@ describe('KvNonceStore', () => { expect(await isolateA.claim('shared-operation', 300, 'operation-2')).toBe(false) }) - it('preserves legacy claims on the standard Cloudflare KV surface', async () => { + it('fails closed when the KV binding cannot make an atomic claim', async () => { const backing = new StubKV() const cloudflareKv: NonceKV = { get: backing.get.bind(backing), @@ -107,8 +110,7 @@ describe('KvNonceStore', () => { } const store = new KvNonceStore(cloudflareKv) - expect(await store.claim('legacy', 300)).toBe(true) - expect(await store.claim('legacy', 300)).toBe(false) + await expect(store.claim('legacy', 300)).rejects.toThrow('atomic putIfAbsent') await expect(store.claim('version-2', 300, 'operation-1')).rejects.toThrow('atomic putIfAbsent') }) }) diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index 0f2e38f..d437d0e 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -557,9 +557,7 @@ describe('POST /:slug/chat/completions — auth paths', () => { it('durably claims an x402-compatible MPP receipt and rejects its replay', async () => { const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) const nonceStore: NonceStore = { - hasSeen: async () => false, claim: async () => true, - markSeen: async () => undefined, } const credential = Buffer.from(JSON.stringify({ payload: { @@ -605,6 +603,53 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(operations.get('x402:0xcommitmentalice:901')?.state).toBe('settled') }) + it('claims an identical generic MPP receipt without a payload nonce only once', async () => { + let executions = 0 + const credential = Buffer.from(JSON.stringify({ receiptId: 'receipt-1' })).toString('base64url') + const { app, settlements } = buildHarness({ + getSandbox: async () => ({ + async *streamPrompt() { + executions += 1 + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'generic' } } + yield { + type: 'sandbox.usage', + data: { + usage: { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0, + budgetEnforced: true, + }, + }, + } + }, + }), + mpp: { + realm: 'agents.tangle.tools', + method: 'stripe', + verifySigner: async () => 'mpp:consumer', + }, + }) + const request = () => app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${credential}`, + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + + const responses = await Promise.all([request(), request()]) + await Promise.all(responses.map((response) => response.text())) + + expect(responses.map((response) => response.status).sort()).toEqual([200, 402]) + expect(executions).toBe(1) + expect(settlements).toHaveLength(1) + }) + it('requires complete receipts only for requests with durable payment ownership', async () => { const operations = new MemoryPaymentOperations() const { app } = buildHarness({ diff --git a/tests/nonce-store.test.ts b/tests/nonce-store.test.ts index 7862853..56963af 100644 --- a/tests/nonce-store.test.ts +++ b/tests/nonce-store.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, afterEach } from 'vitest' -import { KvNonceStore, MemoryNonceStore, nonceTtlSeconds, type KVNamespace } from '../src/nonce-store' +import { claimStoredNonce, KvNonceStore, MemoryNonceStore, nonceTtlSeconds, type NonceStore, type KVNamespace } from '../src/nonce-store' describe('nonceTtlSeconds', () => { it('covers the complete signed validity window', () => { @@ -13,34 +13,34 @@ describe('nonceTtlSeconds', () => { describe('MemoryNonceStore', () => { afterEach(() => vi.useRealTimers()) - it('returns false for unseen nonces — regression: false-positive rejection would break first-time payments', async () => { + it('claims an unseen nonce — regression: false-positive rejection would break first-time payments', async () => { const store = new MemoryNonceStore() - expect(await store.hasSeen('nonce-never-recorded')).toBe(false) + expect(await store.claim('nonce-never-recorded', 60)).toBe(true) }) - it('returns true after markSeen — regression: missed replay detection lets attackers reuse signed payments', async () => { + it('rejects a second claim — regression: missed replay detection lets attackers reuse signed payments', async () => { const store = new MemoryNonceStore() - await store.markSeen('replay-target', 60) - expect(await store.hasSeen('replay-target')).toBe(true) + expect(await store.claim('replay-target', 60)).toBe(true) + expect(await store.claim('replay-target', 60)).toBe(false) }) it('evicts nonces after their TTL expires — regression: infinite retention causes unbounded memory growth', async () => { vi.useFakeTimers() vi.setSystemTime(new Date('2026-01-01T00:00:00Z')) const store = new MemoryNonceStore() - await store.markSeen('short-lived', 60) - expect(await store.hasSeen('short-lived')).toBe(true) + expect(await store.claim('short-lived', 60)).toBe(true) + expect(await store.claim('short-lived', 60)).toBe(false) // Jump 61 seconds — past the TTL vi.advanceTimersByTime(61_000) - expect(await store.hasSeen('short-lived')).toBe(false) + expect(await store.claim('short-lived', 60)).toBe(true) }) it('isolates nonce keys — regression: key collision across commitments would let Alice replay Bob\'s nonce', async () => { const store = new MemoryNonceStore() - await store.markSeen('0xAlice:42', 60) - expect(await store.hasSeen('0xBob:42')).toBe(false) - expect(await store.hasSeen('0xAlice:42')).toBe(true) + expect(await store.claim('0xAlice:42', 60)).toBe(true) + expect(await store.claim('0xBob:42', 60)).toBe(true) + expect(await store.claim('0xAlice:42', 60)).toBe(false) }) it('background eviction removes expired entries — regression: map grows forever without cleanup', async () => { @@ -50,22 +50,33 @@ describe('MemoryNonceStore', () => { // Fill with short-lived nonces for (let i = 0; i < 100; i++) { - await store.markSeen(`n${i}`, 10) + expect(await store.claim(`n${i}`, 10)).toBe(true) } // Advance past TTL + past the 60s eviction throttle vi.advanceTimersByTime(61_000) // Recording a new nonce triggers eviction - await store.markSeen('trigger', 60) + expect(await store.claim('trigger', 60)).toBe(true) - // All old entries should report unseen now + // All old entries can be claimed again now. for (let i = 0; i < 100; i++) { - expect(await store.hasSeen(`n${i}`)).toBe(false) + expect(await store.claim(`n${i}`, 60)).toBe(true) } }) }) +describe('claimStoredNonce', () => { + it('fails closed for a store without an atomic claim method', async () => { + const legacyStore = { + hasSeen: async () => false, + markSeen: async () => undefined, + } as unknown as NonceStore + + await expect(claimStoredNonce(legacyStore, 'legacy', 60)).rejects.toThrow('payment replay') + }) +}) + describe('KvNonceStore', () => { it('allows only one mixed legacy and version 2 claim', async () => { const values = new Map() diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 92b3180..baa2780 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -98,12 +98,10 @@ describe('verifyX402', () => { it('retains a nonce until its signed expiry, beyond the old one-hour cap', async () => { const ttls: number[] = [] const nonceStore: NonceStore = { - hasSeen: async () => false, claim: async (_nonce, ttlSeconds) => { ttls.push(ttlSeconds) return true }, - markSeen: async () => undefined, } const now = Math.floor(Date.now() / 1000) const result = await verifyX402( @@ -117,16 +115,13 @@ describe('verifyX402', () => { expect(ttls[0]).toBeLessThanOrEqual(7200) }) - it('keeps version 1 custom nonce stores source-compatible', async () => { - const seen = new Set() - const nonceStore: NonceStore = { - hasSeen: async (nonce) => seen.has(nonce), - markSeen: async (nonce) => { seen.add(nonce) }, - } - const payload = buildSpendAuth({ nonce: '101' }) + it('fails closed for a custom nonce store without an atomic claim', async () => { + const nonceStore = { + hasSeen: async () => false, + markSeen: async () => undefined, + } as unknown as NonceStore - expect(await verifyX402(payload, baseConfig, nonceStore)).toBe('0xCommitmentAlice') - expect(await verifyX402(payload, baseConfig, nonceStore)).toBeNull() + expect(await verifyX402(buildSpendAuth({ nonce: '101' }), baseConfig, nonceStore)).toBeNull() }) it('isolates nonces per commitment — regression: commitment-less nonce tracking lets Alice replay Bob\'s nonce', async () => { From fd87ffd0b13a7bc182212d10441c93e5700b439f Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 18:23:11 -0600 Subject: [PATCH 17/32] fix(gateway): recover payment acknowledgements safely --- README.md | 1 + docs/a2a-long-horizon.md | 4 + src/a2a/handler.ts | 401 +++++++++++++++++++++++++++----- src/dispatch.ts | 27 ++- src/payment-operations.ts | 4 + src/types.ts | 7 +- tests/a2a-payment-races.test.ts | 218 +++++++++++++++++ tests/middleware.test.ts | 41 ++++ 8 files changed, 647 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index 5b82ee8..289015c 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,7 @@ The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox s An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. Sandbox adapters should emit a complete `sandbox.usage` receipt. Requests with a version 2 payment operation reject missing receipts; API-key and MPP requests keep the legacy visible-token estimate path. +recordUsage must atomically upsert by event.requestId; recovery may retry an event after its acknowledgement is lost. Older custom A2A task stores remain source-compatible through a process-safe fallback. Use an atomic task store for multi-worker production deployments. diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 9e576f5..77fa68a 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -26,6 +26,10 @@ After a restart, the first task read after lease expiry resumes settlement with If settlement acknowledgement fails, the gateway keeps the operation and retries after the lease expires. Cancellation stores the recovery record before it completes the canceled task. The canceled task therefore keeps a retryable lease when settlement acknowledgement fails. +If cancellation must release an unused durable operation, the gateway stores a separate release record before calling the release adapter. +An ambiguous release acknowledgement is retried after its five-minute lease, and the record is cleared only after release acknowledgement. +Usage attribution must atomically upsert by `requestId`. +The finalization record stores whether attribution was acknowledged, so recovery does not repeat an acknowledged usage event. If the record is malformed or has no recoverable operation, the gateway expires the task as failed. Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 591018f..9be5650 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -95,6 +95,19 @@ interface FinalizationRecord { finalState?: FinalizationState maxOutputTokens: number executionBudget: SandboxExecutionBudget + usageRecorded: boolean + recoveryAttempts?: number + recoveryError?: string +} + +interface PaymentReleaseRecord { + version: 1 + lease: { id: string; expiresAt: number } + agentSlug: string + requestId: string + operationId: string + paymentOperation: SerializedPaymentOperation + reason: string recoveryAttempts?: number recoveryError?: string } @@ -254,7 +267,7 @@ async function executeMessageSend( ? task : { ...task, status: { state: 'working', timestamp: nowIso() } } if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { - await releaseOwnedPayment(authz, deps, 'A2A task changed before execution started') + await releaseTaskPayment(authz, task, deps, 'A2A task changed before execution started', false) if (signal.aborted) { const canceled = await deps.taskStore.get(task.id) if (canceled?.status.state === 'canceled') { @@ -302,16 +315,17 @@ async function executeMessageSend( } } } catch (err) { - await releaseOrRetainPayment( + const releasedTask = await releaseTaskPayment( authz, + workingTask, deps, err instanceof Error ? err.message : String(err), workObserved || usage !== undefined, ) - const currentTask = await deps.taskStore.get(task.id) - const failed = currentTask && isTerminal(currentTask.status.state) + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = isTerminal(currentTask.status.state) ? currentTask - : withStatus(workingTask, 'failed') + : withStatus(currentTask, 'failed') try { await deps.taskStore.put(failed) await maybeDeliverPush(failed, deps) @@ -374,8 +388,14 @@ async function executeMessageSend( throw new Error('A2A task changed before payment settlement') } finalizationLeaseId = finalization.lease.id - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) - const settledBase = clearFinalizationMarker(finalizingTask) + let usageRecordedTask = finalizingTask + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + const settledBase = clearFinalizationMarker(usageRecordedTask) const result = inputRequiredSeen ? withStatus( settledBase, @@ -388,7 +408,7 @@ async function executeMessageSend( : withStatus(settledBase, 'completed', undefined, [ responseTextToArtifact(responseText, `${task.id}-artifact-0`), ]) - if (!await compareAndSetTask(deps.taskStore, finalizingTask, result)) { + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, result)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask && (isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required')) { return c.json(ok(req.id, currentTask)) @@ -399,8 +419,9 @@ async function executeMessageSend( await maybeDeliverPush(result, deps) return c.json(ok(req.id, result)) } catch (err) { - await releaseOrRetainPayment( + const releasedTask = await releaseTaskPayment( authz, + workingTask, deps, err instanceof Error ? err.message : String(err), workObserved || usage !== undefined, @@ -414,10 +435,10 @@ async function executeMessageSend( ) return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) } - const currentTask = await deps.taskStore.get(task.id) - const failed = currentTask && isTerminal(currentTask.status.state) + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = isTerminal(currentTask.status.state) ? currentTask - : withStatus(workingTask, 'failed') + : withStatus(currentTask, 'failed') try { await deps.taskStore.put(failed) await maybeDeliverPush(failed, deps) @@ -574,8 +595,9 @@ async function handleMessageStream( }) return } - await releaseOrRetainPayment( + await releaseTaskPayment( authz, + task, deps, 'A2A task changed before payment settlement', workObserved || usage !== undefined, @@ -584,18 +606,24 @@ async function handleMessageStream( } finalizationLeaseId = finalization.lease.id cancels.beginFinalization(task.id) - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) + let usageRecordedTask = finalizingTask + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) if (inputRequiredSeen) { const paused = withStatus( - clearFinalizationMarker(finalizingTask), + clearFinalizationMarker(usageRecordedTask), 'input-required', inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, responseText ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] : task.artifacts, ) - if (!await compareAndSetTask(deps.taskStore, finalizingTask, paused)) { + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, paused)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask) { send({ @@ -620,10 +648,10 @@ async function handleMessageStream( } // Final: persist the terminal task before emitting terminal events. - const completed = withStatus(clearFinalizationMarker(finalizingTask), 'completed', undefined, [ + const completed = withStatus(clearFinalizationMarker(usageRecordedTask), 'completed', undefined, [ responseTextToArtifact(responseText, `${task.id}-artifact-0`), ]) - if (!await compareAndSetTask(deps.taskStore, finalizingTask, completed)) { + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, completed)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask) { send({ @@ -657,8 +685,9 @@ async function handleMessageStream( }) await maybeDeliverPush(completed, deps) } catch (err) { - await releaseOrRetainPayment( + const releasedTask = await releaseTaskPayment( authz, + task, deps, err instanceof Error ? err.message : String(err), workObserved || usage !== undefined, @@ -681,10 +710,10 @@ async function handleMessageStream( } return } - const currentTask = await deps.taskStore.get(task.id) - const failed = currentTask && isTerminal(currentTask.status.state) + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = isTerminal(currentTask.status.state) ? currentTask - : withStatus(task, 'failed') + : withStatus(currentTask, 'failed') try { await deps.taskStore.put(failed) send({ @@ -756,7 +785,7 @@ async function handleTasksGet( } const accessError = await authorizeTaskAccess(c, req, storedTask, deps) if (accessError) return accessError - const task = await recoverFinalizationIfNeeded( + const task = await recoverTaskIfNeeded( storedTask, deps, c.req.param('slug') ?? '', @@ -782,7 +811,7 @@ async function handleTasksCancel( } const accessError = await authorizeTaskAccess(c, req, storedTask, deps) if (accessError) return accessError - const task = await recoverFinalizationIfNeeded( + const task = await recoverTaskIfNeeded( storedTask, deps, c.req.param('slug') ?? '', @@ -862,7 +891,7 @@ async function handleTasksResubscribe( } const accessError = await authorizeTaskAccess(c, req, storedTask, deps) if (accessError) return accessError - const task = await recoverFinalizationIfNeeded( + const task = await recoverTaskIfNeeded( storedTask, deps, c.req.param('slug') ?? '', @@ -1050,7 +1079,7 @@ async function guardMessageRequest( if (storedExisting) { const accessError = await authorizeTaskAccess(c, req, storedExisting, deps) if (accessError) return accessError - const existing = await recoverFinalizationIfNeeded(storedExisting, deps, slug) + const existing = await recoverTaskIfNeeded(storedExisting, deps, slug) if (existing.status.state !== 'input-required') { return c.json( fail( @@ -1119,10 +1148,17 @@ async function claimTaskPayment( await claimPayment(authz, deps.config, deps.state) return undefined } catch { - await releaseOwnedPayment(authz, deps, 'payment authorization failed') - const failed = paymentFailureTask ?? withStatus(task, 'failed') + const releasedTask = await releaseTaskPayment(authz, task, deps, 'payment authorization failed', false) + const releaseRecord = releasedTask.metadata?.[PAYMENT_RELEASE_METADATA_KEY] + const failed = releaseRecord !== undefined + ? isTerminal(releasedTask.status.state) + ? releasedTask + : withStatus(releasedTask, 'failed') + : paymentFailureTask ?? (isTerminal(releasedTask.status.state) + ? releasedTask + : withStatus(releasedTask, 'failed')) try { - if (await compareAndSetTask(deps.taskStore, task, failed) && isTerminal(failed.status.state)) { + if (await compareAndSetTask(deps.taskStore, releasedTask, failed) && isTerminal(failed.status.state)) { await maybeDeliverPush(failed, deps) } } catch (taskError) { @@ -1135,27 +1171,45 @@ async function claimTaskPayment( } } -async function releaseOwnedPayment( +async function releaseTaskPayment( authz: AuthorizedRequest, + task: Task, deps: A2AHandlerDeps, reason: string, -): Promise { - try { - await releasePayment(authz, deps.config, reason) - } catch (releaseError) { - console.error( - `[a2a] payment release failed for ${authz.requestId}:`, - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) + workObserved: boolean, +): Promise { + // Store the operation before release because the adapter acknowledgement can be ambiguous. + if (!workObserved && authz.paymentOperation && deps.config.x402.paymentOperations) { + let marked: Task + try { + marked = await beginPaymentReleaseRecovery(deps.taskStore, task, authz, reason) ?? task + } catch (error) { + console.error( + '[a2a] failed to persist payment release recovery for ' + authz.requestId + ':', + error instanceof Error ? error.message : String(error), + ) + return await deps.taskStore.get(task.id) ?? task + } + const record = readPaymentReleaseRecord(marked) + if (!record) return marked + try { + await releasePayment(authz, deps.config, reason) + } catch (releaseError) { + const retained = await retainPaymentReleaseForRecovery( + deps.taskStore, + task.id, + record.lease.id, + releaseError instanceof Error ? releaseError : new Error(String(releaseError)), + ) + console.error( + '[a2a] payment release retained for ' + authz.requestId + ':', + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + return retained ?? marked + } + return clearPaymentReleaseRecovery(deps.taskStore, marked, record.lease.id) } -} -async function releaseOrRetainPayment( - authz: AuthorizedRequest, - deps: A2AHandlerDeps, - reason: string, - workObserved: boolean, -): Promise { try { await releasePaymentAfterFailure(authz, deps.config, reason, workObserved) } catch (releaseError) { @@ -1164,6 +1218,7 @@ async function releaseOrRetainPayment( releaseError instanceof Error ? releaseError.message : String(releaseError), ) } + return await deps.taskStore.get(task.id) ?? task } /** Keep an owned finalization record durable when settlement acknowledgement is lost. */ @@ -1228,8 +1283,14 @@ async function completeCanceledTask( throw new Error(`A2A task '${task.id}' changed before cancellation settlement`) } + let usageRecordedTask = finalizingTask try { - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs) + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) } catch (settlementError) { const retained = await retainFinalizationForRecovery( deps.taskStore, @@ -1247,20 +1308,20 @@ async function completeCanceledTask( } const canceled = withStatus( - clearFinalizationMarker(finalizingTask), + clearFinalizationMarker(usageRecordedTask), 'canceled', undefined, responseText ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] : finalizingTask.artifacts, ) - if (!await compareAndSetTask(deps.taskStore, finalizingTask, canceled)) { + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, canceled)) { return await deps.taskStore.get(task.id) ?? canceled } await maybeDeliverPush(canceled, deps) return canceled } - await releaseOrRetainPayment(authz, deps, 'a2a task canceled', workObserved) + await releaseTaskPayment(authz, task, deps, 'a2a task canceled', workObserved) const currentTask = await deps.taskStore.get(task.id) const canceledBase = currentTask?.status.state === 'canceled' ? currentTask @@ -1279,7 +1340,9 @@ async function completeCanceledTask( // ── Helpers ─────────────────────────────────────────────────────────────── const FINALIZING_METADATA_KEY = 'gatewayFinalizing' +const PAYMENT_RELEASE_METADATA_KEY = 'gatewayPaymentRelease' const FINALIZATION_LEASE_MS = 5 * 60 * 1000 +const PAYMENT_RELEASE_LEASE_MS = 5 * 60 * 1000 async function authorizeTaskAccess( c: Context, @@ -1389,6 +1452,7 @@ function buildFinalizationRecord( finalState, maxOutputTokens: authz.maxOutputTokens, executionBudget: authz.executionBudget, + usageRecorded: false, } } @@ -1439,6 +1503,144 @@ function withFinalizationRecord(task: Task, record: FinalizationRecord): Task { } } +function markUsageRecordedRecord(task: Task): Task { + const record = readFinalizationRecord(task) + if (!record || record.usageRecorded) return task + return withFinalizationRecord(task, { ...record, usageRecorded: true }) +} + +async function markUsageRecorded(taskStore: TaskStore, task: Task): Promise { + const marked = markUsageRecordedRecord(task) + if (marked === task) return task + if (await compareAndSetTask(taskStore, task, marked)) return marked + return await taskStore.get(task.id) ?? marked +} + +function withPaymentReleaseRecord(task: Task, record: PaymentReleaseRecord): Task { + return { + ...task, + metadata: { ...(task.metadata ?? {}), [PAYMENT_RELEASE_METADATA_KEY]: record }, + } +} + +function clearPaymentReleaseRecord(task: Task): Task { + if (!task.metadata || !(PAYMENT_RELEASE_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[PAYMENT_RELEASE_METADATA_KEY] + return Object.keys(metadata).length > 0 + ? { ...task, metadata } + : (() => { + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata + })() +} + +function readPaymentReleaseRecord(task: Task): PaymentReleaseRecord | undefined { + const raw = task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const record = raw as Partial + if ( + record.version !== 1 || + !record.lease || + typeof record.lease.id !== 'string' || + record.lease.id.length === 0 || + typeof record.lease.expiresAt !== 'number' || + !Number.isFinite(record.lease.expiresAt) || + typeof record.agentSlug !== 'string' || + record.agentSlug.length === 0 || + typeof record.requestId !== 'string' || + record.requestId.length === 0 || + typeof record.operationId !== 'string' || + record.operationId.length === 0 || + !record.paymentOperation || + typeof record.paymentOperation !== 'object' || + typeof record.reason !== 'string' + ) { + return undefined + } + if (record.paymentOperation.operationId !== record.operationId) return undefined + try { + deserializePaymentOperation(record.paymentOperation) + } catch { + return undefined + } + return record as PaymentReleaseRecord +} + +function buildPaymentReleaseRecord( + authz: AuthorizedRequest, + reason: string, +): PaymentReleaseRecord | undefined { + const operation = authz.paymentOperation + if (!operation) return undefined + return { + version: 1, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, + agentSlug: authz.agent.slug, + requestId: authz.requestId, + operationId: operation.operationId, + paymentOperation: serializePaymentOperation({ ...operation, state: 'releasing' }), + reason, + } +} + +async function beginPaymentReleaseRecovery( + taskStore: TaskStore, + task: Task, + authz: AuthorizedRequest, + reason: string, +): Promise { + const record = buildPaymentReleaseRecord(authz, reason) + if (!record) return undefined + for (let attempt = 0; attempt < 8; attempt += 1) { + const current = await taskStore.get(task.id) ?? task + const existing = readPaymentReleaseRecord(current) + if (existing) { + if (existing.operationId !== record.operationId) { + throw new Error('A2A task already has a different payment release recovery') + } + return current + } + const next = withPaymentReleaseRecord(current, record) + if (await compareAndSetTask(taskStore, current, next)) return next + } + throw new Error('A2A task changed before payment release recovery was stored') +} + +async function retainPaymentReleaseForRecovery( + taskStore: TaskStore, + taskId: string, + leaseId: string, + error: Error, +): Promise { + const current = await taskStore.get(taskId) + if (!current) return undefined + const record = readPaymentReleaseRecord(current) + if (!record || record.lease.id !== leaseId) return undefined + const retry: PaymentReleaseRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, + recoveryAttempts: (record.recoveryAttempts ?? 0) + 1, + recoveryError: error.message, + } + const next = withPaymentReleaseRecord(current, retry) + if (await compareAndSetTask(taskStore, current, next)) return next + return await taskStore.get(taskId) +} + +async function clearPaymentReleaseRecovery( + taskStore: TaskStore, + task: Task, + leaseId: string, +): Promise { + const current = await taskStore.get(task.id) ?? task + const record = readPaymentReleaseRecord(current) + if (!record || record.lease.id !== leaseId) return current + const cleared = clearPaymentReleaseRecord(current) + if (await compareAndSetTask(taskStore, current, cleared)) return cleared + return await taskStore.get(task.id) ?? cleared +} + function readFinalizationRecord(task: Task): FinalizationRecord | undefined { const raw = task.metadata?.[FINALIZING_METADATA_KEY] if (!raw || typeof raw !== 'object') return undefined @@ -1459,6 +1661,58 @@ function isTaskFinalizing(task: Task): boolean { return marker === true || (typeof marker === 'object' && marker !== null) } +async function recoverPaymentReleaseIfNeeded( + task: Task, + deps: A2AHandlerDeps, +): Promise { + const raw = task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] + if (raw === undefined) return task + const record = readPaymentReleaseRecord(task) + if (!record) { + return expirePaymentRelease( + task, + deps, + new Error('A2A payment release recovery record is missing'), + ) + } + if (record.lease.expiresAt > Date.now()) return task + + const renewed: PaymentReleaseRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, + } + const leasedTask = withPaymentReleaseRecord(task, renewed) + if (!await compareAndSetTask(deps.taskStore, task, leasedTask)) { + return await deps.taskStore.get(task.id) ?? task + } + + try { + if (!deps.config.x402.paymentOperations) { + throw new Error('A2A payment release recovery is not configured') + } + const operation = deserializePaymentOperation(renewed.paymentOperation) + await deps.config.x402.paymentOperations.releasePayment(operation, renewed.reason) + const recovered = clearPaymentReleaseRecord(leasedTask) + if (!await compareAndSetTask(deps.taskStore, leasedTask, recovered)) { + return await deps.taskStore.get(task.id) ?? recovered + } + return recovered + } catch (error) { + const recoveryError = error instanceof Error ? error : new Error(String(error)) + const retained = await retainPaymentReleaseForRecovery( + deps.taskStore, + task.id, + renewed.lease.id, + recoveryError, + ) + console.error( + '[a2a] payment release recovery failed for ' + task.id + ':', + recoveryError.message, + ) + return retained ?? leasedTask + } +} + async function recoverFinalizationIfNeeded( task: Task, deps: A2AHandlerDeps, @@ -1523,10 +1777,24 @@ async function recoverFinalizationIfNeeded( ? { paymentOperation, paymentOperationAcquired: true } : {}), } - await settleAndRecord(agent, authz, renewed.receipt, deps.config, deps.state.obs) + let usageRecordedTask = leasedTask + await settleAndRecord( + agent, + authz, + renewed.receipt, + deps.config, + deps.state.obs, + { + usageAlreadyRecorded: renewed.usageRecorded === true, + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }, + ) - const recovered = finalizationResultTask(leasedTask, renewed) - if (!await compareAndSetTask(deps.taskStore, leasedTask, recovered)) { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + const recovered = finalizationResultTask(usageRecordedTask, renewed) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, recovered)) { return await deps.taskStore.get(task.id) ?? recovered } await maybeDeliverPush(recovered, deps) @@ -1550,6 +1818,15 @@ async function recoverFinalizationIfNeeded( } } +async function recoverTaskIfNeeded( + task: Task, + deps: A2AHandlerDeps, + requestedAgentSlug: string, +): Promise { + const released = await recoverPaymentReleaseIfNeeded(task, deps) + return recoverFinalizationIfNeeded(released, deps, requestedAgentSlug) +} + function finalizationResultTask(task: Task, record: FinalizationRecord): Task { const cleanTask = clearFinalizationMarker(task) const finalState = record.finalState ?? ( @@ -1607,6 +1884,26 @@ async function expireFinalization( return await deps.taskStore.get(task.id) ?? failed } +async function expirePaymentRelease( + task: Task, + deps: A2AHandlerDeps, + error: Error, +): Promise { + const cleanTask = clearPaymentReleaseRecord(task) + const failed: Task = { + ...withStatus(cleanTask, 'failed'), + metadata: { + ...(cleanTask.metadata ?? {}), + gatewayPaymentReleaseRecovery: { error: error.message }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await maybeDeliverPush(failed, deps) + return failed + } + return await deps.taskStore.get(task.id) ?? failed +} + function clearFinalizationMarker(task: Task): Task { if (!task.metadata || !(FINALIZING_METADATA_KEY in task.metadata)) return task const metadata = { ...task.metadata } diff --git a/src/dispatch.ts b/src/dispatch.ts index 8b14f89..76fe242 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -66,6 +66,7 @@ export interface AuthorizedRequest { requiredPaymentAmount: bigint paymentPayload: Record | null paymentNonceKey?: string + mppMethod?: string paymentOperation?: PaymentOperation paymentOperationAcquired?: boolean } @@ -243,6 +244,7 @@ export async function authenticateAndGuard( let keyInfo: ApiKeyInfo | null = null let x402Payload: Record | null = null let paymentNonceKey: string | undefined + let mppMethod: string | undefined if (spendAuthHeader) { const signer = await verifyX402( @@ -314,6 +316,7 @@ export async function authenticateAndGuard( } consumerId = signer paymentMethod = 'mpp' + mppMethod = authHeader.match(/^Payment\s+(\S+)\s+/i)?.[1]?.toLowerCase() x402Payload = mppPaymentPayload(authHeader) ?? null paymentNonceKey = mppReplayNonceKey(authHeader) } else if (authHeader.startsWith('Bearer ')) { @@ -494,6 +497,7 @@ export async function authenticateAndGuard( requiredPaymentAmount, paymentPayload: x402Payload, paymentNonceKey, + mppMethod, } } @@ -577,8 +581,11 @@ export async function claimPayment( if (!authz.paymentNonceKey) { throw new Error('MPP payment has no replay identity') } + const mppMethod = (authz.mppMethod ?? config.mpp?.method ?? 'blueprintevm').toLowerCase() const durablePayload = durableMppPaymentPayload(authz.paymentPayload) - if (durablePayload && config.x402.paymentOperations) { + // Only BlueprinTEVM carries x402 authorization fields; other MPP methods + // must stay on their verifier and method-specific settlement path. + if (mppMethod === 'blueprintevm' && durablePayload && config.x402.paymentOperations) { const context = { requestId: authz.requestId, agentId: authz.agent.id, @@ -967,12 +974,20 @@ async function closeSandboxIterator(iterator: AsyncIterator) * formats call this once their stream has drained, so settlement happens * exactly once per request regardless of protocol. */ +export interface SettleAndRecordOptions { + /** Skip attribution after a durable finalization marker confirms it ran. */ + usageAlreadyRecorded?: boolean + /** Persist the caller's recovery marker after attribution succeeds. */ + onUsageRecorded?: () => Promise +} + export async function settleAndRecord( agent: AgentMeta, authz: AuthorizedRequest, usage: SandboxUsageReceipt, config: GatewayConfig, obs: GatewayObserver | undefined, + options: SettleAndRecordOptions = {}, ): Promise { const tokenCost = ( usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens @@ -1019,11 +1034,17 @@ export async function settleAndRecord( ) // Durable settlement happens first. If attribution storage is // unavailable, recovery must never refund delivered work. - await config.recordUsage(usageEvent) + if (!options.usageAlreadyRecorded) { + await config.recordUsage(usageEvent) + await options.onUsageRecorded?.() + } } else { // Legacy adapters retain attribution-before-charge because their // settlement callback may resolve that usage row. - await config.recordUsage(usageEvent) + if (!options.usageAlreadyRecorded) { + await config.recordUsage(usageEvent) + await options.onUsageRecorded?.() + } if (config.settlePayment) { await config.settlePayment( { diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 9a52305..7c8d5bf 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -71,6 +71,10 @@ export interface PaymentOperations { operation: PaymentOperation, input: PaymentSettlementInput, ): Promise + /** + * Release an unused authorization. + * Repeated calls must recover an ambiguous acknowledgement by operationId. + */ releasePayment(operation: PaymentOperation, reason: string): Promise reclaimPayment(operationId: string): Promise } diff --git a/src/types.ts b/src/types.ts index 8af750a..a36b8a4 100644 --- a/src/types.ts +++ b/src/types.ts @@ -273,7 +273,12 @@ export interface GatewayConfig { consumer: { method: PaymentMethod; consumerId: string; keyId?: string; requestId: string }, ) => Promise<{ allow: true } | { allow: false; reason: string; code: string }> - /** Record a usage event after request completes. */ + /** + * Record a usage event after request completes. + * The implementation must atomically upsert by requestId and return + * success when the row already exists. Recovery may retry after an + * acknowledgement is lost, so one request ID must produce one usage row. + */ recordUsage: (event: GatewayUsageEvent) => Promise /** x402 payment configuration */ diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index a8b0f5b..f7e9662 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -750,4 +750,222 @@ describe('A2A payment ownership races', () => { expect(recoveryAttempts).toBe(1) expect(records).toBe(1) }) + + it('persists and recovers an ambiguous release acknowledgement', async () => { + const taskStore = new InMemoryTaskStore() + let executionStarted!: () => void + const executionReady = new Promise((resolve) => { executionStarted = resolve }) + let releaseExecution!: () => void + const executionReleased = new Promise((resolve) => { releaseExecution = resolve }) + let sandboxStarted = false + let releaseCalls = 0 + let recoveryCalls = 0 + + class BlockingExecutionOperations extends MemoryPaymentOperations { + override async beginPaymentExecution(operation: PaymentOperation): Promise { + const executing = await super.beginPaymentExecution(operation) + executionStarted() + await executionReleased + return executing + } + } + + const operations = new BlockingExecutionOperations({ + onRelease: async () => { + releaseCalls += 1 + throw new Error('release acknowledgement lost after refund') + }, + onReclaim: async () => { recoveryCalls += 1 }, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + streamPrompt() { + sandboxStarted = true + return (async function* () { + yield { type: 'sandbox.usage', data: { usage: usage() } } + })() + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('84') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-release-recovery') }, + }), + }) + await executionReady + + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-release-recovery' }, + }), + }) + expect(cancel.status).toBe(200) + releaseExecution() + + const response = await send + const body = await response.json() as { result?: { status?: { state?: string } }; error?: unknown } + expect(body.error).toBeUndefined() + expect(body.result?.status?.state).toBe('canceled') + expect(sandboxStarted).toBe(false) + expect(operations.get(`x402:${commitment}:84`)?.state).toBe('releasing') + expect(releaseCalls).toBe(1) + + const retained = await taskStore.get('task-release-recovery') + const marker = retained?.metadata?.gatewayPaymentRelease as { + lease: { id: string; expiresAt: number } + operationId: string + recoveryAttempts?: number + } + expect(retained?.status.state).toBe('canceled') + expect(marker.operationId).toBe(`x402:${commitment}:84`) + expect(marker.lease.expiresAt).toBeGreaterThan(Date.now()) + expect(marker.recoveryAttempts).toBe(1) + + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayPaymentRelease: { + ...marker, + lease: { ...marker.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + const recovered = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 3, + method: 'tasks/get', + params: { id: 'task-release-recovery' }, + }), + }) + const recoveredBody = await recovered.json() as { + result?: { status?: { state?: string } } + } + expect(recoveredBody.result?.status?.state).toBe('canceled') + expect((await taskStore.get('task-release-recovery'))?.metadata?.gatewayPaymentRelease).toBeUndefined() + expect(operations.get(`x402:${commitment}:84`)?.state).toBe('released') + expect(releaseCalls).toBe(1) + expect(recoveryCalls).toBe(1) + }) + + it('records usage once when an inserted acknowledgement is lost', async () => { + const taskStore = new InMemoryTaskStore() + const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) + const usageRequestIds = new Set() + let recordCalls = 0 + let rows = 0 + let firstCall = true + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid output' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async (event) => { + recordCalls += 1 + if (!usageRequestIds.has(event.requestId)) { + usageRequestIds.add(event.requestId) + rows += 1 + } + if (firstCall) { + firstCall = false + throw new Error('usage acknowledgement lost after insert') + } + }, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + + const send = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('85') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-usage-recovery') }, + }), + }) + const body = await send.json() as { error?: { code?: number } } + expect(body.error?.code).toBe(-32603) + expect(recordCalls).toBe(1) + expect(rows).toBe(1) + expect(operations.get(`x402:${commitment}:85`)?.state).toBe('settled') + + const retained = await taskStore.get('task-usage-recovery') + const marker = retained?.metadata?.gatewayFinalizing as { + lease: { id: string; expiresAt: number } + usageRecorded: boolean + } + expect(retained?.status.state).toBe('working') + expect(marker.usageRecorded).toBe(false) + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayFinalizing: { + ...marker, + lease: { ...marker.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + + const recovered = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/get', + params: { id: 'task-usage-recovery' }, + }), + }) + const recoveredBody = await recovered.json() as { + result?: { status?: { state?: string } } + } + expect(recoveredBody.result?.status?.state).toBe('completed') + expect(recordCalls).toBe(2) + expect(rows).toBe(1) + expect(usageRequestIds.size).toBe(1) + expect((await taskStore.get('task-usage-recovery'))?.metadata?.gatewayFinalizing).toBeUndefined() + expect(operations.get(`x402:${commitment}:85`)?.state).toBe('settled') + }) }) diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index d437d0e..a6a3568 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -603,6 +603,47 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(operations.get('x402:0xcommitmentalice:901')?.state).toBe('settled') }) + it('keeps a generic Stripe MPP receipt on its method-specific path', async () => { + const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) + const credential = Buffer.from(JSON.stringify({ + from: 'stripe-customer', + amount: fundedRequestAmount, + nonce: '902', + expiry: String(Math.floor(Date.now() / 1000) + 600), + receiptId: 'stripe-receipt-902', + })).toString('base64url') + const { app, settlements, usage } = buildHarness({ + mpp: { + realm: 'agents.tangle.tools', + verifySigner: async () => 'mpp:stripe-customer', + }, + x402: { + operatorAddress, + chainId: 3799, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + }) + const response = await app.request('/v1/agents/test-agent/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${credential}`, + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'hi' }] }), + }) + const streamed = await readSse(response) + + expect(response.status).toBe(200) + expect(streamed.combinedText).toBe('Hello, world!') + expect(streamed.done).toBe(true) + expect(operations.get('x402:stripe-customer:902')).toBeUndefined() + expect(settlements).toHaveLength(1) + expect(settlements[0]?.method).toBe('mpp') + expect(usage).toHaveLength(1) + }) + it('claims an identical generic MPP receipt without a payload nonce only once', async () => { let executions = 0 const credential = Buffer.from(JSON.stringify({ receiptId: 'receipt-1' })).toString('base64url') From 187194d94f3c214b6b807877567a168c06fcccdb Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 18:28:50 -0600 Subject: [PATCH 18/32] fix(payments): retain failed nonce release ownership --- src/dispatch.ts | 4 +- tests/a2a-payment-races.test.ts | 99 +++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/src/dispatch.ts b/src/dispatch.ts index 76fe242..35aa294 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -556,6 +556,9 @@ export async function claimPayment( if (!authz.paymentOperationAcquired) { throw new Error('payment operation was already claimed') } + // Attach owned state before the shared nonce claim. If that claim fails, + // the caller can persist release recovery after an ambiguous refund. + authz.paymentOperation = operation } if (operation && authz.paymentNonceKey) { const claimed = await claimPaymentNonce( @@ -576,7 +579,6 @@ export async function claimPayment( throw new Error('payment nonce was already consumed') } } - authz.paymentOperation = operation } else if (authz.paymentMethod === 'mpp') { if (!authz.paymentNonceKey) { throw new Error('MPP payment has no replay identity') diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index f7e9662..1eef97a 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -875,6 +875,105 @@ describe('A2A payment ownership races', () => { expect(recoveryCalls).toBe(1) }) + it('recovers an ambiguous release after shared nonce ownership fails', async () => { + const taskStore = new InMemoryTaskStore() + const nonceStore = { + hasSeen: async () => false, + claim: async () => false, + } + let sandboxStarted = false + let releaseCalls = 0 + let recoveryCalls = 0 + let failRecovery = true + const operations = new MemoryPaymentOperations({ + onRelease: async () => { + releaseCalls += 1 + throw new Error('release acknowledgement lost after refund') + }, + onReclaim: async () => { + recoveryCalls += 1 + if (failRecovery) { + failRecovery = false + throw new Error('release recovery acknowledgement also lost') + } + }, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + streamPrompt() { + sandboxStarted = true + return (async function* () { + yield { type: 'sandbox.usage', data: { usage: usage() } } + })() + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + + const response = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('86') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-nonce-release-recovery') }, + }), + }) + const body = await response.json() as { error?: { code?: number } } + expect(body.error?.code).toBe(-32603) + expect(sandboxStarted).toBe(false) + expect(operations.get(`x402:${commitment}:86`)?.state).toBe('releasing') + expect(releaseCalls).toBe(1) + expect(recoveryCalls).toBe(1) + + const retained = await taskStore.get('task-nonce-release-recovery') + const marker = retained?.metadata?.gatewayPaymentRelease as { + lease: { id: string; expiresAt: number } + operationId: string + } + expect(retained?.status.state).toBe('failed') + expect(marker.operationId).toBe(`x402:${commitment}:86`) + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayPaymentRelease: { + ...marker, + lease: { ...marker.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + + const recovered = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/get', + params: { id: 'task-nonce-release-recovery' }, + }), + }) + expect(recovered.status).toBe(200) + expect((await taskStore.get('task-nonce-release-recovery'))?.metadata?.gatewayPaymentRelease) + .toBeUndefined() + expect(operations.get(`x402:${commitment}:86`)?.state).toBe('released') + expect(recoveryCalls).toBe(2) + }) + it('records usage once when an inserted acknowledgement is lost', async () => { const taskStore = new InMemoryTaskStore() const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) From e96fb7d1cb514dd12b2a874d91add8b628c39a66 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 22:13:54 -0600 Subject: [PATCH 19/32] refactor(gateway): consolidate payment recovery ownership --- .agent/deep-clean-baseline.json | 91 ++++ .agent/skill-runs.jsonl | 2 + README.md | 40 +- docs/a2a-long-horizon.md | 17 +- src/a2a/handler.ts | 422 +++++++++++---- src/a2a/task-store-sql.ts | 26 +- src/a2a/task-store.ts | 17 +- src/dispatch.ts | 443 ++++++++++++++-- src/index.ts | 32 ++ src/middleware.ts | 100 +++- src/mpp-payment.ts | 117 +++++ src/payment-operations.ts | 7 + src/payment-recovery-sql.ts | 108 ++++ src/payment-recovery-worker.ts | 457 +++++++++++++++++ src/payment-recovery.ts | 330 ++++++++++++ src/types.ts | 25 +- src/verify.ts | 111 ++-- tests/a2a-atomicity.test.ts | 3 + tests/a2a-durability.test.ts | 36 +- tests/a2a-payment-races.test.ts | 439 ++++++++++++++-- tests/a2a.test.ts | 5 +- tests/middleware.test.ts | 82 ++- tests/payment-operations.test.ts | 16 +- tests/payment-recovery-sql.test.ts | 196 +++++++ tests/payment-recovery.test.ts | 789 +++++++++++++++++++++++++++++ tests/protocol-guards.test.ts | 13 +- tests/verify.test.ts | 24 +- 27 files changed, 3653 insertions(+), 295 deletions(-) create mode 100644 .agent/deep-clean-baseline.json create mode 100644 .agent/skill-runs.jsonl create mode 100644 src/mpp-payment.ts create mode 100644 src/payment-recovery-sql.ts create mode 100644 src/payment-recovery-worker.ts create mode 100644 src/payment-recovery.ts create mode 100644 tests/payment-recovery-sql.test.ts create mode 100644 tests/payment-recovery.test.ts diff --git a/.agent/deep-clean-baseline.json b/.agent/deep-clean-baseline.json new file mode 100644 index 0000000..ed1fb66 --- /dev/null +++ b/.agent/deep-clean-baseline.json @@ -0,0 +1,91 @@ +{ + "schema_version": 1, + "recorded_at_utc": "2026-08-15T03:39:51.140Z", + "scope": { + "repository": "/home/drew/code/agent-gateway-x402-priced-reservation", + "base_ref": "origin/fix/x402-priced-reservation", + "head": "187194d94f3c214b6b807877567a168c06fcccdb", + "local_commits_ahead": 7, + "working_tree_paths": 26, + "working_tree_status": [ + " M README.md", + " M docs/a2a-long-horizon.md", + " M src/a2a/handler.ts", + " M src/a2a/task-store-sql.ts", + " M src/a2a/task-store.ts", + " M src/dispatch.ts", + " M src/index.ts", + " M src/middleware.ts", + " M src/payment-operations.ts", + " M src/types.ts", + " M src/verify.ts", + " M tests/a2a-atomicity.test.ts", + " M tests/a2a-durability.test.ts", + " M tests/a2a-payment-races.test.ts", + " M tests/a2a.test.ts", + " M tests/middleware.test.ts", + " M tests/payment-operations.test.ts", + " M tests/protocol-guards.test.ts", + " M tests/verify.test.ts", + "?? .agent/", + "?? src/mpp-payment.ts", + "?? src/payment-recovery-sql.ts", + "?? src/payment-recovery-worker.ts", + "?? src/payment-recovery.ts", + "?? tests/payment-recovery-sql.test.ts", + "?? tests/payment-recovery.test.ts" + ], + "committed_diff_stat": " .github/workflows/publish.yml | 1 +\n README.md | 1 +\n docs/a2a-long-horizon.md | 17 +-\n src/a2a/handler.ts | 1012 +++++++++++++++++++++++++++++++++-----\n src/a2a/task-store-sql.ts | 5 +-\n src/a2a/task-store.ts | 4 +-\n src/dispatch.ts | 135 ++++-\n src/middleware.ts | 4 +-\n src/nonce-store.ts | 52 +-\n src/payment-operations.ts | 8 +-\n src/types.ts | 7 +-\n src/verify.ts | 92 ++--\n tests/a2a-atomicity.test.ts | 433 ++++++++++++++++\n tests/a2a-payment-races.test.ts | 527 +++++++++++++++++++-\n tests/kv-stores.test.ts | 46 +-\n tests/middleware.test.ts | 187 ++++++-\n tests/nonce-store.test.ts | 74 ++-\n tests/payment-operations.test.ts | 142 +++++-\n tests/protocol-guards.test.ts | 11 +-\n tests/verify.test.ts | 46 +-\n 20 files changed, 2535 insertions(+), 269 deletions(-)\n", + "uncommitted_diff_stat": " README.md | 40 +++-\n docs/a2a-long-horizon.md | 17 +-\n src/a2a/handler.ts | 355 +++++++++++++++++++++++++++----\n src/a2a/task-store-sql.ts | 26 ++-\n src/a2a/task-store.ts | 17 +-\n src/dispatch.ts | 439 ++++++++++++++++++++++++++++++++++++---\n src/index.ts | 33 ++-\n src/middleware.ts | 77 ++++++-\n src/payment-operations.ts | 7 +\n src/types.ts | 25 ++-\n src/verify.ts | 119 +++++++----\n tests/a2a-atomicity.test.ts | 3 +\n tests/a2a-durability.test.ts | 36 +++-\n tests/a2a-payment-races.test.ts | 439 ++++++++++++++++++++++++++++++++++-----\n tests/a2a.test.ts | 5 +-\n tests/middleware.test.ts | 78 ++++++-\n tests/payment-operations.test.ts | 10 +-\n tests/protocol-guards.test.ts | 13 +-\n tests/verify.test.ts | 24 ++-\n 19 files changed, 1523 insertions(+), 240 deletions(-)\n" + }, + "config_read": { + "tsconfig_strict": true, + "tsconfig_strict_flags_off": [], + "knip_config_files": [], + "top_10_churn": " 14 src/middleware.ts\n 14 src/dispatch.ts\n 14 src/a2a/handler.ts\n 14 package.json\n 13 src/types.ts\n 11 tests/middleware.test.ts\n 11 tests/a2a-payment-races.test.ts\n 9 src/verify.ts\n 9 README.md\n 8 tests/verify.test.ts\n" + }, + "metrics": { + "type_errors": { + "command": "npx tsc --noEmit", + "exit_code": 0, + "value": 0, + "raw_output": "(node:1416767) Warning: The 'NO_COLOR' env is ignored due to the 'FORCE_COLOR' env being set.\n(Use `node --trace-warnings ...` to show where the warning was created)\n" + }, + "knip_issues": { + "command": "npx knip --reporter json", + "exit_code": 1, + "value": 5, + "raw_output": "(node:1420486) Warning: The 'NO_COLOR' env is ignored due to the 'FORCE_COLOR' env being set.\n(Use `node --trace-warnings ...` to show where the warning was created)\n(node:1422626) Warning: The 'NO_COLOR' env is ignored due to the 'FORCE_COLOR' env being set.\n(Use `node --trace-warnings ...` to show where the warning was created)\n{\"issues\":[{\"file\":\"coverage/block-navigation.js\",\"binaries\":[],\"catalog\":[],\"catalogReferences\":[],\"dependencies\":[],\"devDependencies\":[],\"duplicates\":[],\"enumMembers\":[],\"exports\":[],\"files\":[{\"name\":\"coverage/block-navigation.js\"}],\"namespaceMembers\":[],\"optionalPeerDependencies\":[],\"types\":[],\"unlisted\":[],\"unresolved\":[]},{\"file\":\"coverage/prettify.js\",\"binaries\":[],\"catalog\":[],\"catalogReferences\":[],\"dependencies\":[],\"devDependencies\":[],\"duplicates\":[],\"enumMembers\":[],\"exports\":[],\"files\":[{\"name\":\"coverage/prettify.js\"}],\"namespaceMembers\":[],\"optionalPeerDependencies\":[],\"types\":[],\"unlisted\":[],\"unresolved\":[]},{\"file\":\"coverage/sorter.js\",\"binaries\":[],\"catalog\":[],\"catalogReferences\":[],\"dependencies\":[],\"devDependencies\":[],\"duplicates\":[],\"enumMembers\":[],\"exports\":[],\"files\":[{\"name\":\"coverage/sorter.js\"}],\"namespaceMembers\":[],\"optionalPeerDependencies\":[],\"types\":[],\"unlisted\":[],\"unresolved\":[]},{\"file\":\"src/dispatch.ts\",\"binaries\":[],\"catalog\":[],\"catalogReferences\":[],\"dependencies\":[],\"devDependencies\":[],\"duplicates\":[],\"enumMembers\":[],\"exports\":[{\"name\":\"dispatchSandboxStream\",\"line\":1061,\"col\":24,\"pos\":36966},{\"name\":\"estimateTokens\",\"line\":1625,\"col\":17,\"pos\":57245},{\"name\":\"estimateBillableInputTokens\",\"line\":1630,\"col\":17,\"pos\":57418}],\"files\":[],\"namespaceMembers\":[],\"optionalPeerDependencies\":[],\"types\":[],\"unlisted\":[],\"unresolved\":[]},{\"file\":\"src/payment-recovery.ts\",\"binaries\":[],\"catalog\":[],\"catalogReferences\":[],\"dependencies\":[],\"devDependencies\":[],\"duplicates\":[],\"enumMembers\":[],\"exports\":[{\"name\":\"DEFAULT_STALE_REQUEST_MS\",\"line\":111,\"col\":14,\"pos\":2998},{\"name\":\"DEFAULT_RECEIPT_TIMEOUT_MS\",\"line\":112,\"col\":14,\"pos\":3045},{\"name\":\"DEFAULT_RECOVERY_RETRY_MS\",\"line\":113,\"col\":14,\"pos\":3098},{\"name\":\"DEFAULT_RECOVERY_LEASE_MS\",\"line\":114,\"col\":14,\"pos\":3146}],\"files\":[],\"namespaceMembers\":[],\"optionalPeerDependencies\":[],\"types\":[],\"unlisted\":[],\"unresolved\":[]}]}\n" + }, + "cycles": { + "command": "npx madge --circular src/", + "exit_code": 0, + "value": 0, + "raw_output": "(node:1423767) Warning: The 'NO_COLOR' env is ignored due to the 'FORCE_COLOR' env being set.\n(Use `node --trace-warnings ...` to show where the warning was created)\n(node:1427702) Warning: The 'NO_COLOR' env is ignored due to the 'FORCE_COLOR' env being set.\n(Use `node --trace-warnings ...` to show where the warning was created)\n- Finding files\nProcessed \u001b[1m0\u001b[22m files \u001b[2m(3s)\u001b[22m \n\n\u001b[32m✔\u001b[39m \u001b[1mNo circular dependency found!\u001b[22m\n\n" + }, + "duplication": { + "command": "npx jscpd --min-lines 6 --min-tokens 50 src/", + "exit_code": 0, + "clones": 21, + "total_lines": 8741, + "duplicated_lines": 262, + "duplicated_line_percent": "3.00%", + "total_tokens": 44037, + "duplicated_tokens": 1759, + "duplicated_token_percent": "3.99%", + "raw_output": "(node:1447741) Warning: The 'NO_COLOR' env is ignored due to the 'FORCE_COLOR' env being set.\n(Use `node --trace-warnings ...` to show where the warning was created)\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [246:33 - 257:47] (12 lines, 75 tokens)\n a2a/handler.ts [463:35 - 475:47]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [301:7 - 312:28] (12 lines, 53 tokens)\n a2a/handler.ts [538:22 - 549:32]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [331:5 - 347:40] (17 lines, 100 tokens)\n a2a/handler.ts [444:5 - 457:63]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [373:5 - 391:10] (19 lines, 85 tokens)\n a2a/handler.ts [598:9 - 619:14]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [397:5 - 403:83] (7 lines, 59 tokens)\n a2a/handler.ts [640:9 - 646:87]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [657:81 - 671:18] (15 lines, 76 tokens)\n a2a/handler.ts [685:82 - 699:16]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [808:30 - 829:4] (22 lines, 148 tokens)\n a2a/handler.ts [914:38 - 935:4]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [811:23 - 829:4] (19 lines, 136 tokens)\n a2a/handler.ts [837:26 - 855:4]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [962:29 - 970:31] (9 lines, 54 tokens)\n a2a/handler.ts [997:29 - 1005:31]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [962:29 - 970:31] (9 lines, 54 tokens)\n a2a/handler.ts [1030:30 - 1038:31]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [962:29 - 985:31] (24 lines, 160 tokens)\n a2a/handler.ts [1052:32 - 1071:38]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [1009:5 - 1017:8] (9 lines, 66 tokens)\n a2a/handler.ts [1040:85 - 1048:8]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [1661:63 - 1678:2] (18 lines, 116 tokens)\n payment-recovery.ts [247:1 - 264:2]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32ma2a/handler.ts\u001b[39m\u001b[22m [1683:71 - 1699:2] (17 lines, 111 tokens)\n payment-recovery.ts [271:67 - 287:2]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mdispatch.ts\u001b[39m\u001b[22m [603:29 - 617:6] (15 lines, 68 tokens)\n dispatch.ts [644:23 - 658:6]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mdispatch.ts\u001b[39m\u001b[22m [973:17 - 980:10] (8 lines, 54 tokens)\n dispatch.ts [1438:24 - 1445:18]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mdispatch.ts\u001b[39m\u001b[22m [1417:42 - 1425:16] (9 lines, 58 tokens)\n dispatch.ts [1437:27 - 1444:15]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mnonce-store.ts\u001b[39m\u001b[22m [69:56 - 85:87] (17 lines, 58 tokens)\n rate-limit.ts [58:55 - 70:87]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mpayment-recovery-worker.ts\u001b[39m\u001b[22m [449:73 - 457:2] (9 lines, 61 tokens)\n payment-recovery.ts [235:1 - 243:2]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mpublish.ts\u001b[39m\u001b[22m [56:13 - 64:10] (9 lines, 85 tokens)\n publish.ts [69:14 - 77:10]\n\u001b[1mClone found (typescript)\u001b[22m\n - \u001b[1m\u001b[32mpublish.ts\u001b[39m\u001b[22m [56:36 - 62:58] (7 lines, 82 tokens)\n publish.ts [103:39 - 109:58]\n\u001b[90m┌────────────┬────────────────┬─────────────┬──────────────┬──────────────┬──────────────────┬───────────────────┐\u001b[39m\n\u001b[90m│\u001b[39m\u001b[31m Format \u001b[39m\u001b[90m│\u001b[39m\u001b[31m Files analyzed \u001b[39m\u001b[90m│\u001b[39m\u001b[31m Total lines \u001b[39m\u001b[90m│\u001b[39m\u001b[31m Total tokens \u001b[39m\u001b[90m│\u001b[39m\u001b[31m Clones found \u001b[39m\u001b[90m│\u001b[39m\u001b[31m Duplicated lines \u001b[39m\u001b[90m│\u001b[39m\u001b[31m Duplicated tokens \u001b[39m\u001b[90m│\u001b[39m\n\u001b[90m├────────────┼────────────────┼─────────────┼──────────────┼──────────────┼──────────────────┼───────────────────┤\u001b[39m\n\u001b[90m│\u001b[39m typescript \u001b[90m│\u001b[39m 24 \u001b[90m│\u001b[39m 8741 \u001b[90m│\u001b[39m 44037 \u001b[90m│\u001b[39m 21 \u001b[90m│\u001b[39m 262 (3.00%) \u001b[90m│\u001b[39m 1759 (3.99%) \u001b[90m│\u001b[39m\n\u001b[90m├────────────┼────────────────┼─────────────┼──────────────┼──────────────┼──────────────────┼───────────────────┤\u001b[39m\n\u001b[90m│\u001b[39m \u001b[1mTotal:\u001b[22m \u001b[90m│\u001b[39m 24 \u001b[90m│\u001b[39m 8741 \u001b[90m│\u001b[39m 44037 \u001b[90m│\u001b[39m 21 \u001b[90m│\u001b[39m 262 (3.00%) \u001b[90m│\u001b[39m 1759 (3.99%) \u001b[90m│\u001b[39m\n\u001b[90m└────────────┴────────────────┴─────────────┴──────────────┴──────────────┴──────────────────┴───────────────────┘\u001b[39m\n\u001b[90mFound 21 clones.\u001b[39m\n\u001b[90mtime: 290.216ms\u001b[39m\n\n\u001b[90m💡 Auto-refactor with AI: \u001b[1m\u001b[39mnpx skills add https://github.com/kucherenko/jscpd --skill dry-refactoring\u001b[90m\u001b[22m\n\u001b[90m🎩 New: Gangsta Agents — discipline your AI coding → gangsta.page\u001b[39m\n\u001b[90m💖 Support jscpd project → https://opencollective.com/jscpd\u001b[39m\n" + }, + "escape_hatches": { + "command": "grep -rEn 'as any|: any|@ts-(ignore|expect-error)' src/ | wc -l", + "exit_code": 1, + "value": 0, + "raw_output": "" + }, + "debt_markers": { + "command": "grep -rEn 'TODO|FIXME|HACK|XXX' src/ | wc -l", + "exit_code": 1, + "value": 0, + "raw_output": "" + } + } +} diff --git a/.agent/skill-runs.jsonl b/.agent/skill-runs.jsonl new file mode 100644 index 0000000..29a68a7 --- /dev/null +++ b/.agent/skill-runs.jsonl @@ -0,0 +1,2 @@ +{"skill":"/deep-clean","ts":"2026-08-15T04:14:22Z","project":"agent-gateway-x402-priced-reservation","target":"PR #11 payment lifecycle and recovery: 26 files","operatorPrompt":"","durationMin":null,"verdict":"60 LOC and 2249 bytes of proven duplicate code removed; all gates green","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/orchestrate","ts":"2026-08-15T04:14:22Z","project":"agent-gateway-x402-priced-reservation","target":"agent-gateway PR #11 deep-clean","operatorPrompt":"","durationMin":null,"verdict":"six read-only Luna Max audits synthesized; local commit b83f1f2; all gates green","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} diff --git a/README.md b/README.md index 289015c..f92dc77 100644 --- a/README.md +++ b/README.md @@ -14,11 +14,15 @@ npm install @tangle-network/agent-gateway ```ts import { createAgentGateway, + recoverPayments, + SqlPaymentRecoveryStore, verifyApiKeyFromStore, } from '@tangle-network/agent-gateway' import { Hono } from 'hono' const app = new Hono() +const paymentRecoveryStore = new SqlPaymentRecoveryStore(sqlAdapter) +await paymentRecoveryStore.migrate() app.route('/v1/agents', createAgentGateway({ resolveAgent: loadPublishedAgent, getSandbox: openAgentSandbox, @@ -28,8 +32,11 @@ app.route('/v1/agents', createAgentGateway({ chainId: 3799, currencyDecimals: 6, verifySigner: verifySpendAuthSignature, + paymentProtocolVersion: 2, + paymentOperations, authorizePayment: reserveSpendAuthorization, }, + paymentRecovery: { store: paymentRecoveryStore }, defaultOutputTokens: 1024, maxOutputTokens: 4096, verifyApiKey: (authHeader) => verifyApiKeyFromStore(authHeader, apiKeyStore), @@ -41,23 +48,46 @@ Set `x402.demoMode: true` only for local development and tests; that explicit mo Keep `verifySigner` free of side effects. Use `authorizePayment` to reserve or claim funds after rate limits, content checks, and product authorization succeed. For production version 2, set `x402.paymentProtocolVersion: 2`, provide `paymentOperations`, and return its operation from `authorizePayment`. +Production version 2 also requires a durable `paymentRecovery.store`. +Run `recoverPayments(config)` from a private scheduled worker. +Every live request and worker uses a unique durable fence token. +A stale request or worker cannot update a row after another worker takes its lease. +Provider settlement, recovery, and release methods must still use the operation ID idempotently. The operation store owns claim, execution start, receipt retention, partial settle, release, and expiry reclaim. An executing or retained operation cannot expire into a refund. -A retained operation can settle later when recovery obtains its usage receipt. +A retained operation settles from its receipt when one exists. +If the receipt does not arrive before `receiptTimeoutMs`, recovery settles the original quoted ceiling. +The fallback never settles the payer's larger authorization amount. Keep version 1 explicitly configured while old and new gateways coexist; shared nonce storage must reject a version 1 claim owned by a version 2 operation. Before it calls the verifier, the gateway requires the signed amount to cover filtered input plus the requested output limit. The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. Sandbox adapters should emit a complete `sandbox.usage` receipt. -Requests with a version 2 payment operation reject missing receipts; API-key and MPP requests keep the legacy visible-token estimate path. +Requests with a version 2 operation or generic MPP charge reject missing receipts. +API-key requests keep the legacy visible-token estimate path. recordUsage must atomically upsert by event.requestId; recovery may retry an event after its acknowledgement is lost. Older custom A2A task stores remain source-compatible through a process-safe fallback. Use an atomic task store for multi-worker production deployments. MPP is method-specific. -Configure `mpp.verifySigner` for production MPP credentials; it receives the decoded JSON payload when available plus the original decoded credential, and returns the authenticated consumer ID or `null`. +Configure `mpp.authenticateCredential` for production MPP credentials. +This callback receives the decoded payload and live credential. +It returns `{ consumerId, paymentIdentity }` or `null`. +`paymentIdentity` must be a stable, non-secret processor identity. +Equivalent encodings of one credential must return the same payment identity. +It must not reserve, confirm, or consume payment. The default `blueprintevm` method may reuse `x402.verifySigner` when its credential has the compatible x402 payload shape. -Other methods are not accepted until they have their own verifier. +Every other method requires an `mpp.charge` lifecycle. +The lifecycle confirms payment after all request denials. +The gateway then acquires its execution fence before it returns a response or starts sandbox work. +`confirmPayment` must bind the provider operation to the supplied `operationId` before confirmation. +It must return only after it verifies final payment success. +`recoverPayment` must inspect that operation ID and must never create another charge. +An authoritative `not-found` result must fence the operation ID against a later charge. +`releasePayment` must perform an idempotent refund or release. +The live credential is passed to `confirmPayment` only on the original request. +The nonce and recovery stores persist only the SHA-256 digest of `paymentIdentity`. +`Payment-Receipt` values must contain visible ASCII only. The same authentication, authorization, rate-limit, filtering, sandbox, settlement, and usage-recording pipeline is used by the OpenAI-compatible and A2A endpoints. Wire protocol handlers only translate their request and response shapes. @@ -67,6 +97,8 @@ Wire protocol handlers only translate their request and response shapes. The gateway speaks Google's A2A protocol alongside its OpenAI-compatible surface: discovery via `.well-known/agent.json`, JSON-RPC 2.0 dispatch for `message/send`, `message/stream`, `tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and the four `tasks/pushNotificationConfig/*` methods. Long-horizon agents — durable tasks across worker restarts, webhook delivery on terminal state, `input-required` pauses with multi-turn continuation — are documented in [`docs/a2a-long-horizon.md`](./docs/a2a-long-horizon.md). Production A2A task control requires `a2a.authorizeTaskAccess`; explicit demo mode is the local-test exception. Custom production task stores must implement atomic `createIfAbsent` and `compareAndSet` methods. +Task stores must retain gateway recovery metadata until reconciliation clears it. +The bundled memory and SQL stores enforce this rule even after the normal task TTL. Push destinations must use HTTPS without URL credentials. ## Tier diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 77fa68a..99094dd 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -20,17 +20,21 @@ The gateway rejects a production task store that lacks either method before it s Explicit `x402.demoMode` keeps a read-then-write fallback for local tests only. Use `SqlTaskStore` or another atomic adapter for multi-worker production deployments. -Before payment settlement, the gateway stores a recovery record in task metadata. +Before a payment provider can mutate state, the gateway stores the recovery identity in task metadata. The record contains the payment operation, usage receipt, output artifact, and a five-minute lease. +The task also points to the shared payment recovery outbox. +This outbox survives task-handler crashes and supports scheduled recovery without a task read. After a restart, the first task read after lease expiry resumes settlement with the same operation. If settlement acknowledgement fails, the gateway keeps the operation and retries after the lease expires. Cancellation stores the recovery record before it completes the canceled task. The canceled task therefore keeps a retryable lease when settlement acknowledgement fails. -If cancellation must release an unused durable operation, the gateway stores a separate release record before calling the release adapter. -An ambiguous release acknowledgement is retried after its five-minute lease, and the record is cleared only after release acknowledgement. +If cancellation must release an unused durable operation, the shared outbox enters `releasing` before the adapter call. +An ambiguous release acknowledgement remains in that outbox and is retried by operation ID. Usage attribution must atomically upsert by `requestId`. The finalization record stores whether attribution was acknowledged, so recovery does not repeat an acknowledged usage event. -If the record is malformed or has no recoverable operation, the gateway expires the task as failed. +If a legacy record is malformed or has no recoverable operation, the gateway expires the task as failed. +If work exists without a receipt, the outbox settles the original quoted ceiling after its configured timeout. +The task-store TTL cannot delete a task while any gateway recovery marker remains. Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. The hook receives the task and request headers so the application can enforce task ownership. @@ -132,7 +136,10 @@ CREATE TABLE IF NOT EXISTS a2a_tasks ( CREATE INDEX IF NOT EXISTS idx_a2a_tasks_context ON a2a_tasks (context_id, updated_at); ``` -One table, JSON payload. TTL is enforced at read time (default 1 hour, configurable via `new SqlTaskStore(db, { ttlMs })`); expired rows are lazily deleted so callers see consistent "expired" semantics regardless of when the GC actually runs. +One table stores the JSON payload. +TTL is enforced at read time and defaults to one hour. +Configure it with `new SqlTaskStore(db, { ttlMs })`. +Expired rows are lazily deleted only when they have no payment recovery marker. `SqlTaskStore` also exposes `listByContext(contextId)` for surfacing all tasks in a conversation when the consumer's UI wants to show a thread view. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 9be5650..157f3dd 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -25,6 +25,12 @@ import { settleAndRecord, } from '../dispatch' import type { PaymentOperation } from '../payment-operations' +import { + deserializePaymentOperation, + serializePaymentOperation, + type SerializedPaymentOperation, +} from '../payment-recovery' +import { recoverPayment as recoverDurablePayment } from '../payment-recovery-worker' import type { GatewayConfig, PaymentMethod, @@ -61,21 +67,6 @@ export interface A2AHandlerDeps { pushStore?: PushNotificationStore } -interface SerializedPaymentOperation { - protocolVersion: 2 - operationId: string - acquiredByRequestId: string - executionStartedAt?: number - retentionReason?: string - nonceKey: string - authorizationId: string - reservedAmount: string - settledAmount: string - refundAmount: string - expiresAt: number - state: PaymentOperation['state'] -} - type FinalizationState = 'completed' | 'input-required' | 'canceled' interface FinalizationRecord { @@ -112,6 +103,11 @@ interface PaymentReleaseRecord { recoveryError?: string } +interface TaskPaymentRecoveryMarker { + version: 1 + id: string +} + /** Terminal task states — fire-once push delivery occurs on these transitions. */ const TERMINAL_STATES: ReadonlySet = new Set([ 'completed', @@ -247,6 +243,7 @@ async function handleMessageSend( const guard = await guardMessageRequest(c, slug, req, deps) if (guard instanceof Response) return guard const { authz, task } = guard + setPaymentResponseHeaders(c, authz) const controller = cancels.register(task.id) try { return await executeMessageSend(c, req, deps, authz, task, controller.signal) @@ -285,6 +282,7 @@ async function executeMessageSend( let inputRequiredSeen = false let finalizationLeaseId: string | undefined try { + await beginPaymentExecution(authz, deps.config) for await (const event of dispatchSandboxStreamRich( authz.agent, authz.userMessage, @@ -293,10 +291,8 @@ async function executeMessageSend( signal, task.id, authz.maxOutputTokens, - async () => { - await beginPaymentExecution(authz, deps.config) - }, - authz.paymentOperation !== undefined, + undefined, + authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, () => { workObserved = true }, @@ -395,7 +391,7 @@ async function executeMessageSend( }, }) usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - const settledBase = clearFinalizationMarker(usageRecordedTask) + const settledBase = clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)) const result = inputRequiredSeen ? withStatus( settledBase, @@ -464,8 +460,49 @@ async function handleMessageStream( const guard = await guardMessageRequest(c, slug, req, deps) if (guard instanceof Response) return guard const { authz, task } = guard + setPaymentResponseHeaders(c, authz) const controller = cancels.register(task.id) + const workingStatus: TaskStatusUpdateEvent = { + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: { state: 'working', timestamp: nowIso() }, + final: false, + } + const workingTask: Task = task.status.state === 'working' + ? task + : { ...task, status: workingStatus.status } + if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { + cancels.clear(task.id) + const current = await releaseTaskPayment( + authz, + await deps.taskStore.get(task.id) ?? task, + deps, + 'A2A task changed before execution started', + false, + ) + if (current.status.state === 'canceled') return c.json(ok(req.id, current)) + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) + } + try { + await beginPaymentExecution(authz, deps.config) + } catch (error) { + cancels.clear(task.id) + await releaseTaskPayment( + authz, + workingTask, + deps, + error instanceof Error ? error.message : String(error), + false, + ) + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment execution authorization failed')) + } + if (controller.signal.aborted) { + cancels.clear(task.id) + const canceled = await completeCanceledTask(authz, workingTask, '', undefined, false, deps) + return c.json(ok(req.id, canceled)) + } let responseText = '' let usage: SandboxUsageReceipt | undefined let workObserved = false @@ -477,24 +514,10 @@ async function handleMessageStream( ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) } - // Status: working - const workingStatus: TaskStatusUpdateEvent = { - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: { state: 'working', timestamp: nowIso() }, - final: false, - } - const workingTask: Task = task.status.state === 'working' - ? task - : { ...task, status: workingStatus.status } let inputRequiredPrompt: string | undefined let inputRequiredSeen = false let finalizationLeaseId: string | undefined try { - if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { - throw new Error('A2A task changed before execution started') - } send(workingStatus) for await (const event of dispatchSandboxStreamRich( @@ -505,10 +528,8 @@ async function handleMessageStream( controller.signal, task.id, authz.maxOutputTokens, - async () => { - await beginPaymentExecution(authz, deps.config) - }, - authz.paymentOperation !== undefined, + undefined, + authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, () => { workObserved = true }, @@ -616,7 +637,7 @@ async function handleMessageStream( if (inputRequiredSeen) { const paused = withStatus( - clearFinalizationMarker(usageRecordedTask), + clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), 'input-required', inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, responseText @@ -648,7 +669,7 @@ async function handleMessageStream( } // Final: persist the terminal task before emitting terminal events. - const completed = withStatus(clearFinalizationMarker(usageRecordedTask), 'completed', undefined, [ + const completed = withStatus(clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), 'completed', undefined, [ responseTextToArtifact(responseText, `${task.id}-artifact-0`), ]) if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, completed)) { @@ -762,6 +783,12 @@ async function handleMessageStream( 'X-Request-Id': authz.requestId, 'X-Agent-Slug': authz.agent.slug, 'X-Task-Id': task.id, + ...(authz.mppChargeOperation + ? { 'Payment-Receipt': authz.mppChargeOperation.receipt } + : {}), + ...(authz.paymentRecoveryId + ? { 'X-Payment-Operation-Id': authz.paymentRecoveryId } + : {}), }, }) } @@ -1104,9 +1131,9 @@ async function guardMessageRequest( fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${existing.id}' changed before continuation`), ) } - const claimError = await claimTaskPayment(c, req, continued, authz, deps, existing) - if (claimError) return claimError - return { authz, task: continued } + const claimedTask = await claimTaskPayment(c, req, continued, authz, deps, existing) + if (claimedTask instanceof Response) return claimedTask + return { authz, task: claimedTask } } // Unknown taskId in params: fall through and mint a fresh task with that // exact id so callers that pre-allocate ids (idempotency) get them. @@ -1131,9 +1158,9 @@ async function guardMessageRequest( fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' already exists`), ) } - const claimError = await claimTaskPayment(c, req, task, authz, deps) - if (claimError) return claimError - return { authz, task } + const claimedTask = await claimTaskPayment(c, req, task, authz, deps) + if (claimedTask instanceof Response) return claimedTask + return { authz, task: claimedTask } } async function claimTaskPayment( @@ -1143,20 +1170,49 @@ async function claimTaskPayment( authz: AuthorizedRequest, deps: A2AHandlerDeps, paymentFailureTask?: Task, -): Promise { +): Promise { + let paymentTask = task try { - await claimPayment(authz, deps.config, deps.state) - return undefined + await claimPayment(authz, deps.config, deps.state, { + onRecoveryPrepared: async (recoveryId) => { + paymentTask = await attachPaymentRecoveryMarker( + deps.taskStore, + paymentTask, + recoveryId, + ) + }, + }) } catch { - const releasedTask = await releaseTaskPayment(authz, task, deps, 'payment authorization failed', false) + let recoveryTask = paymentTask + try { + recoveryTask = await retainPaymentRecoveryMarker( + deps.taskStore, + paymentTask, + authz.paymentRecoveryId, + ) + } catch (error) { + console.error( + `[a2a] failed to attach payment recovery to ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + } + const releasedTask = await releaseTaskPayment( + authz, + recoveryTask, + deps, + 'payment authorization failed', + false, + ) const releaseRecord = releasedTask.metadata?.[PAYMENT_RELEASE_METADATA_KEY] const failed = releaseRecord !== undefined ? isTerminal(releasedTask.status.state) ? releasedTask : withStatus(releasedTask, 'failed') - : paymentFailureTask ?? (isTerminal(releasedTask.status.state) - ? releasedTask - : withStatus(releasedTask, 'failed')) + : paymentFailureTask + ? preservePaymentRecoveryMarker(paymentFailureTask, releasedTask) + : isTerminal(releasedTask.status.state) + ? releasedTask + : withStatus(releasedTask, 'failed') try { if (await compareAndSetTask(deps.taskStore, releasedTask, failed) && isTerminal(failed.status.state)) { await maybeDeliverPush(failed, deps) @@ -1169,6 +1225,38 @@ async function claimTaskPayment( } return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment authorization failed')) } + + const current = await deps.taskStore.get(task.id) + if (current && JSON.stringify(current) === JSON.stringify(paymentTask)) { + return paymentTask + } + { + const changedTask = current ?? paymentTask + let recoveryTask = changedTask + try { + recoveryTask = await retainPaymentRecoveryMarker( + deps.taskStore, + changedTask, + authz.paymentRecoveryId, + ) + } catch (error) { + console.error( + `[a2a] failed to retain payment recovery after task race for ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + } + const released = await releaseTaskPayment( + authz, + recoveryTask, + deps, + 'A2A task changed during payment confirmation', + false, + ) + if (released.status.state === 'canceled') return c.json(ok(req.id, released)) + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed during payment confirmation`), + ) + } } async function releaseTaskPayment( @@ -1179,7 +1267,12 @@ async function releaseTaskPayment( workObserved: boolean, ): Promise { // Store the operation before release because the adapter acknowledgement can be ambiguous. - if (!workObserved && authz.paymentOperation && deps.config.x402.paymentOperations) { + if ( + !workObserved && + !authz.paymentRecoveryId && + authz.paymentOperation && + deps.config.x402.paymentOperations + ) { let marked: Task try { marked = await beginPaymentReleaseRecovery(deps.taskStore, task, authz, reason) ?? task @@ -1218,7 +1311,10 @@ async function releaseTaskPayment( releaseError instanceof Error ? releaseError.message : String(releaseError), ) } - return await deps.taskStore.get(task.id) ?? task + const current = await deps.taskStore.get(task.id) ?? task + return workObserved + ? current + : clearReconciledPaymentRecoveryMarker(current, deps) } /** Keep an owned finalization record durable when settlement acknowledgement is lost. */ @@ -1292,6 +1388,12 @@ async function completeCanceledTask( }) usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) } catch (settlementError) { + await releasePaymentAfterFailure( + authz, + deps.config, + settlementError instanceof Error ? settlementError.message : String(settlementError), + true, + ) const retained = await retainFinalizationForRecovery( deps.taskStore, task.id, @@ -1308,7 +1410,7 @@ async function completeCanceledTask( } const canceled = withStatus( - clearFinalizationMarker(usageRecordedTask), + clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), 'canceled', undefined, responseText @@ -1341,9 +1443,19 @@ async function completeCanceledTask( const FINALIZING_METADATA_KEY = 'gatewayFinalizing' const PAYMENT_RELEASE_METADATA_KEY = 'gatewayPaymentRelease' +const PAYMENT_RECOVERY_METADATA_KEY = 'gatewayPaymentRecovery' const FINALIZATION_LEASE_MS = 5 * 60 * 1000 const PAYMENT_RELEASE_LEASE_MS = 5 * 60 * 1000 +function setPaymentResponseHeaders(c: Context, authz: AuthorizedRequest): void { + if (authz.mppChargeOperation) { + c.header('Payment-Receipt', authz.mppChargeOperation.receipt) + } + if (authz.paymentRecoveryId) { + c.header('X-Payment-Operation-Id', authz.paymentRecoveryId) + } +} + async function authorizeTaskAccess( c: Context, req: JSONRPCRequest, @@ -1399,6 +1511,86 @@ async function compareAndSetTask(taskStore: TaskStore, expected: Task, next: Tas return taskStore.compareAndSet(expected, next) } +async function attachPaymentRecoveryMarker( + taskStore: TaskStore, + task: Task, + recoveryId: string | undefined, +): Promise { + if (!recoveryId) return task + const existing = readPaymentRecoveryMarker(task) + if (existing) { + if (existing.id !== recoveryId) { + throw new Error('A2A task already has a different payment recovery identity') + } + return task + } + const next = withPaymentRecoveryMarker(task, recoveryId) + if (await compareAndSetTask(taskStore, task, next)) return next + throw new Error('A2A task changed while payment recovery was attached') +} + +/** Attach only as a retention marker. The returned task must never execute. */ +async function retainPaymentRecoveryMarker( + taskStore: TaskStore, + task: Task, + recoveryId: string | undefined, +): Promise { + if (!recoveryId) return task + let current = await taskStore.get(task.id) ?? task + for (let attempt = 0; attempt < 16; attempt += 1) { + const existing = readPaymentRecoveryMarker(current) + if (existing) { + if (existing.id !== recoveryId) { + throw new Error('A2A task already has a different payment recovery identity') + } + return current + } + const next = withPaymentRecoveryMarker(current, recoveryId) + if (await compareAndSetTask(taskStore, current, next)) return next + const latest = await taskStore.get(task.id) + if (!latest) throw new Error('A2A task disappeared while payment recovery was retained') + current = latest + } + throw new Error('A2A task changed too many times while payment recovery was retained') +} + +function withPaymentRecoveryMarker(task: Task, recoveryId: string): Task { + return { + ...task, + metadata: { + ...(task.metadata ?? {}), + [PAYMENT_RECOVERY_METADATA_KEY]: { version: 1, id: recoveryId }, + }, + } +} + +function preservePaymentRecoveryMarker(base: Task, source: Task): Task { + const marker = readPaymentRecoveryMarker(source) + return marker ? withPaymentRecoveryMarker(base, marker.id) : base +} + +function readPaymentRecoveryMarker(task: Task): TaskPaymentRecoveryMarker | undefined { + const raw = task.metadata?.[PAYMENT_RECOVERY_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const marker = raw as Partial + if (marker.version !== 1 || typeof marker.id !== 'string' || marker.id.length === 0) { + return undefined + } + return marker as TaskPaymentRecoveryMarker +} + +function clearPaymentRecoveryMarker(task: Task): Task { + if (!task.metadata || !(PAYMENT_RECOVERY_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[PAYMENT_RECOVERY_METADATA_KEY] + return Object.keys(metadata).length > 0 + ? { ...task, metadata } + : (() => { + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata + })() +} + function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): TaskStore { if (typeof taskStore.createIfAbsent === 'function' && typeof taskStore.compareAndSet === 'function') { return taskStore @@ -1456,46 +1648,6 @@ function buildFinalizationRecord( } } -function serializePaymentOperation(operation: PaymentOperation): SerializedPaymentOperation { - return { - protocolVersion: 2, - operationId: operation.operationId, - acquiredByRequestId: operation.acquiredByRequestId, - ...(operation.executionStartedAt !== undefined - ? { executionStartedAt: operation.executionStartedAt } - : {}), - ...(operation.retentionReason ? { retentionReason: operation.retentionReason } : {}), - nonceKey: operation.nonceKey, - authorizationId: operation.authorizationId, - reservedAmount: operation.reservedAmount.toString(), - settledAmount: operation.settledAmount.toString(), - refundAmount: operation.refundAmount.toString(), - expiresAt: operation.expiresAt, - state: operation.state, - } -} - -function deserializePaymentOperation(value: SerializedPaymentOperation): PaymentOperation { - if (value.protocolVersion !== 2) throw new Error('unsupported A2A payment operation version') - if (!value.operationId || !value.acquiredByRequestId || !value.nonceKey || !value.authorizationId) { - throw new Error('incomplete A2A payment operation recovery record') - } - return { - protocolVersion: 2, - operationId: value.operationId, - acquiredByRequestId: value.acquiredByRequestId, - ...(value.executionStartedAt !== undefined ? { executionStartedAt: value.executionStartedAt } : {}), - ...(value.retentionReason ? { retentionReason: value.retentionReason } : {}), - nonceKey: value.nonceKey, - authorizationId: value.authorizationId, - reservedAmount: BigInt(value.reservedAmount), - settledAmount: BigInt(value.settledAmount), - refundAmount: BigInt(value.refundAmount), - expiresAt: value.expiresAt, - state: value.state, - } -} - function withFinalizationRecord(task: Task, record: FinalizationRecord): Task { return { ...task, @@ -1636,7 +1788,7 @@ async function clearPaymentReleaseRecovery( const current = await taskStore.get(task.id) ?? task const record = readPaymentReleaseRecord(current) if (!record || record.lease.id !== leaseId) return current - const cleared = clearPaymentReleaseRecord(current) + const cleared = clearPaymentRecoveryMarker(clearPaymentReleaseRecord(current)) if (await compareAndSetTask(taskStore, current, cleared)) return cleared return await taskStore.get(task.id) ?? cleared } @@ -1692,6 +1844,12 @@ async function recoverPaymentReleaseIfNeeded( } const operation = deserializePaymentOperation(renewed.paymentOperation) await deps.config.x402.paymentOperations.releasePayment(operation, renewed.reason) + if (deps.config.paymentRecovery) { + const recovered = await recoverDurablePayment(renewed.operationId, deps.config, { force: true }) + if (recovered && recovered.state !== 'reconciled') { + throw new Error('durable payment release is still pending') + } + } const recovered = clearPaymentReleaseRecord(leasedTask) if (!await compareAndSetTask(deps.taskStore, leasedTask, recovered)) { return await deps.taskStore.get(task.id) ?? recovered @@ -1744,6 +1902,24 @@ async function recoverFinalizationIfNeeded( const agent = await deps.config.resolveAgent(agentSlug) if (!agent || !agent.enabled) throw new Error('A2A recovery agent is unavailable') + const paymentRecovery = readPaymentRecoveryMarker(leasedTask) + if (paymentRecovery && deps.config.paymentRecovery) { + const recovery = await recoverDurablePayment(paymentRecovery.id, deps.config, { + force: true, + usage: renewed.receipt, + }) + if (recovery?.state !== 'reconciled') { + throw new Error('durable payment finalization is still pending') + } + const usageRecordedTask = await markUsageRecorded(deps.taskStore, leasedTask) + const recoveredTask = finalizationResultTask(usageRecordedTask, renewed) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, recoveredTask)) { + return await deps.taskStore.get(task.id) ?? recoveredTask + } + await maybeDeliverPush(recoveredTask, deps) + return recoveredTask + } + let paymentOperation: PaymentOperation | undefined if (renewed.operationId || renewed.paymentOperation) { if (!renewed.operationId || !renewed.paymentOperation) { @@ -1773,6 +1949,9 @@ async function recoverFinalizationIfNeeded( executionBudget: renewed.executionBudget, requiredPaymentAmount: 0n, paymentPayload: null, + ...(readPaymentRecoveryMarker(leasedTask) + ? { paymentRecoveryId: readPaymentRecoveryMarker(leasedTask)!.id } + : {}), ...(paymentOperation ? { paymentOperation, paymentOperationAcquired: true } : {}), @@ -1805,7 +1984,10 @@ async function recoverFinalizationIfNeeded( `[a2a] finalization recovery failed for ${task.id}:`, recoveryError.message, ) - if (renewed.operationId && renewed.paymentOperation) { + if ( + (readPaymentRecoveryMarker(leasedTask) && deps.config.paymentRecovery) || + (renewed.operationId && renewed.paymentOperation) + ) { const retained = await retainFinalizationForRecovery( deps.taskStore, task.id, @@ -1824,11 +2006,44 @@ async function recoverTaskIfNeeded( requestedAgentSlug: string, ): Promise { const released = await recoverPaymentReleaseIfNeeded(task, deps) - return recoverFinalizationIfNeeded(released, deps, requestedAgentSlug) + const finalized = await recoverFinalizationIfNeeded(released, deps, requestedAgentSlug) + return recoverPaymentMarkerIfNeeded(finalized, deps) +} + +async function recoverPaymentMarkerIfNeeded( + task: Task, + deps: A2AHandlerDeps, +): Promise { + const marker = readPaymentRecoveryMarker(task) + if (!marker || !deps.config.paymentRecovery) return task + try { + const record = await recoverDurablePayment(marker.id, deps.config) + if (record?.state !== 'reconciled') return task + return clearReconciledPaymentRecoveryMarker(task, deps) + } catch (error) { + console.error( + `[a2a] durable payment recovery failed for ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + return task + } +} + +async function clearReconciledPaymentRecoveryMarker( + task: Task, + deps: A2AHandlerDeps, +): Promise { + const marker = readPaymentRecoveryMarker(task) + if (!marker || !deps.config.paymentRecovery) return task + const record = await deps.config.paymentRecovery.store.get(marker.id) + if (record?.state !== 'reconciled') return task + const cleared = clearPaymentRecoveryMarker(task) + if (await compareAndSetTask(deps.taskStore, task, cleared)) return cleared + return await deps.taskStore.get(task.id) ?? cleared } function finalizationResultTask(task: Task, record: FinalizationRecord): Task { - const cleanTask = clearFinalizationMarker(task) + const cleanTask = clearPaymentRecoveryMarker(clearFinalizationMarker(task)) const finalState = record.finalState ?? ( task.status.state === 'canceled' ? 'canceled' @@ -1917,12 +2132,7 @@ function clearFinalizationMarker(task: Task): Task { } function isTerminal(state: Task['status']['state']): boolean { - return ( - state === 'completed' || - state === 'canceled' || - state === 'failed' || - state === 'rejected' - ) + return TERMINAL_STATES.has(state) } function asError(error: unknown): Error { diff --git a/src/a2a/task-store-sql.ts b/src/a2a/task-store-sql.ts index 58a3baf..96e3cc2 100644 --- a/src/a2a/task-store-sql.ts +++ b/src/a2a/task-store-sql.ts @@ -40,7 +40,7 @@ * await store.migrate() */ -import type { TaskStore } from './task-store' +import { hasPendingPaymentRecovery, type TaskStore } from './task-store' import type { Task } from './types' /** @@ -136,15 +136,19 @@ export class SqlTaskStore implements TaskStore { ) const row = rows[0] if (!row) return undefined - if (Date.now() - row.updated_at > this.ttlMs) { + const task = JSON.parse(row.payload) as Task + if (Date.now() - row.updated_at > this.ttlMs && !hasPendingPaymentRecovery(task)) { // Lazy GC. If the delete loses a race with another reader, that reader // observes either the stale-then-deleted task (returning undefined here) // or, after this delete commits, observes undefined directly — either // way callers see consistent "expired" semantics. - void this.db.exec(`DELETE FROM ${this.table} WHERE id = ?`, [id]) + void this.db.exec( + `DELETE FROM ${this.table} WHERE id = ? AND payload = ?`, + [id, row.payload], + ) return undefined } - return JSON.parse(row.payload) as Task + return task } async put(task: Task): Promise { @@ -192,7 +196,12 @@ export class SqlTaskStore implements TaskStore { } async delete(id: string): Promise { - await this.db.exec(`DELETE FROM ${this.table} WHERE id = ?`, [id]) + const task = await this.get(id) + if (!task || hasPendingPaymentRecovery(task)) return + await this.db.exec( + `DELETE FROM ${this.table} WHERE id = ? AND payload = ?`, + [id, JSON.stringify(task)], + ) } /** @@ -208,7 +217,10 @@ export class SqlTaskStore implements TaskStore { ) const now = Date.now() return rows - .filter((r) => now - r.updated_at <= this.ttlMs) - .map((r) => JSON.parse(r.payload) as Task) + .map((r) => ({ task: JSON.parse(r.payload) as Task, updatedAt: r.updated_at })) + .filter(({ task, updatedAt }) => + now - updatedAt <= this.ttlMs || hasPendingPaymentRecovery(task), + ) + .map(({ task }) => task) } } diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index 85fa1ff..c986273 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -19,6 +19,17 @@ export interface TaskStore { const DEFAULT_TTL_MS = 60 * 60 * 1000 +const PAYMENT_RECOVERY_KEYS = [ + 'gatewayFinalizing', + 'gatewayPaymentRelease', + 'gatewayPaymentRecovery', +] as const + +/** Recovery-bearing tasks must remain readable until reconciliation clears the marker. */ +export function hasPendingPaymentRecovery(task: Task): boolean { + return PAYMENT_RECOVERY_KEYS.some((key) => task.metadata?.[key] !== undefined) +} + export class InMemoryTaskStore implements TaskStore { private readonly entries = new Map() @@ -52,6 +63,8 @@ export class InMemoryTaskStore implements TaskStore { } async delete(id: string): Promise { + const entry = this.entries.get(id) + if (entry && hasPendingPaymentRecovery(entry.task)) return this.entries.delete(id) } @@ -62,7 +75,9 @@ export class InMemoryTaskStore implements TaskStore { private gc(): void { const now = Date.now() for (const [id, entry] of this.entries) { - if (entry.expiresAt <= now) this.entries.delete(id) + if (entry.expiresAt <= now && !hasPendingPaymentRecovery(entry.task)) { + this.entries.delete(id) + } } } } diff --git a/src/dispatch.ts b/src/dispatch.ts index 35aa294..47b9760 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -11,10 +11,26 @@ import type { Context } from 'hono' import { filterConsumerMessagesStrict, redactSystemPromptFromOutput } from './filter' +import { + assertMppChargeOperation, + mppPaymentOperationId, + type MppChargeOperation, +} from './mpp-payment' import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' import { type RateLimitStore, checkRateLimit } from './rate-limit' import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' -import type { PaymentOperation } from './payment-operations' +import { paymentNonceKey, type PaymentOperation } from './payment-operations' +import { + PAYMENT_RECOVERY_VERSION, + PaymentRecoveryFenceError, + PaymentRecoveryReplayError, + recoveryTiming, + serializePaymentOperation, + updateOwnedPaymentRecovery, + type PaymentRecoveryRecord, + type PaymentRecoveryTarget, + type PaymentSettlementBasis, +} from './payment-recovery' import type { AgentMeta, ApiKeyInfo, @@ -30,7 +46,7 @@ import { isApiKeyAuthEnabled, isMppAuthEnabled, mppPaymentPayload, - mppReplayNonceKey, + mppPaymentCredential, verifyMpp, verifyX402, } from './verify' @@ -67,8 +83,21 @@ export interface AuthorizedRequest { paymentPayload: Record | null paymentNonceKey?: string mppMethod?: string + /** Live generic MPP credential. Never written to the recovery store. */ + mppCredential?: string + /** Stable method identity. The gateway persists only its digest. */ + mppPaymentIdentity?: string + mppChargeOperation?: MppChargeOperation paymentOperation?: PaymentOperation paymentOperationAcquired?: boolean + paymentRecoveryId?: string + /** Unique ownership fence for live or recovery transitions. */ + paymentRecoveryFence?: string +} + +export interface PaymentClaimHooks { + /** Persist the recovery identity before the provider can mutate payment state. */ + onRecoveryPrepared?: (recoveryId: string) => Promise } function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { @@ -245,6 +274,8 @@ export async function authenticateAndGuard( let x402Payload: Record | null = null let paymentNonceKey: string | undefined let mppMethod: string | undefined + let mppCredential: string | undefined + let mppPaymentIdentity: string | undefined if (spendAuthHeader) { const signer = await verifyX402( @@ -281,7 +312,7 @@ export async function authenticateAndGuard( consumerId = signer paymentMethod = 'x402' } else if (isMppAuthEnabled(config) && authHeader.toLowerCase().startsWith('payment ')) { - const signer = await verifyMpp( + const authenticated = await verifyMpp( authHeader, config.mpp!, config.x402, @@ -289,7 +320,7 @@ export async function authenticateAndGuard( requiredPaymentAmount, false, ) - if (!signer) { + if (!authenticated) { const realm = config.mpp!.realm const method = config.mpp!.method ?? 'blueprintevm' await state.obs?.onAuthFailure?.(ctx, { @@ -314,11 +345,13 @@ export async function authenticateAndGuard( }, ) } - consumerId = signer + consumerId = authenticated.consumerId paymentMethod = 'mpp' mppMethod = authHeader.match(/^Payment\s+(\S+)\s+/i)?.[1]?.toLowerCase() + mppCredential = mppPaymentCredential(authHeader) + mppPaymentIdentity = authenticated.paymentIdentity x402Payload = mppPaymentPayload(authHeader) ?? null - paymentNonceKey = mppReplayNonceKey(authHeader) + paymentNonceKey = authenticated.replayKey } else if (authHeader.startsWith('Bearer ')) { const verify = config.verifyApiKey ?? (config.x402.demoMode ? defaultVerifyApiKey : null) if (!verify || !isApiKeyAuthEnabled(config)) { @@ -498,6 +531,8 @@ export async function authenticateAndGuard( paymentPayload: x402Payload, paymentNonceKey, mppMethod, + mppCredential, + mppPaymentIdentity, } } @@ -506,14 +541,15 @@ export async function claimPayment( authz: AuthorizedRequest, config: GatewayConfig, state: GatewayState, + hooks: PaymentClaimHooks = {}, ): Promise { if (authz.paymentMethod === 'x402' && authz.paymentPayload) { - const context = { - requestId: authz.requestId, - agentId: authz.agent.id, - requiredAmount: authz.requiredPaymentAmount, - maxOutputTokens: authz.maxOutputTokens, - executionBudget: authz.executionBudget, + const context = paymentAuthorizationContext(authz) + if (config.x402.paymentProtocolVersion === 2) { + await preparePaymentRecovery(authz, config, { + kind: 'x402', + operationId: `x402:${paymentNonceKey(authz.paymentPayload)}`, + }, hooks) } let operation: PaymentOperation | undefined if (config.x402.authorizePayment) { @@ -559,6 +595,7 @@ export async function claimPayment( // Attach owned state before the shared nonce claim. If that claim fails, // the caller can persist release recovery after an ambiguous refund. authz.paymentOperation = operation + await markRecoveryClaimed(authz, config) } if (operation && authz.paymentNonceKey) { const claimed = await claimPaymentNonce( @@ -569,7 +606,7 @@ export async function claimPayment( ) if (!claimed) { try { - await config.x402.paymentOperations!.releasePayment(operation, 'shared payment nonce was already owned') + await releasePayment(authz, config, 'shared payment nonce was already owned') } catch (releaseError) { console.error( `[agent-gateway] payment release failed for ${authz.requestId}:`, @@ -585,16 +622,14 @@ export async function claimPayment( } const mppMethod = (authz.mppMethod ?? config.mpp?.method ?? 'blueprintevm').toLowerCase() const durablePayload = durableMppPaymentPayload(authz.paymentPayload) - // Only BlueprinTEVM carries x402 authorization fields; other MPP methods - // must stay on their verifier and method-specific settlement path. + // Only BlueprinTEVM carries x402 authorization fields. Other MPP methods + // use the isolated immediate-charge lifecycle below. if (mppMethod === 'blueprintevm' && durablePayload && config.x402.paymentOperations) { - const context = { - requestId: authz.requestId, - agentId: authz.agent.id, - requiredAmount: authz.requiredPaymentAmount, - maxOutputTokens: authz.maxOutputTokens, - executionBudget: authz.executionBudget, - } + const context = paymentAuthorizationContext(authz) + await preparePaymentRecovery(authz, config, { + kind: 'x402', + operationId: `x402:${paymentNonceKey(durablePayload)}`, + }, hooks) const operation = await config.x402.paymentOperations.claimPayment(durablePayload, context) if (operation.protocolVersion !== 2) throw new Error('payment operation protocol version mismatch') if (operation.acquiredByRequestId !== context.requestId) { @@ -603,6 +638,7 @@ export async function claimPayment( authz.paymentPayload = durablePayload authz.paymentOperation = operation authz.paymentOperationAcquired = true + await markRecoveryClaimed(authz, config) const claimed = await claimPaymentNonce( state.nonceStore, authz.paymentNonceKey, @@ -611,7 +647,7 @@ export async function claimPayment( ) if (!claimed) { try { - await config.x402.paymentOperations.releasePayment(operation, 'shared payment nonce was already owned') + await releasePayment(authz, config, 'shared payment nonce was already owned') } catch (releaseError) { console.error( `[agent-gateway] payment release failed for ${authz.requestId}:`, @@ -620,9 +656,55 @@ export async function claimPayment( } throw new Error('payment nonce was already consumed') } - } else { + } else if (mppMethod === 'blueprintevm') { const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) if (!claimed) throw new Error('payment nonce was already consumed') + } else { + const lifecycle = config.mpp?.charge + if (!lifecycle || lifecycle.protocolVersion !== 1) { + throw new Error('MPP charge lifecycle is not configured') + } + if (!authz.mppCredential) throw new Error('MPP payment credential is unavailable') + if (!authz.mppPaymentIdentity) throw new Error('MPP payment identity is unavailable') + const operationId = await mppPaymentOperationId(mppMethod, authz.mppPaymentIdentity) + await preparePaymentRecovery(authz, config, { + kind: 'mpp-charge', + method: mppMethod, + operationId, + }, hooks) + const claimed = await claimPaymentNonce( + state.nonceStore, + authz.paymentNonceKey, + authz.paymentPayload ?? {}, + `${operationId}:${authz.requestId}`, + ) + if (!claimed) { + await markRecoveryReconciled(authz, config) + throw new Error('payment nonce was already consumed') + } + const operation = await lifecycle.confirmPayment({ + operationId, + requestId: authz.requestId, + agentId: authz.agent.id, + consumerId: authz.consumerId, + method: mppMethod, + credential: authz.mppCredential, + amount: authz.requiredPaymentAmount, + currencyDecimals: config.x402.currencyDecimals ?? 6, + }) + assertMppChargeOperation( + operation, + { operationId, requestId: authz.requestId, method: mppMethod }, + ['confirmed'], + false, + ) + authz.mppChargeOperation = operation + await markRecoveryClaimed(authz, config) + assertMppChargeOperation( + operation, + { operationId, requestId: authz.requestId, method: mppMethod }, + ['confirmed'], + ) } } @@ -648,14 +730,153 @@ export async function claimPayment( } } +function paymentAuthorizationContext(authz: AuthorizedRequest) { + return { + requestId: authz.requestId, + agentId: authz.agent.id, + requiredAmount: authz.requiredPaymentAmount, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + } +} + +async function preparePaymentRecovery( + authz: AuthorizedRequest, + config: GatewayConfig, + payment: PaymentRecoveryTarget, + hooks: PaymentClaimHooks, +): Promise { + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const now = Date.now() + const fenceId = globalThis.crypto.randomUUID() + const leaseExpiresAt = now + recoveryTiming(recovery).staleRequestMs + const record: PaymentRecoveryRecord = { + version: PAYMENT_RECOVERY_VERSION, + id: payment.operationId, + revision: 0, + state: 'claiming', + payment, + attribution: { + requestId: authz.requestId, + agentId: authz.agent.id, + agentSlug: authz.agent.slug, + consumerId: authz.consumerId, + paymentMethod: authz.paymentMethod, + startMs: authz.startMs, + pricePerTokenUsd: authz.agent.pricePerTokenUsd, + platformFeePercent: authz.agent.platformFeePercent, + requiredAmount: authz.requiredPaymentAmount.toString(), + currencyDecimals: config.x402.currencyDecimals ?? 6, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + }, + workStarted: false, + usageRecorded: false, + attempts: 0, + nextAttemptAt: leaseExpiresAt, + lease: { id: fenceId, expiresAt: leaseExpiresAt }, + createdAt: now, + updatedAt: now, + } + if (!await recovery.store.createIfAbsent(record)) { + if ((await recovery.store.get(record.id))?.state === 'reconciled') { + throw new PaymentRecoveryReplayError(record.id) + } + throw new Error('payment recovery identity was already claimed') + } + authz.paymentRecoveryId = record.id + authz.paymentRecoveryFence = fenceId + await hooks.onRecoveryPrepared?.(record.id) +} + +async function markRecoveryClaimed( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const fenceId = requirePaymentRecoveryFence(authz) + const now = Date.now() + const leaseExpiresAt = now + recoveryTiming(recovery).staleRequestMs + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + state: 'claimed', + payment: recoveryTarget(authz, record.payment), + lease: { id: fenceId, expiresAt: leaseExpiresAt }, + nextAttemptAt: leaseExpiresAt, + }), now) +} + +function requirePaymentRecoveryFence(authz: AuthorizedRequest): string { + if (!authz.paymentRecoveryFence) { + throw new Error('payment recovery fence is unavailable') + } + return authz.paymentRecoveryFence +} + +function recoveryTarget( + authz: AuthorizedRequest, + current: PaymentRecoveryTarget, +): PaymentRecoveryTarget { + if (authz.paymentOperation) { + return { + kind: 'x402', + operationId: authz.paymentOperation.operationId, + operation: serializePaymentOperation(authz.paymentOperation), + } + } + if (authz.mppChargeOperation) { + return { + kind: 'mpp-charge', + method: authz.mppChargeOperation.method, + operationId: authz.mppChargeOperation.operationId, + operation: authz.mppChargeOperation, + } + } + return current +} + /** Release an owned operation when execution cannot produce a valid receipt. */ export async function releasePayment( authz: AuthorizedRequest, config: GatewayConfig, reason: string, ): Promise { - if (!authz.paymentOperation || authz.paymentOperationAcquired !== true || !config.x402.paymentOperations) return - await config.x402.paymentOperations.releasePayment(authz.paymentOperation, reason) + const ownsX402 = authz.paymentOperation && + authz.paymentOperationAcquired === true && + config.x402.paymentOperations + const ownsMpp = authz.mppChargeOperation && config.mpp?.charge + if (!ownsX402 && !ownsMpp) { + await relinquishPaymentRecovery(authz, config, Date.now()) + return + } + await markRecoveryReleasing(authz, config, reason) + try { + if (ownsX402) { + authz.paymentOperation = await config.x402.paymentOperations!.releasePayment( + authz.paymentOperation!, + reason, + ) + } else { + const operation = await config.mpp!.charge!.releasePayment(authz.mppChargeOperation!, reason) + assertMppChargeOperation( + operation, + { + operationId: authz.mppChargeOperation!.operationId, + requestId: authz.requestId, + method: authz.mppChargeOperation!.method, + }, + ['released'], + false, + ) + authz.mppChargeOperation = operation + } + } catch (error) { + await relinquishPaymentRecovery(authz, config, Date.now()) + throw error + } + await markRecoveryReconciled(authz, config) } /** Mark a durable reservation active immediately before sandbox execution. */ @@ -663,8 +884,29 @@ export async function beginPaymentExecution( authz: AuthorizedRequest, config: GatewayConfig, ): Promise { - if (!authz.paymentOperation || authz.paymentOperationAcquired !== true || !config.x402.paymentOperations) return - authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) + if (!authz.paymentRecoveryId) { + if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { + authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) + } + return + } + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const fenceId = requirePaymentRecoveryFence(authz) + const now = Date.now() + const fallbackAt = now + recoveryTiming(recovery).receiptTimeoutMs + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + state: 'executing', + payment: recoveryTarget(authz, record.payment), + workStarted: true, + fallbackAt, + lease: { id: fenceId, expiresAt: fallbackAt }, + nextAttemptAt: fallbackAt, + }), now) + if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { + authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) + } } /** @@ -678,6 +920,27 @@ export async function releasePaymentAfterFailure( workObserved: boolean, ): Promise { if (workObserved) { + const recovery = config.paymentRecovery + if (recovery && authz.paymentRecoveryId) { + const fenceId = requirePaymentRecoveryFence(authz) + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => { + if (record.state === 'settling' && record.usage) { + return { ...record, lease: undefined, nextAttemptAt: Date.now() } + } + const fallbackAt = record.fallbackAt ?? + Date.now() + recoveryTiming(recovery).receiptTimeoutMs + return { + ...record, + state: 'retained', + payment: recoveryTarget(authz, record.payment), + workStarted: true, + fallbackAt, + reason, + lease: undefined, + nextAttemptAt: fallbackAt, + } + }) + } if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { authz.paymentOperation = await config.x402.paymentOperations.retainPayment(authz.paymentOperation, reason) } @@ -689,6 +952,61 @@ export async function releasePaymentAfterFailure( await releasePayment(authz, config, reason) } +async function relinquishPaymentRecovery( + authz: AuthorizedRequest, + config: GatewayConfig, + nextAttemptAt: number, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId || !authz.paymentRecoveryFence) return + try { + await updateOwnedPaymentRecovery( + recovery.store, + authz.paymentRecoveryId, + authz.paymentRecoveryFence, + (record) => ({ ...record, lease: undefined, nextAttemptAt }), + ) + } catch (error) { + if (!(error instanceof PaymentRecoveryFenceError)) throw error + } +} + +async function markRecoveryReleasing( + authz: AuthorizedRequest, + config: GatewayConfig, + reason: string, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const fenceId = requirePaymentRecoveryFence(authz) + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + state: 'releasing', + payment: recoveryTarget(authz, record.payment), + reason, + nextAttemptAt: Date.now(), + })) +} + +async function markRecoveryReconciled( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const fenceId = requirePaymentRecoveryFence(authz) + const now = Date.now() + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + state: 'reconciled', + payment: recoveryTarget(authz, record.payment), + lease: undefined, + lastError: undefined, + nextAttemptAt: Number.MAX_SAFE_INTEGER, + reconciledAt: now, + }), now) +} + export async function reclaimPayment( operationId: string, config: GatewayConfig, @@ -835,6 +1153,8 @@ export async function* dispatchSandboxStreamRich( if (signal?.aborted) return await onExecutionStart?.() if (signal?.aborted) return + // A stream adapter can start paid work synchronously during this call. + onSandboxStart?.() const promptStream = box.streamPrompt(userMessage, { sessionId: sessionId ?? `consumer:${consumerId}`, systemPrompt: agent.systemPrompt, @@ -842,7 +1162,6 @@ export async function* dispatchSandboxStreamRich( executionBudget, signal, }) - onSandboxStart?.() const iterator = promptStream[Symbol.asyncIterator]() try { while (true) { @@ -981,6 +1300,10 @@ export interface SettleAndRecordOptions { usageAlreadyRecorded?: boolean /** Persist the caller's recovery marker after attribution succeeds. */ onUsageRecorded?: () => Promise + /** Recovery uses the original quoted ceiling when no receipt arrives. */ + settlementBasis?: PaymentSettlementBasis + /** Exact base-unit charge selected by the recovery policy. */ + paymentAmount?: bigint } export async function settleAndRecord( @@ -991,6 +1314,9 @@ export async function settleAndRecord( obs: GatewayObserver | undefined, options: SettleAndRecordOptions = {}, ): Promise { + const settlementBasis = options.settlementBasis ?? 'usage-receipt' + await markRecoverySettling(authz, usage, settlementBasis, config) + if (options.usageAlreadyRecorded) await markRecoveryUsageRecorded(authz, config) const tokenCost = ( usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens ) * agent.pricePerTokenUsd @@ -1013,6 +1339,7 @@ export async function settleAndRecord( ownerEarnedUsd: ownerEarned, platformFeeUsd: platformFee, durationMs: Date.now() - authz.startMs, + settlementBasis, } const ctx: RequestContext = { requestId: authz.requestId, @@ -1021,7 +1348,7 @@ export async function settleAndRecord( } try { if (authz.paymentOperation && config.x402.paymentOperations) { - const amount = actualX402Amount( + const amount = options.paymentAmount ?? actualX402Amount( agent.pricePerTokenUsd, usage.inputTokens, usage.outputTokens, @@ -1032,13 +1359,22 @@ export async function settleAndRecord( ) authz.paymentOperation = await config.x402.paymentOperations.settlePayment( authz.paymentOperation, - { amount, totalCostUsd: totalCost, usage }, + { amount, totalCostUsd: totalCost, usage, basis: settlementBasis }, ) // Durable settlement happens first. If attribution storage is // unavailable, recovery must never refund delivered work. if (!options.usageAlreadyRecorded) { await config.recordUsage(usageEvent) await options.onUsageRecorded?.() + await markRecoveryUsageRecorded(authz, config) + } + } else if (authz.mppChargeOperation) { + // Generic MPP charge methods confirm before the response. Finalization + // records attribution only; it never invokes the legacy settlement hook. + if (!options.usageAlreadyRecorded) { + await config.recordUsage(usageEvent) + await options.onUsageRecorded?.() + await markRecoveryUsageRecorded(authz, config) } } else { // Legacy adapters retain attribution-before-charge because their @@ -1058,7 +1394,7 @@ export async function settleAndRecord( ) } } - await obs?.onRequestComplete?.(ctx, usageEvent) + await markRecoveryReconciled(authz, config) } catch (err) { const msg = err instanceof Error ? err.message : String(err) console.error(`[agent-gateway] settlement failed for ${authz.consumerId}: ${msg}`) @@ -1069,6 +1405,49 @@ export async function settleAndRecord( }) throw err } + try { + await obs?.onRequestComplete?.(ctx, usageEvent) + } catch (error) { + console.error( + `[agent-gateway] completion observer failed for ${authz.requestId}:`, + error instanceof Error ? error.message : String(error), + ) + } +} + +async function markRecoverySettling( + authz: AuthorizedRequest, + usage: SandboxUsageReceipt, + settlementBasis: PaymentSettlementBasis, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const fenceId = requirePaymentRecoveryFence(authz) + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => { + return { + ...record, + state: 'settling', + payment: recoveryTarget(authz, record.payment), + workStarted: true, + usage, + settlementBasis, + nextAttemptAt: Date.now(), + } + }) +} + +async function markRecoveryUsageRecorded( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const fenceId = requirePaymentRecoveryFence(authz) + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + usageRecorded: true, + })) } function mergeUsage( diff --git a/src/index.ts b/src/index.ts index 2684d1c..9bf2506 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,5 +1,12 @@ export { createAgentGateway } from './middleware' export { reclaimPayment } from './dispatch' +export { + recoverPayment, + recoverPayments, + type RecoverPaymentOptions, + type RecoverPaymentsOptions, + type PaymentRecoveryRun, +} from './payment-recovery-worker' export { verifyX402, verifyMpp, @@ -7,7 +14,32 @@ export { isApiKeyAuthEnabled, isMppAuthEnabled, mppReplayNonceKey, + mppPaymentCredential, + type VerifiedMppCredential, } from './verify' +export { + MPP_CHARGE_PROTOCOL_VERSION, + mppPaymentOperationId, + type MppAuthenticatedCredential, + type MppChargeLifecycle, + type MppChargeOperation, + type MppChargeOperationState, + type MppChargeRecoveryResult, + type MppChargeRequest, +} from './mpp-payment' +export { + PAYMENT_RECOVERY_VERSION, + MemoryPaymentRecoveryStore, + PaymentRecoveryFenceError, + type PaymentRecoveryAttribution, + type PaymentRecoveryConfig, + type PaymentRecoveryRecord, + type PaymentRecoveryState, + type PaymentRecoveryStore, + type PaymentRecoveryTarget, + type PaymentSettlementBasis, +} from './payment-recovery' +export { SqlPaymentRecoveryStore } from './payment-recovery-sql' export { PAYMENT_PROTOCOL_VERSION, MemoryPaymentOperations, diff --git a/src/middleware.ts b/src/middleware.ts index 70656f5..b7def31 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -15,6 +15,11 @@ import { } from './dispatch' import { MemoryNonceStore } from './nonce-store' import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' +import { + MemoryPaymentRecoveryStore, + PaymentRecoveryReplayError, + assertPaymentRecoveryConfig, +} from './payment-recovery' import { MemoryRateLimitStore, type RateLimitStore } from './rate-limit' import type { ChatCompletionChunk, ChatCompletionRequest, GatewayConfig } from './types' import { isApiKeyAuthEnabled, isMppAuthEnabled } from './verify' @@ -29,7 +34,8 @@ import { isApiKeyAuthEnabled, isMppAuthEnabled } from './verify' * GET /:slug/chat/completions — agent discovery metadata (Tangle-native shape) * POST /:slug/chat/completions — OpenAI-compatible chat endpoint (paid) */ -export function createAgentGateway(config: GatewayConfig) { +export function createAgentGateway(inputConfig: GatewayConfig) { + let config = inputConfig // Production gateways must verify x402 signatures. Tests and local // dev can opt into the explicit demo path. if (!config.x402.verifySigner && !config.x402.demoMode) { @@ -86,6 +92,37 @@ export function createAgentGateway(config: GatewayConfig) { if (config.x402.paymentProtocolVersion === 1 && config.x402.paymentOperations) { throw new Error('createAgentGateway: version 1 cannot be combined with version 2 payment operations') } + const mppMethod = (config.mpp?.method ?? 'blueprintevm').toLowerCase() + if (config.mpp?.charge && config.mpp.charge.protocolVersion !== 1) { + throw new Error('createAgentGateway: unsupported MPP charge lifecycle version') + } + if ( + config.mpp?.charge && + mppMethod !== 'blueprintevm' && + !config.mpp.authenticateCredential + ) { + throw new Error('createAgentGateway: generic MPP methods require pure credential authentication') + } + if ( + config.mpp && + mppMethod !== 'blueprintevm' && + config.mpp.authenticateCredential && + !config.mpp.charge + ) { + throw new Error('createAgentGateway: generic MPP methods require a charge lifecycle') + } + const needsRecovery = config.x402.paymentProtocolVersion === 2 || + (mppMethod !== 'blueprintevm' && config.mpp?.charge !== undefined) + if (needsRecovery && !config.paymentRecovery) { + if (!config.x402.demoMode) { + throw new Error('createAgentGateway: durable payment recovery is required in production') + } + config = { + ...config, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + } + } + if (config.paymentRecovery) assertPaymentRecoveryConfig(config.paymentRecovery) const gw = new Hono() const rateLimitStore: RateLimitStore = config.rateLimitStore ?? new MemoryRateLimitStore() const state: GatewayState = { @@ -194,7 +231,10 @@ export function createAgentGateway(config: GatewayConfig) { const authz = guard try { await claimPayment(authz, config, state) - } catch { + } catch (error) { + const replayedGenericMpp = error instanceof PaymentRecoveryReplayError && + authz.paymentMethod === 'mpp' && + authz.mppMethod !== 'blueprintevm' try { await releasePayment(authz, config, 'payment authorization failed') } catch (releaseError) { @@ -209,14 +249,50 @@ export function createAgentGateway(config: GatewayConfig) { agentSlug: authz.agent.slug, startMs: authz.startMs, }, - { method: authz.paymentMethod, code: 'payment_authorization_failed', httpStatus: 402 }, + { + method: authz.paymentMethod, + code: replayedGenericMpp ? 'invalid_mpp_credential' : 'payment_authorization_failed', + httpStatus: replayedGenericMpp ? 401 : 402, + }, ) + const status = replayedGenericMpp ? 401 : 402 + return c.json( + { + error: { + message: replayedGenericMpp ? 'Invalid Payment credential' : 'Payment authorization failed', + type: replayedGenericMpp ? 'authentication_error' : 'payment_required', + code: replayedGenericMpp ? 'invalid_mpp_credential' : 'payment_authorization_failed', + }, + }, + replayedGenericMpp + ? { + status, + headers: { + 'WWW-Authenticate': `Payment realm="${config.mpp!.realm}", method="${config.mpp!.method ?? 'blueprintevm'}"`, + 'X-Request-Id': authz.requestId, + }, + } + : { status, headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': authz.requestId } }, + ) + } + + try { + await beginPaymentExecution(authz, config) + } catch { + try { + await releasePayment(authz, config, 'payment execution fence was lost') + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } return c.json( { error: { message: 'Payment authorization failed', type: 'payment_required', - code: 'payment_authorization_failed', + code: 'payment_execution_fence_lost', }, }, { status: 402, headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': authz.requestId } }, @@ -309,12 +385,10 @@ function streamChatCompletions( abortController.signal, undefined, maxOutputTokens, - async () => { - await beginPaymentExecution(authz, config) - }, - authz.paymentOperation !== undefined, + undefined, + authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, () => { - if (authz.paymentOperation) workObserved = true + if (authz.paymentRecoveryId) workObserved = true }, )) { if (event.kind === 'text') { @@ -394,7 +468,13 @@ function streamChatCompletions( 'X-Agent-Slug': agent.slug, 'X-Agent-Hosting': agent.sandboxEndpoint ? 'sovereign' : 'centralized', 'X-Payment-Method': paymentMethod, - 'X-Payment-Settled': paymentMethod === 'x402' ? 'pending' : 'true', + 'X-Payment-Settled': paymentMethod === 'x402' || authz.paymentOperation ? 'pending' : 'true', + ...(authz.mppChargeOperation + ? { 'Payment-Receipt': authz.mppChargeOperation.receipt } + : {}), + ...(authz.paymentRecoveryId + ? { 'X-Payment-Operation-Id': authz.paymentRecoveryId } + : {}), ...(rateLimitRemaining !== undefined ? { 'X-RateLimit-Remaining': String(rateLimitRemaining) } : {}), diff --git a/src/mpp-payment.ts b/src/mpp-payment.ts new file mode 100644 index 0000000..2f6f431 --- /dev/null +++ b/src/mpp-payment.ts @@ -0,0 +1,117 @@ +/** Version of the gateway's method-specific MPP charge contract. */ +export const MPP_CHARGE_PROTOCOL_VERSION = 1 as const + +export type MppChargeOperationState = 'confirmed' | 'releasing' | 'released' + +/** Pure authentication result for one method credential. */ +export interface MppAuthenticatedCredential { + consumerId: string + /** + * Stable, non-secret processor identity for this payment credential. + * Return the same value for equivalent encodings of one credential. + * The gateway hashes this value before it persists or claims it. + */ + paymentIdentity: string +} + +/** Durable result of one immediate MPP charge. */ +export interface MppChargeOperation { + protocolVersion: typeof MPP_CHARGE_PROTOCOL_VERSION + operationId: string + acquiredByRequestId: string + method: string + /** A complete Payment-Receipt header value. */ + receipt: string + state: MppChargeOperationState +} + +export interface MppChargeRequest { + /** + * Stable provider idempotency key. The adapter must bind every processor + * operation to this value before it attempts confirmation. + */ + operationId: string + requestId: string + agentId: string + consumerId: string + method: string + /** Original decoded credential. It is available only on the live request. */ + credential: string + amount: bigint + currencyDecimals: number +} + +export type MppChargeRecoveryResult = + | MppChargeOperation + /** `not-found` is final and must fence this operation ID against a later charge. */ + | { operationId: string; state: 'not-found' | 'pending' } + +/** + * Immediate-charge lifecycle for a non-BlueprinTEVM MPP method. + * + * `confirmPayment` runs after every request denial and before a response or + * sandbox call. It must use `operationId` as its processor idempotency key, + * confirm payment, verify final success, and only then return `confirmed`. + * + * Every method must also support id-only recovery and an idempotent release. + * Recovery must inspect the existing processor operation. It must never + * create a second charge when an acknowledgement is ambiguous. + */ +export interface MppChargeLifecycle { + readonly protocolVersion: typeof MPP_CHARGE_PROTOCOL_VERSION + confirmPayment(request: MppChargeRequest): Promise + releasePayment(operation: MppChargeOperation, reason: string): Promise + recoverPayment(operationId: string): Promise +} + +/** Stable gateway identity. Neither the credential nor adapter identity is persisted. */ +export async function mppPaymentOperationId( + method: string, + paymentIdentity: string, +): Promise { + const normalizedMethod = method.trim().toLowerCase() + if (!normalizedMethod || normalizedMethod.length > 128) { + throw new Error('MPP payment method is invalid') + } + if (!paymentIdentity || paymentIdentity.length > 8192) { + throw new Error('MPP payment identity is invalid') + } + const digest = await globalThis.crypto.subtle.digest( + 'SHA-256', + new TextEncoder().encode(`${normalizedMethod}\0${paymentIdentity}`), + ) + const fingerprint = [...new Uint8Array(digest)] + .map((byte) => byte.toString(16).padStart(2, '0')) + .join('') + return `mpp:${normalizedMethod}:${fingerprint}` +} + +export function assertMppChargeOperation( + operation: MppChargeOperation, + expected: { operationId: string; requestId: string; method: string }, + allowedStates: readonly MppChargeOperationState[], + requireReceipt = true, +): void { + if (operation.protocolVersion !== MPP_CHARGE_PROTOCOL_VERSION) { + throw new Error('MPP charge operation protocol version mismatch') + } + if (operation.operationId !== expected.operationId) { + throw new Error('MPP charge operation id mismatch') + } + if (operation.acquiredByRequestId !== expected.requestId) { + throw new Error('MPP charge operation request owner mismatch') + } + if (operation.method.toLowerCase() !== expected.method.toLowerCase()) { + throw new Error('MPP charge operation method mismatch') + } + if (!allowedStates.includes(operation.state)) { + throw new Error(`MPP charge operation is in invalid state ${operation.state}`) + } + if (requireReceipt && ( + !operation.receipt || + operation.receipt.length > 8192 || + /[^\x20-\x7e]/.test(operation.receipt) + )) { + throw new Error('MPP charge operation has an invalid payment receipt') + } +} diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 7c8d5bf..77f10ce 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -1,4 +1,5 @@ import type { SandboxUsageReceipt } from './types' +import type { PaymentSettlementBasis } from './payment-recovery' /** Version negotiated by gateways that use durable payment operations. */ export const PAYMENT_PROTOCOL_VERSION = 2 as const @@ -51,6 +52,8 @@ export interface PaymentSettlementInput { amount: bigint totalCostUsd: number usage: SandboxUsageReceipt + /** Distinguishes a provider receipt from the bounded missing-receipt fallback. */ + basis: PaymentSettlementBasis } /** @@ -260,6 +263,7 @@ export class MemoryPaymentOperations implements PaymentOperations { const recovery = this.recoverSettlement(current, { amount: current.settledAmount, totalCostUsd: 0, + basis: 'usage-receipt', usage: { inputTokens: 0, outputTokens: 0, @@ -409,4 +413,7 @@ function validateSettlement(operation: PaymentOperation, input: PaymentSettlemen throw new Error('settlement cost must be finite and non-negative') } if (!input.usage.budgetEnforced) throw new Error('sandbox usage receipt is not budget-enforced') + if (input.basis !== 'usage-receipt' && input.basis !== 'quoted-ceiling') { + throw new Error('payment settlement basis is invalid') + } } diff --git a/src/payment-recovery-sql.ts b/src/payment-recovery-sql.ts new file mode 100644 index 0000000..79c5fb8 --- /dev/null +++ b/src/payment-recovery-sql.ts @@ -0,0 +1,108 @@ +import type { SqlAdapter } from './a2a/task-store-sql' +import type { + PaymentRecoveryRecord, + PaymentRecoveryStore, +} from './payment-recovery' + +const TABLE_DDL = (table: string) => ` + CREATE TABLE IF NOT EXISTS ${table} ( + id TEXT PRIMARY KEY, + state TEXT NOT NULL, + next_attempt_at INTEGER NOT NULL, + revision INTEGER NOT NULL, + payload TEXT NOT NULL, + updated_at INTEGER NOT NULL + ) +` + +const DUE_INDEX_DDL = (table: string) => ` + CREATE INDEX IF NOT EXISTS idx_${table}_due + ON ${table} (state, next_attempt_at) +` + +/** Durable recovery outbox for D1, sqlite, libSQL, or an adapted SQL driver. */ +export class SqlPaymentRecoveryStore implements PaymentRecoveryStore { + private readonly table: string + + constructor( + private readonly db: SqlAdapter, + options: { table?: string } = {}, + ) { + this.table = options.table ?? 'payment_recovery' + if (!/^[A-Za-z_][A-Za-z0-9_]*$/.test(this.table)) { + throw new Error('payment recovery table name is invalid') + } + } + + /** Idempotent. Run this before the gateway starts accepting traffic. */ + async migrate(): Promise { + await this.db.exec(TABLE_DDL(this.table)) + await this.db.exec(DUE_INDEX_DDL(this.table)) + } + + async createIfAbsent(record: PaymentRecoveryRecord): Promise { + try { + const result = await this.db.exec( + `INSERT INTO ${this.table} (id, state, next_attempt_at, revision, payload, updated_at) VALUES (?, ?, ?, ?, ?, ?)`, + [ + record.id, + record.state, + record.nextAttemptAt, + record.revision, + JSON.stringify(record), + record.updatedAt, + ], + ) + return result.rowsAffected === 1 + } catch (error) { + if (await this.get(record.id)) return false + throw error + } + } + + async get(id: string): Promise { + const rows = await this.db.query<{ payload: string }>( + `SELECT payload FROM ${this.table} WHERE id = ?`, + [id], + ) + return rows[0] ? parseRecord(rows[0].payload) : undefined + } + + async compareAndSet( + expected: PaymentRecoveryRecord, + next: PaymentRecoveryRecord, + ): Promise { + const result = await this.db.exec( + `UPDATE ${this.table} SET state = ?, next_attempt_at = ?, revision = ?, payload = ?, updated_at = ? WHERE id = ? AND revision = ?`, + [ + next.state, + next.nextAttemptAt, + next.revision, + JSON.stringify(next), + next.updatedAt, + expected.id, + expected.revision, + ], + ) + return result.rowsAffected === 1 + } + + async listDue(now: number, limit: number): Promise { + if (!Number.isSafeInteger(limit) || limit <= 0) { + throw new Error('payment recovery scan limit must be a positive safe integer') + } + const rows = await this.db.query<{ payload: string }>( + `SELECT payload FROM ${this.table} WHERE state <> ? AND next_attempt_at <= ? ORDER BY next_attempt_at ASC LIMIT ?`, + ['reconciled', now, limit], + ) + return rows.map((row) => parseRecord(row.payload)) + } +} + +function parseRecord(payload: string): PaymentRecoveryRecord { + const value = JSON.parse(payload) as PaymentRecoveryRecord + if (value.version !== 1 || !value.id || !Number.isSafeInteger(value.revision)) { + throw new Error('stored payment recovery record is invalid') + } + return value +} diff --git a/src/payment-recovery-worker.ts b/src/payment-recovery-worker.ts new file mode 100644 index 0000000..1496881 --- /dev/null +++ b/src/payment-recovery-worker.ts @@ -0,0 +1,457 @@ +import { + type AuthorizedRequest, + releasePayment, + settleAndRecord, +} from './dispatch' +import { assertMppChargeOperation } from './mpp-payment' +import { + deserializePaymentOperation, + PaymentRecoveryFenceError, + recoveryTiming, + updateOwnedPaymentRecovery, + type PaymentRecoveryRecord, +} from './payment-recovery' +import type { AgentMeta, GatewayConfig, SandboxUsageReceipt } from './types' + +interface RecoveryWorkerOptions { + now?: number + workerId?: string +} + +export interface RecoverPaymentOptions extends RecoveryWorkerOptions { + /** Process one requested row even when its normal retry time is later. */ + force?: boolean + /** Exact receipt recovered from a durable protocol record, such as an A2A task. */ + usage?: SandboxUsageReceipt +} + +/** Batch recovery never accepts receipt data because each receipt belongs to one row. */ +export interface RecoverPaymentsOptions extends RecoveryWorkerOptions { + limit?: number +} + +export interface PaymentRecoveryRun { + scanned: number + reconciled: number + deferred: number + failed: number +} + +/** Scan and reconcile due payment rows. Safe to run concurrently on many workers. */ +export async function recoverPayments( + config: GatewayConfig, + options: RecoverPaymentsOptions = {}, +): Promise { + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const now = options.now ?? Date.now() + const limit = options.limit ?? 100 + if (!Number.isSafeInteger(limit) || limit <= 0) { + throw new Error('payment recovery limit must be a positive safe integer') + } + const due = await recovery.store.listDue(now, limit) + const run: PaymentRecoveryRun = { scanned: due.length, reconciled: 0, deferred: 0, failed: 0 } + for (const candidate of due) { + try { + const result = await recoverPayment(candidate.id, config, { + now, + force: false, + ...(options.workerId ? { workerId: options.workerId } : {}), + }) + if (result?.state === 'reconciled') run.reconciled += 1 + else run.deferred += 1 + } catch { + run.failed += 1 + } + } + return run +} + +/** Reconcile one payment identity. Hosts can expose this through a private worker API. */ +export async function recoverPayment( + recoveryId: string, + config: GatewayConfig, + options: RecoverPaymentOptions = {}, +): Promise { + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const now = options.now ?? Date.now() + const current = await recovery.store.get(recoveryId) + if (!current || current.state === 'reconciled') return current + if (!options.force && current.nextAttemptAt > now) return current + + const leased = await acquireLease( + current.id, + config, + options.workerId ? `${options.workerId}:${randomId()}` : randomId(), + now, + ) + if (!leased) return recovery.store.get(recoveryId) + const fenceId = requireFence(leased) + + try { + const ready = options.usage + ? await persistRecoveredUsage(leased, options.usage, config, now) + : leased + await reconcileLeased(ready, config, now) + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + try { + await updateOwnedPaymentRecovery(recovery.store, recoveryId, fenceId, (record) => ({ + ...record, + attempts: record.attempts + 1, + lastError: message, + lease: undefined, + nextAttemptAt: now + recoveryTiming(recovery).retryDelayMs, + }), now) + } catch (updateError) { + if (!(updateError instanceof PaymentRecoveryFenceError)) throw updateError + } + throw error + } + return recovery.store.get(recoveryId) +} + +async function reconcileLeased( + record: PaymentRecoveryRecord, + config: GatewayConfig, + now: number, +): Promise { + if (record.payment.kind === 'mpp-charge' && !record.payment.operation) { + record = await recoverUnknownMppCharge(record, config) + } + if (record.state === 'reconciled') return + + if (record.state === 'claiming') { + if (record.payment.kind === 'x402') { + if (!config.x402.paymentOperations) { + throw new Error('x402 payment recovery is not configured') + } + const recovered = await config.x402.paymentOperations.reclaimPayment(record.payment.operationId) + if (recovered.state !== 'released' && recovered.state !== 'reclaimed') { + throw new Error(`ambiguous x402 claim recovered in state ${recovered.state}`) + } + await completeRecord(record, config) + return + } + throw new Error('MPP charge confirmation remains unresolved') + } + + if (record.state === 'claimed' && !record.workStarted) { + await releaseRecoveredPayment(record, config, 'request ended before sandbox execution') + return + } + + if (record.state === 'releasing') { + await releaseRecoveredPayment(record, config, record.reason ?? 'payment recovery release') + return + } + + if ( + record.state === 'executing' || + record.state === 'retained' || + record.state === 'settling' || + (record.state === 'claimed' && record.workStarted) + ) { + if (!record.usage && record.fallbackAt !== undefined && record.fallbackAt > now) { + await updateOwnedPaymentRecovery( + config.paymentRecovery!.store, + record.id, + requireFence(record), + (current) => ({ + ...current, + lease: undefined, + nextAttemptAt: record.fallbackAt!, + }), + now, + ) + return + } + await settleRecoveredPayment(record, config) + return + } + + throw new Error(`payment recovery cannot reconcile state ${record.state}`) +} + +async function recoverUnknownMppCharge( + record: PaymentRecoveryRecord, + config: GatewayConfig, +): Promise { + if (record.payment.kind !== 'mpp-charge') return record + const method = record.payment.method + const operationId = record.payment.operationId + const lifecycle = config.mpp?.charge + if (!lifecycle || lifecycle.protocolVersion !== 1) { + throw new Error('MPP charge recovery is not configured') + } + const result = await lifecycle.recoverPayment(operationId) + if (!('protocolVersion' in result)) { + if (result.operationId !== operationId) { + throw new Error('MPP recovery operation id mismatch') + } + if (result.state === 'not-found') { + await completeRecord(record, config) + return requireRecord(record.id, config) + } + throw new Error('MPP charge confirmation is still pending') + } + assertMppChargeOperation( + result, + { + operationId, + requestId: record.attribution.requestId, + method, + }, + ['confirmed', 'releasing', 'released'], + false, + ) + return updateOwnedPaymentRecovery( + config.paymentRecovery!.store, + record.id, + requireFence(record), + (current) => ({ + ...current, + state: result.state === 'released' + ? 'reconciled' + : result.state === 'releasing' + ? 'releasing' + : current.workStarted + ? 'executing' + : 'claimed', + payment: { kind: 'mpp-charge', method, operationId, operation: result }, + ...(result.state === 'released' + ? { reconciledAt: Date.now(), nextAttemptAt: Number.MAX_SAFE_INTEGER, lease: undefined } + : {}), + }), + ) +} + +async function releaseRecoveredPayment( + record: PaymentRecoveryRecord, + config: GatewayConfig, + reason: string, +): Promise { + const authz = authorizedRequest(record) + if (record.payment.kind === 'x402') { + if (!record.payment.operation) { + if (!config.x402.paymentOperations) throw new Error('x402 payment recovery is not configured') + const recovered = await config.x402.paymentOperations.reclaimPayment(record.payment.operationId) + if (recovered.state !== 'released' && recovered.state !== 'reclaimed') { + throw new Error(`x402 release recovered in state ${recovered.state}`) + } + await completeRecord(record, config) + return + } + authz.paymentOperation = deserializePaymentOperation(record.payment.operation) + authz.paymentOperationAcquired = true + } else { + if (!record.payment.operation) throw new Error('MPP release operation is missing') + authz.mppChargeOperation = record.payment.operation + } + await releasePayment(authz, config, reason) +} + +async function settleRecoveredPayment( + record: PaymentRecoveryRecord, + config: GatewayConfig, +): Promise { + const authz = authorizedRequest(record) + if (record.payment.kind === 'x402') { + if (!record.payment.operation) throw new Error('x402 settlement operation is missing') + authz.paymentOperation = deserializePaymentOperation(record.payment.operation) + authz.paymentOperationAcquired = true + } else { + if (!record.payment.operation) throw new Error('MPP settlement operation is missing') + authz.mppChargeOperation = record.payment.operation + } + + const fallback = !record.usage + const usage = record.usage ?? quotedCeilingUsage(record) + await settleAndRecord( + recoveryAgent(record), + authz, + usage, + config, + config.observer, + { + usageAlreadyRecorded: record.usageRecorded, + settlementBasis: fallback ? 'quoted-ceiling' : record.settlementBasis ?? 'usage-receipt', + ...(fallback && record.payment.kind === 'x402' + ? { paymentAmount: BigInt(record.attribution.requiredAmount) } + : {}), + }, + ) +} + +function quotedCeilingUsage(record: PaymentRecoveryRecord): SandboxUsageReceipt { + const amount = BigInt(record.attribution.requiredAmount) + const providerCostUsd = baseUnitsToNumber(amount, record.attribution.currencyDecimals) + if (!Number.isFinite(providerCostUsd) || providerCostUsd < 0) { + throw new Error('quoted payment ceiling cannot be represented as USD') + } + return { + inputTokens: 0, + outputTokens: 0, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd, + budgetEnforced: true, + } +} + +function baseUnitsToNumber(amount: bigint, decimals: number): number { + const digits = amount.toString().padStart(decimals + 1, '0') + if (decimals === 0) return Number(digits) + const split = digits.length - decimals + return Number(`${digits.slice(0, split)}.${digits.slice(split)}`) +} + +function authorizedRequest(record: PaymentRecoveryRecord): AuthorizedRequest { + const attribution = record.attribution + return { + agent: recoveryAgent(record), + consumerId: attribution.consumerId, + paymentMethod: attribution.paymentMethod, + keyInfo: null, + userMessage: '[payment recovery]', + rateLimitRemaining: undefined, + requestId: attribution.requestId, + startMs: attribution.startMs, + maxOutputTokens: attribution.maxOutputTokens, + executionBudget: attribution.executionBudget, + requiredPaymentAmount: BigInt(attribution.requiredAmount), + paymentPayload: null, + paymentRecoveryId: record.id, + paymentRecoveryFence: requireFence(record), + } +} + +function recoveryAgent(record: PaymentRecoveryRecord): AgentMeta { + const attribution = record.attribution + return { + id: attribution.agentId, + ownerId: '[payment recovery]', + slug: attribution.agentSlug, + pricePerTokenUsd: attribution.pricePerTokenUsd, + platformFeePercent: attribution.platformFeePercent, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, + } +} + +async function acquireLease( + id: string, + config: GatewayConfig, + workerId: string, + now: number, +): Promise { + const recovery = config.paymentRecovery! + for (let attempt = 0; attempt < 16; attempt += 1) { + const current = await recovery.store.get(id) + if (!current || current.state === 'reconciled') return current + if (current.lease && current.lease.expiresAt > now) return undefined + const next: PaymentRecoveryRecord = { + ...current, + revision: current.revision + 1, + lease: { id: workerId, expiresAt: now + recoveryTiming(recovery).leaseMs }, + updatedAt: now, + } + if (await recovery.store.compareAndSet(current, next)) return next + } + throw new Error(`payment recovery record ${id} changed too many times`) +} + +async function completeRecord(record: PaymentRecoveryRecord, config: GatewayConfig): Promise { + const now = Date.now() + await updateOwnedPaymentRecovery( + config.paymentRecovery!.store, + record.id, + requireFence(record), + (current) => ({ + ...current, + state: 'reconciled', + lease: undefined, + lastError: undefined, + nextAttemptAt: Number.MAX_SAFE_INTEGER, + reconciledAt: now, + }), + now, + ) +} + +async function requireRecord(id: string, config: GatewayConfig): Promise { + const record = await config.paymentRecovery!.store.get(id) + if (!record) throw new Error(`payment recovery record ${id} was not found`) + return record +} + +function randomId(): string { + return globalThis.crypto.randomUUID() +} + +function requireFence(record: PaymentRecoveryRecord): string { + if (!record.lease?.id) throw new Error(`payment recovery record ${record.id} has no owner fence`) + return record.lease.id +} + +async function persistRecoveredUsage( + record: PaymentRecoveryRecord, + usage: SandboxUsageReceipt, + config: GatewayConfig, + now: number, +): Promise { + assertRecoveryUsage(usage) + return updateOwnedPaymentRecovery( + config.paymentRecovery!.store, + record.id, + requireFence(record), + (current) => { + if (current.usage && !sameUsage(current.usage, usage)) { + throw new Error('payment recovery receipt does not match persisted usage') + } + return { + ...current, + state: 'settling', + workStarted: true, + usage, + settlementBasis: 'usage-receipt', + nextAttemptAt: now, + } + }, + now, + ) +} + +function assertRecoveryUsage(usage: SandboxUsageReceipt): void { + for (const [name, value] of [ + ['inputTokens', usage.inputTokens], + ['outputTokens', usage.outputTokens], + ['reasoningTokens', usage.reasoningTokens], + ['toolTokens', usage.toolTokens], + ['toolCallCount', usage.toolCallCount], + ] as const) { + if (!Number.isSafeInteger(value) || value < 0) { + throw new Error(`payment recovery ${name} must be a non-negative safe integer`) + } + } + if (!Number.isFinite(usage.providerCostUsd) || usage.providerCostUsd < 0) { + throw new Error('payment recovery providerCostUsd must be finite and non-negative') + } + if (usage.budgetEnforced !== true) { + throw new Error('payment recovery receipt must confirm budget enforcement') + } +} + +function sameUsage(left: SandboxUsageReceipt, right: SandboxUsageReceipt): boolean { + return left.inputTokens === right.inputTokens && + left.outputTokens === right.outputTokens && + left.reasoningTokens === right.reasoningTokens && + left.toolTokens === right.toolTokens && + left.toolCallCount === right.toolCallCount && + left.providerCostUsd === right.providerCostUsd && + left.budgetEnforced === right.budgetEnforced +} diff --git a/src/payment-recovery.ts b/src/payment-recovery.ts new file mode 100644 index 0000000..1f54b93 --- /dev/null +++ b/src/payment-recovery.ts @@ -0,0 +1,330 @@ +import type { MppChargeOperation } from './mpp-payment' +import type { PaymentOperation } from './payment-operations' +import type { + PaymentMethod, + SandboxExecutionBudget, + SandboxUsageReceipt, +} from './types' + +export const PAYMENT_RECOVERY_VERSION = 1 as const + +export type PaymentRecoveryState = + | 'claiming' + | 'claimed' + | 'executing' + | 'retained' + | 'settling' + | 'releasing' + | 'reconciled' + +export type PaymentSettlementBasis = 'usage-receipt' | 'quoted-ceiling' + +export interface SerializedPaymentOperation { + protocolVersion: 2 + operationId: string + acquiredByRequestId: string + executionStartedAt?: number + retentionReason?: string + nonceKey: string + authorizationId: string + reservedAmount: string + settledAmount: string + refundAmount: string + expiresAt: number + state: PaymentOperation['state'] +} + +export interface PaymentRecoveryAttribution { + requestId: string + agentId: string + agentSlug: string + consumerId: string + paymentMethod: PaymentMethod + startMs: number + pricePerTokenUsd: number + platformFeePercent: number + requiredAmount: string + currencyDecimals: number + maxOutputTokens: number + executionBudget: SandboxExecutionBudget +} + +export type PaymentRecoveryTarget = + | { + kind: 'x402' + operationId: string + operation?: SerializedPaymentOperation + } + | { + kind: 'mpp-charge' + method: string + operationId: string + operation?: MppChargeOperation + } + +/** Durable outbox row for one payment identity. Rows are never deleted here. */ +export interface PaymentRecoveryRecord { + version: typeof PAYMENT_RECOVERY_VERSION + id: string + revision: number + state: PaymentRecoveryState + payment: PaymentRecoveryTarget + attribution: PaymentRecoveryAttribution + workStarted: boolean + /** Earliest time a missing receipt may settle at the quoted ceiling. */ + fallbackAt?: number + usage?: SandboxUsageReceipt + usageRecorded: boolean + settlementBasis?: PaymentSettlementBasis + reason?: string + attempts: number + lastError?: string + nextAttemptAt: number + lease?: { id: string; expiresAt: number } + createdAt: number + updatedAt: number + reconciledAt?: number +} + +export interface PaymentRecoveryStore { + createIfAbsent(record: PaymentRecoveryRecord): Promise + get(id: string): Promise + compareAndSet( + expected: PaymentRecoveryRecord, + next: PaymentRecoveryRecord, + ): Promise + listDue(now: number, limit: number): Promise +} + +export interface PaymentRecoveryConfig { + store: PaymentRecoveryStore + /** Claimed payment with no execution becomes recoverable after this delay. */ + staleRequestMs?: number + /** Work without a final receipt settles at the quoted ceiling after this delay. */ + receiptTimeoutMs?: number + /** Failed provider recovery waits this long before its next attempt. */ + retryDelayMs?: number + /** One recovery worker owns a row for this duration. */ + leaseMs?: number +} + +export const DEFAULT_STALE_REQUEST_MS = 30_000 +export const DEFAULT_RECEIPT_TIMEOUT_MS = 5 * 60_000 +export const DEFAULT_RECOVERY_RETRY_MS = 30_000 +export const DEFAULT_RECOVERY_LEASE_MS = 30_000 + +/** Atomic single-process store for tests and explicit local demo mode. */ +export class MemoryPaymentRecoveryStore implements PaymentRecoveryStore { + private readonly records = new Map() + + async createIfAbsent(record: PaymentRecoveryRecord): Promise { + if (this.records.has(record.id)) return false + this.records.set(record.id, clone(record)) + return true + } + + async get(id: string): Promise { + const record = this.records.get(id) + return record ? clone(record) : undefined + } + + async compareAndSet( + expected: PaymentRecoveryRecord, + next: PaymentRecoveryRecord, + ): Promise { + const current = this.records.get(expected.id) + if (!current || current.revision !== expected.revision) return false + this.records.set(expected.id, clone(next)) + return true + } + + async listDue(now: number, limit: number): Promise { + return [...this.records.values()] + .filter((record) => record.state !== 'reconciled' && record.nextAttemptAt <= now) + .sort((left, right) => left.nextAttemptAt - right.nextAttemptAt) + .slice(0, limit) + .map(clone) + } +} + +export class PaymentRecoveryFenceError extends Error { + constructor(id: string) { + super(`payment recovery fence was lost for ${id}`) + this.name = 'PaymentRecoveryFenceError' + } +} + +export class PaymentRecoveryReplayError extends Error { + constructor(id: string) { + super(`payment recovery identity was already reconciled: ${id}`) + this.name = 'PaymentRecoveryReplayError' + } +} + +/** Update only while the caller still owns the durable row fence. */ +export async function updateOwnedPaymentRecovery( + store: PaymentRecoveryStore, + id: string, + fenceId: string, + update: (record: PaymentRecoveryRecord) => PaymentRecoveryRecord, + now = Date.now(), +): Promise { + for (let attempt = 0; attempt < 16; attempt += 1) { + const current = await store.get(id) + if (!current) throw new Error(`payment recovery record ${id} was not found`) + if (current.state === 'reconciled' || current.lease?.id !== fenceId) { + throw new PaymentRecoveryFenceError(id) + } + const candidate = update(clone(current)) + assertRecoveryUpdate(current, candidate) + const next: PaymentRecoveryRecord = { + ...candidate, + id: current.id, + version: PAYMENT_RECOVERY_VERSION, + revision: current.revision + 1, + createdAt: current.createdAt, + updatedAt: now, + } + if (await store.compareAndSet(current, next)) return next + } + throw new Error(`payment recovery record ${id} changed too many times`) +} + +const PAYMENT_RECOVERY_TRANSITIONS: Record> = { + claiming: new Set(['claiming', 'claimed', 'releasing', 'reconciled']), + claimed: new Set(['claimed', 'executing', 'releasing', 'reconciled']), + executing: new Set(['executing', 'retained', 'settling', 'releasing', 'reconciled']), + retained: new Set(['retained', 'settling', 'reconciled']), + settling: new Set(['settling', 'reconciled']), + releasing: new Set(['releasing', 'reconciled']), + reconciled: new Set(['reconciled']), +} + +function assertRecoveryUpdate( + current: PaymentRecoveryRecord, + candidate: PaymentRecoveryRecord, +): void { + if (!PAYMENT_RECOVERY_TRANSITIONS[current.state].has(candidate.state)) { + throw new Error(`invalid payment recovery transition ${current.state} -> ${candidate.state}`) + } + if (JSON.stringify(current.attribution) !== JSON.stringify(candidate.attribution)) { + throw new Error('payment recovery attribution is immutable') + } + if ( + current.payment.kind !== candidate.payment.kind || + current.payment.operationId !== candidate.payment.operationId || + (current.payment.kind === 'mpp-charge' && + candidate.payment.kind === 'mpp-charge' && + current.payment.method !== candidate.payment.method) + ) { + throw new Error('payment recovery identity is immutable') + } + if (current.workStarted && !candidate.workStarted) { + throw new Error('payment recovery cannot forget started work') + } + if (current.usageRecorded && !candidate.usageRecorded) { + throw new Error('payment recovery cannot forget usage attribution') + } + if (current.usage && (!candidate.usage || !sameRecoveryUsage(current.usage, candidate.usage))) { + throw new Error('payment recovery usage receipt is immutable') + } + if ( + current.settlementBasis && + current.settlementBasis !== candidate.settlementBasis + ) { + throw new Error('payment recovery settlement basis is immutable') + } +} + +function sameRecoveryUsage( + left: SandboxUsageReceipt, + right: SandboxUsageReceipt, +): boolean { + return left.inputTokens === right.inputTokens && + left.outputTokens === right.outputTokens && + left.reasoningTokens === right.reasoningTokens && + left.toolTokens === right.toolTokens && + left.toolCallCount === right.toolCallCount && + left.providerCostUsd === right.providerCostUsd && + left.budgetEnforced === right.budgetEnforced +} + +export function serializePaymentOperation( + operation: PaymentOperation, +): SerializedPaymentOperation { + return { + protocolVersion: 2, + operationId: operation.operationId, + acquiredByRequestId: operation.acquiredByRequestId, + ...(operation.executionStartedAt !== undefined + ? { executionStartedAt: operation.executionStartedAt } + : {}), + ...(operation.retentionReason ? { retentionReason: operation.retentionReason } : {}), + nonceKey: operation.nonceKey, + authorizationId: operation.authorizationId, + reservedAmount: operation.reservedAmount.toString(), + settledAmount: operation.settledAmount.toString(), + refundAmount: operation.refundAmount.toString(), + expiresAt: operation.expiresAt, + state: operation.state, + } +} + +export function deserializePaymentOperation( + value: SerializedPaymentOperation, +): PaymentOperation { + if (value.protocolVersion !== 2) throw new Error('unsupported payment operation version') + if (!value.operationId || !value.acquiredByRequestId || !value.nonceKey || !value.authorizationId) { + throw new Error('incomplete payment operation recovery record') + } + return { + protocolVersion: 2, + operationId: value.operationId, + acquiredByRequestId: value.acquiredByRequestId, + ...(value.executionStartedAt !== undefined ? { executionStartedAt: value.executionStartedAt } : {}), + ...(value.retentionReason ? { retentionReason: value.retentionReason } : {}), + nonceKey: value.nonceKey, + authorizationId: value.authorizationId, + reservedAmount: BigInt(value.reservedAmount), + settledAmount: BigInt(value.settledAmount), + refundAmount: BigInt(value.refundAmount), + expiresAt: value.expiresAt, + state: value.state, + } +} + +export function recoveryTiming(config: PaymentRecoveryConfig): { + staleRequestMs: number + receiptTimeoutMs: number + retryDelayMs: number + leaseMs: number +} { + return { + staleRequestMs: config.staleRequestMs ?? DEFAULT_STALE_REQUEST_MS, + receiptTimeoutMs: config.receiptTimeoutMs ?? DEFAULT_RECEIPT_TIMEOUT_MS, + retryDelayMs: config.retryDelayMs ?? DEFAULT_RECOVERY_RETRY_MS, + leaseMs: config.leaseMs ?? DEFAULT_RECOVERY_LEASE_MS, + } +} + +export function assertPaymentRecoveryConfig(config: PaymentRecoveryConfig): void { + for (const [name, value] of Object.entries(recoveryTiming(config))) { + if (!Number.isSafeInteger(value) || value <= 0) { + throw new Error(`payment recovery ${name} must be a positive safe integer`) + } + } + const store = config.store + if ( + !store || + typeof store.createIfAbsent !== 'function' || + typeof store.get !== 'function' || + typeof store.compareAndSet !== 'function' || + typeof store.listDue !== 'function' + ) { + throw new Error('payment recovery store must provide atomic create, compare-and-set, and due scans') + } +} + +function clone(value: T): T { + return JSON.parse(JSON.stringify(value)) as T +} diff --git a/src/types.ts b/src/types.ts index a36b8a4..182ba9c 100644 --- a/src/types.ts +++ b/src/types.ts @@ -3,6 +3,8 @@ import type { PaymentOperation, PaymentOperations, } from './payment-operations' +import type { MppAuthenticatedCredential, MppChargeLifecycle } from './mpp-payment' +import type { PaymentRecoveryConfig, PaymentSettlementBasis } from './payment-recovery' // --- Agent resolution --- @@ -134,17 +136,16 @@ export interface MppConfig { /** MPP method name (default: "blueprintevm") */ method?: string /** - * Production verifier for the method-specific credential. Return the - * authenticated consumer id, or null when the credential is invalid. - * The callback receives the decoded JSON payload when one exists plus the - * original decoded credential so non-JSON methods can verify their own form. - * Omit only when x402.demoMode is explicitly enabled for local testing, or - * when x402.verifySigner handles an x402-compatible MPP credential. + * Pure credential authentication. Return stable method-owned identity, or null. + * This callback must not consume a credential, create a processor object, + * reserve funds, confirm payment, or perform any other financial mutation. */ - verifySigner?: ( + authenticateCredential?: ( payload: Record, context: { method: string; credential: string }, - ) => Promise + ) => Promise + /** Required immediate-charge lifecycle for every non-BlueprinTEVM method. */ + charge?: MppChargeLifecycle } export interface PaymentResult { @@ -203,6 +204,8 @@ export interface GatewayUsageEvent { ownerEarnedUsd: number platformFeeUsd: number durationMs: number + /** Exact receipt in normal operation; quoted ceiling only after receipt timeout. */ + settlementBasis?: PaymentSettlementBasis } // --- Sandbox interface --- @@ -287,6 +290,12 @@ export interface GatewayConfig { /** MPP (Machine Payments Protocol) configuration. It is advertised only when a production verifier or explicit demo mode is available. */ mpp?: MppConfig + /** + * Durable payment recovery outbox. Production payment protocol version 2 + * and generic MPP charge methods require this configuration. + */ + paymentRecovery?: PaymentRecoveryConfig + /** * Verify an API key. Return key info if valid, null if invalid. * In explicit x402 demo mode, the built-in verifier accepts `sk_agent_*` keys. diff --git a/src/verify.ts b/src/verify.ts index 4d77863..aaf6634 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -1,7 +1,16 @@ import type { X402Config, MppConfig, ApiKeyInfo, GatewayConfig } from './types' import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' +import { + mppPaymentOperationId, + type MppAuthenticatedCredential, +} from './mpp-payment' -/** Return the canonical opaque nonce key used by the final payment claim. */ +export interface VerifiedMppCredential extends MppAuthenticatedCredential { + /** Opaque replay key. BlueprinTEVM shares the x402 nonce namespace. */ + replayKey: string +} + +/** Return the legacy opaque nonce key used by older consumers. */ export function mppReplayNonceKey(authHeader: string): string | undefined { const decoded = decodeMppCredential(authHeader) return decoded ? canonicalMppNonceKey(decoded.method, decoded.payload, decoded.credential) : undefined @@ -12,6 +21,11 @@ export function mppPaymentPayload(authHeader: string): Record | return decodeMppCredential(authHeader)?.payload } +/** Return the decoded method credential for the post-guard charge lifecycle. */ +export function mppPaymentCredential(authHeader: string): string | undefined { + return decodeMppCredential(authHeader)?.credential +} + interface DecodedMppCredential { method: string credential: string @@ -50,20 +64,33 @@ function canonicalMppNonceKey( credential: string, ): string { if (payload.nonce === undefined) { - // Generic MPP methods may not expose a numeric nonce. The signed receipt - // itself is still the replay identity and must be claimed exactly once. return `mpp:${method.toLowerCase()}:receipt:${Buffer.from(credential).toString('base64url')}` } const nonce = BigInt(String(payload.nonce)).toString() const commitment = payload.commitment - // BlueprinTEVM carries the same SpendAuth identity as x402. Keep one - // namespace so a credential cannot cross the two HTTP transports. if (method.toLowerCase() === 'blueprintevm' && typeof commitment === 'string' && commitment.length > 0) { return `${commitment.toLowerCase()}:${nonce}` } return `mpp:${method.toLowerCase()}:${String(payload.commitment ?? payload.from ?? 'unknown').toLowerCase()}:${nonce}` } +function blueprintevmNonceKey(payload: Record): string | undefined { + if (payload.nonce === undefined) return undefined + const nonce = BigInt(String(payload.nonce)).toString() + const identity = payload.commitment ?? payload.from + if (typeof identity !== 'string' || identity.length === 0) return undefined + return `${identity.toLowerCase()}:${nonce}` +} + +function stableJson(value: unknown): string { + if (Array.isArray(value)) return `[${value.map(stableJson).join(',')}]` + if (value && typeof value === 'object') { + const record = value as Record + return `{${Object.keys(record).sort().map((key) => `${JSON.stringify(key)}:${stableJson(record[key])}`).join(',')}}` + } + return JSON.stringify(value) +} + /** Pure capability checks shared by discovery and every request protocol. */ export function isApiKeyAuthEnabled( config: Pick, @@ -71,17 +98,17 @@ export function isApiKeyAuthEnabled( return config.verifyApiKey !== undefined || config.x402.demoMode === true } -/** MPP is enabled only when a real verifier or explicit demo mode exists. */ +/** MPP is enabled only when authentication and method settlement are complete. */ export function isMppAuthEnabled( config: Pick, ): boolean { const method = (config.mpp?.method ?? 'blueprintevm').toLowerCase() - return Boolean( - config.mpp && - (config.mpp.verifySigner !== undefined || - (method === 'blueprintevm' && config.x402.verifySigner !== undefined) || - config.x402.demoMode === true), - ) + if (!config.mpp) return false + const authenticated = config.mpp.authenticateCredential !== undefined || + (method === 'blueprintevm' && config.x402.verifySigner !== undefined) || + (method === 'blueprintevm' && config.x402.demoMode === true) + if (!authenticated) return false + return method === 'blueprintevm' || config.mpp.charge !== undefined } /** @@ -150,11 +177,11 @@ export async function verifyX402( * Verify MPP (Machine Payments Protocol) Authorization: Payment header. * * MPP uses `Authorization: Payment ` format. The - * credential is method-specific; `MppConfig.verifySigner` owns verification - * and returns the consumer identity. The built-in `blueprintevm` path can + * credential is method-specific; `MppConfig.authenticateCredential` owns authentication + * and returns the consumer plus stable payment identity. The built-in `blueprintevm` path can * reuse the x402 verifier for credentials with the compatible payload shape. * - * Returns the signer address if valid, null otherwise. + * Returns authenticated identity if valid, null otherwise. * In demo mode, accepts any well-formed Payment header with an identity. */ export async function verifyMpp( @@ -164,23 +191,23 @@ export async function verifyMpp( nonceStore?: NonceStore, minimumAmount = 1n, markNonce = true, -): Promise { +): Promise { // MPP format: "Payment " const match = authHeader.match(/^Payment\s+(\S+)\s+(\S+)$/i) if (!match) return null const [, rawMethod] = match const method = rawMethod.toLowerCase() - if (config.method && method !== config.method.toLowerCase()) return null + if (method !== (config.method ?? 'blueprintevm').toLowerCase()) return null try { const decodedCredential = decodeMppCredential(authHeader) if (!decodedCredential) return null const { credential: decoded, payload } = decodedCredential - // Validate common EVM fields before a production verifier can reserve or - // settle funds. BlueprinTEVM carries x402-equivalent token amounts and - // must cover the same request ceiling as the X-Payment-Signature path. + // Validate common EVM fields before pure credential authentication. + // BlueprinTEVM carries x402-equivalent token amounts and must cover the + // same request ceiling as the X-Payment-Signature path. const operator = payload.operator ?? payload.to if (operator !== undefined) { if (typeof operator !== 'string' || operator.toLowerCase() !== x402Config.operatorAddress.toLowerCase()) { @@ -199,25 +226,49 @@ export async function verifyMpp( return null } - const nonceKey = canonicalMppNonceKey(method, payload, decoded) - if (nonceStore?.hasSeen && await nonceStore.hasSeen(nonceKey)) return null + const blueprintevmKey = method === 'blueprintevm' + ? blueprintevmNonceKey(payload) + : undefined + if (markNonce && blueprintevmKey && nonceStore?.hasSeen && await nonceStore.hasSeen(blueprintevmKey)) { + return null + } - let consumerId: string | null = null - if (config.verifySigner) { - consumerId = await config.verifySigner(payload, { method, credential: decoded }) + let authenticated: MppAuthenticatedCredential | null = null + if (config.authenticateCredential) { + authenticated = await config.authenticateCredential(payload, { method, credential: decoded }) } else if (method === 'blueprintevm' && x402Config.verifySigner && payload.commitment) { const verified = await x402Config.verifySigner(payload, { protocolVersion: x402Config.paymentProtocolVersion ?? (x402Config.paymentOperations ? 2 : 1), }) - consumerId = verified ? String(payload.commitment) : null + const paymentIdentity = blueprintevmNonceKey(payload) ?? stableJson(payload) + authenticated = verified + ? { consumerId: String(payload.commitment), paymentIdentity } + : null } else if (x402Config.demoMode) { const identity = payload.commitment ?? payload.from if (typeof identity !== 'string' || identity.length === 0) return null - consumerId = identity + const paymentIdentity = method === 'blueprintevm' + ? blueprintevmNonceKey(payload) ?? stableJson(payload) + : '' + if (!paymentIdentity) return null + authenticated = { consumerId: identity, paymentIdentity } } else { return null } - if (!consumerId) return null + if ( + !authenticated || + typeof authenticated.consumerId !== 'string' || + authenticated.consumerId.length === 0 || + typeof authenticated.paymentIdentity !== 'string' || + authenticated.paymentIdentity.length === 0 + ) return null + + const replayKey = method === 'blueprintevm' + ? blueprintevmKey ?? + await mppPaymentOperationId(method, authenticated.paymentIdentity) + : await mppPaymentOperationId(method, authenticated.paymentIdentity) + if (!replayKey) return null + if (markNonce && !blueprintevmKey && nonceStore?.hasSeen && await nonceStore.hasSeen(replayKey)) return null if (nonceStore && markNonce) { const expiry = payload.expiry === undefined @@ -225,11 +276,11 @@ export async function verifyMpp( : BigInt(String(payload.expiry)) const ttl = nonceTtlSeconds(expiry) if (ttl === undefined) return null - const claimed = await claimStoredNonce(nonceStore, nonceKey, ttl) + const claimed = await claimStoredNonce(nonceStore, replayKey, ttl) if (!claimed) return null } - return consumerId + return { ...authenticated, replayKey } } catch { return null } diff --git a/tests/a2a-atomicity.test.ts b/tests/a2a-atomicity.test.ts index e988a7c..5dbf022 100644 --- a/tests/a2a-atomicity.test.ts +++ b/tests/a2a-atomicity.test.ts @@ -5,6 +5,7 @@ import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { createAgentGateway } from '../src/middleware' import { MemoryNonceStore } from '../src/nonce-store' import { MemoryPaymentOperations, type PaymentOperation } from '../src/payment-operations' +import { MemoryPaymentRecoveryStore } from '../src/payment-recovery' import type { AgentMeta, GatewayConfig, SandboxBox, SandboxUsageReceipt } from '../src/types' import type { Artifact, Task } from '../src/a2a/types' @@ -274,6 +275,7 @@ describe('A2A task atomicity and restart recovery', () => { paymentOperations: operations, }, nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, a2a: { taskStore, authorizeTaskAccess: async () => true, @@ -385,6 +387,7 @@ describe('A2A task atomicity and restart recovery', () => { paymentOperations: operations, }, nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() diff --git a/tests/a2a-durability.test.ts b/tests/a2a-durability.test.ts index d5de8ac..8404976 100644 --- a/tests/a2a-durability.test.ts +++ b/tests/a2a-durability.test.ts @@ -69,7 +69,10 @@ function makeTaskAdapter(): SqlAdapter & { rows: Map } { return { rowsAffected: 1 } } if (s.startsWith('DELETE')) { - const [id] = params as [string] + const [id, expectedPayload] = params as [string, string | undefined] + if (expectedPayload !== undefined && rows.get(id)?.payload !== expectedPayload) { + return { rowsAffected: 0 } + } const had = rows.delete(id) return { rowsAffected: had ? 1 : 0 } } @@ -155,6 +158,25 @@ describe('SqlTaskStore', () => { expect(await store.get('t1')).toBeUndefined() }) + it.each([ + 'gatewayFinalizing', + 'gatewayPaymentRelease', + 'gatewayPaymentRecovery', + ])('does not expire a task while %s reconciliation is pending', async (key) => { + const db = makeTaskAdapter() + const store = new SqlTaskStore(db, { ttlMs: 1 }) + const task = { ...makeTask('t1'), metadata: { [key]: { version: 1 } } } + await store.put(task) + db.rows.get('t1')!.updated_at = Date.now() - 100 + + expect(await store.get('t1')).toEqual(task) + expect(db.rows.has('t1')).toBe(true) + + await store.put({ ...task, metadata: undefined }) + db.rows.get('t1')!.updated_at = Date.now() - 100 + expect(await store.get('t1')).toBeUndefined() + }) + it('delete removes the row', async () => { const db = makeTaskAdapter() const store = new SqlTaskStore(db) @@ -163,6 +185,18 @@ describe('SqlTaskStore', () => { expect(await store.get('t1')).toBeUndefined() }) + it('does not delete a task while payment reconciliation is pending', async () => { + const db = makeTaskAdapter() + const store = new SqlTaskStore(db) + const task = { + ...makeTask('t1'), + metadata: { gatewayPaymentRecovery: { version: 1, id: 'payment-1' } }, + } + await store.put(task) + await store.delete(task.id) + expect(await store.get(task.id)).toEqual(task) + }) + it('listByContext returns most-recent-first for tasks sharing a context', async () => { const db = makeTaskAdapter() const store = new SqlTaskStore(db) diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 1eef97a..9f00225 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -3,13 +3,23 @@ import { describe, expect, it } from 'vitest' import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { createAgentGateway } from '../src/middleware' +import type { MppChargeLifecycle, MppChargeOperation } from '../src/mpp-payment' import { MemoryNonceStore } from '../src/nonce-store' import { MemoryPaymentOperations, type PaymentOperation } from '../src/payment-operations' +import { MemoryPaymentRecoveryStore } from '../src/payment-recovery' +import { recoverPayment } from '../src/payment-recovery-worker' import type { AgentMeta, GatewayConfig, SandboxBox } from '../src/types' const operatorAddress = '0x1111111111111111111111111111111111111111' const commitment = `0x${'ab'.repeat(32)}` +function createTestGateway(config: GatewayConfig) { + return createAgentGateway({ + ...config, + paymentRecovery: config.paymentRecovery ?? { store: new MemoryPaymentRecoveryStore() }, + }) +} + const agent: AgentMeta = { id: 'agent-a2a-races', ownerId: 'owner', @@ -34,6 +44,53 @@ function paymentHeader(nonce: string): string { }) } +function stripePaymentHeader(token: string): string { + const credential = Buffer.from(JSON.stringify({ sharedPaymentToken: token })).toString('base64url') + return `Payment stripe ${credential}` +} + +class A2AChargeLifecycle implements MppChargeLifecycle { + readonly protocolVersion = 1 as const + readonly operations = new Map() + confirmations = 0 + releases = 0 + confirmationGate?: Promise + confirmationStarted?: () => void + loseConfirmationAcknowledgement = false + + async confirmPayment(request: Parameters[0]) { + this.confirmations += 1 + this.confirmationStarted?.() + await this.confirmationGate + const operation: MppChargeOperation = { + protocolVersion: 1, + operationId: request.operationId, + acquiredByRequestId: request.requestId, + method: request.method, + receipt: `stripe-receipt=${request.operationId}`, + state: 'confirmed', + } + this.operations.set(request.operationId, operation) + if (this.loseConfirmationAcknowledgement) { + throw new Error('Stripe confirmation acknowledgement lost') + } + return operation + } + + async releasePayment(operation: MppChargeOperation) { + const current = this.operations.get(operation.operationId) + if (current?.state === 'released') return current + this.releases += 1 + const released: MppChargeOperation = { ...operation, state: 'released' } + this.operations.set(operation.operationId, released) + return released + } + + async recoverPayment(operationId: string) { + return this.operations.get(operationId) ?? { operationId, state: 'not-found' as const } + } +} + function message(text: string, taskId: string) { return { kind: 'message', @@ -95,7 +152,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const continuation = app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('76') }, @@ -121,13 +178,92 @@ describe('A2A payment ownership races', () => { expect(canceled.status).toBe(200) finishAuthorization() const response = await continuation - const body = await response.json() as { error?: { code?: number } } + const body = await response.json() as { result?: { status?: { state?: string } } } - expect(body.error?.code).toBe(-32602) + expect(body.result?.status?.state).toBe('canceled') expect(runs).toBe(0) expect((await taskStore.get('task-payment-cancel'))?.status.state).toBe('canceled') }) + it('retains a continuation task until ambiguous Stripe confirmation reconciles', async () => { + const taskStore = new InMemoryTaskStore() + await taskStore.put({ + kind: 'task', + id: 'task-stripe-continuation-recovery', + contextId: 'ctx-race', + status: { state: 'input-required', timestamp: new Date().toISOString() }, + history: [message('initial', 'task-stripe-continuation-recovery')], + }) + const lifecycle = new A2AChargeLifecycle() + lifecycle.loseConfirmationAcknowledgement = true + const recoveryStore = new MemoryPaymentRecoveryStore() + let runs = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const continuation = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: stripePaymentHeader('spt_continuation_ack'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('continue', 'task-stripe-continuation-recovery') }, + }), + }) + const failed = await continuation.json() as { error?: { code?: number } } + const retained = await taskStore.get('task-stripe-continuation-recovery') + const marker = retained?.metadata?.gatewayPaymentRecovery as { id: string } + + expect(failed.error?.code).toBe(-32603) + expect(retained?.status.state).toBe('input-required') + expect(marker.id).toMatch(/^mpp:stripe:/) + await taskStore.delete('task-stripe-continuation-recovery') + expect(await taskStore.get('task-stripe-continuation-recovery')).toBeDefined() + expect(runs).toBe(0) + expect(lifecycle.confirmations).toBe(1) + + expect((await recoverPayment(marker.id, config, { force: true }))?.state).toBe('reconciled') + const fetched = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/get', + params: { id: 'task-stripe-continuation-recovery' }, + }), + }) + expect(fetched.status).toBe(200) + expect(lifecycle.releases).toBe(1) + expect((await taskStore.get('task-stripe-continuation-recovery'))?.metadata?.gatewayPaymentRecovery) + .toBeUndefined() + }) + it('releases payment when cancellation interrupts execution before sandbox start', async () => { const taskStore = new InMemoryTaskStore() let executionStarted!: () => void @@ -168,7 +304,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = app.request('/v1/agents/a2a-races', { method: 'POST', @@ -234,7 +370,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const request = (text: string) => app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json', Authorization: 'Bearer sk_agent_race' }, @@ -288,7 +424,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('75') }, @@ -366,7 +502,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('78') }, @@ -453,7 +589,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const stream = await app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('80') }, @@ -524,7 +660,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': paymentHeader('79') }, @@ -588,7 +724,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const streamPromise = app.request('/v1/agents/a2a-races', { method: 'POST', headers: { @@ -678,7 +814,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = app.request('/v1/agents/a2a-races', { method: 'POST', @@ -753,6 +889,7 @@ describe('A2A payment ownership races', () => { it('persists and recovers an ambiguous release acknowledgement', async () => { const taskStore = new InMemoryTaskStore() + const recoveryStore = new MemoryPaymentRecoveryStore() let executionStarted!: () => void const executionReady = new Promise((resolve) => { executionStarted = resolve }) let releaseExecution!: () => void @@ -796,10 +933,11 @@ describe('A2A payment ownership races', () => { paymentOperations: operations, }, nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore }, a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = app.request('/v1/agents/a2a-races', { method: 'POST', @@ -835,26 +973,11 @@ describe('A2A payment ownership races', () => { expect(releaseCalls).toBe(1) const retained = await taskStore.get('task-release-recovery') - const marker = retained?.metadata?.gatewayPaymentRelease as { - lease: { id: string; expiresAt: number } - operationId: string - recoveryAttempts?: number - } + const marker = retained?.metadata?.gatewayPaymentRecovery as { id: string } expect(retained?.status.state).toBe('canceled') - expect(marker.operationId).toBe(`x402:${commitment}:84`) - expect(marker.lease.expiresAt).toBeGreaterThan(Date.now()) - expect(marker.recoveryAttempts).toBe(1) - - await taskStore.put({ - ...retained!, - metadata: { - ...retained!.metadata, - gatewayPaymentRelease: { - ...marker, - lease: { ...marker.lease, expiresAt: Date.now() - 1 }, - }, - }, - }) + expect(marker.id).toBe(`x402:${commitment}:84`) + expect((await recoveryStore.get(marker.id))?.state).toBe('releasing') + expect((await recoverPayment(marker.id, config, { force: true }))?.state).toBe('reconciled') const recovered = await app.request('/v1/agents/a2a-races', { method: 'POST', headers: { 'Content-Type': 'application/json' }, @@ -869,7 +992,7 @@ describe('A2A payment ownership races', () => { result?: { status?: { state?: string } } } expect(recoveredBody.result?.status?.state).toBe('canceled') - expect((await taskStore.get('task-release-recovery'))?.metadata?.gatewayPaymentRelease).toBeUndefined() + expect((await taskStore.get('task-release-recovery'))?.metadata?.gatewayPaymentRecovery).toBeUndefined() expect(operations.get(`x402:${commitment}:84`)?.state).toBe('released') expect(releaseCalls).toBe(1) expect(recoveryCalls).toBe(1) @@ -877,6 +1000,7 @@ describe('A2A payment ownership races', () => { it('recovers an ambiguous release after shared nonce ownership fails', async () => { const taskStore = new InMemoryTaskStore() + const recoveryStore = new MemoryPaymentRecoveryStore() const nonceStore = { hasSeen: async () => false, claim: async () => false, @@ -917,10 +1041,11 @@ describe('A2A payment ownership races', () => { paymentOperations: operations, }, nonceStore, + paymentRecovery: { store: recoveryStore }, a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const response = await app.request('/v1/agents/a2a-races', { method: 'POST', @@ -937,25 +1062,17 @@ describe('A2A payment ownership races', () => { expect(sandboxStarted).toBe(false) expect(operations.get(`x402:${commitment}:86`)?.state).toBe('releasing') expect(releaseCalls).toBe(1) - expect(recoveryCalls).toBe(1) + expect(recoveryCalls).toBe(0) const retained = await taskStore.get('task-nonce-release-recovery') - const marker = retained?.metadata?.gatewayPaymentRelease as { - lease: { id: string; expiresAt: number } - operationId: string - } + const marker = retained?.metadata?.gatewayPaymentRecovery as { id: string } expect(retained?.status.state).toBe('failed') - expect(marker.operationId).toBe(`x402:${commitment}:86`) - await taskStore.put({ - ...retained!, - metadata: { - ...retained!.metadata, - gatewayPaymentRelease: { - ...marker, - lease: { ...marker.lease, expiresAt: Date.now() - 1 }, - }, - }, - }) + expect(marker.id).toBe(`x402:${commitment}:86`) + await expect(recoverPayment(marker.id, config, { force: true })).rejects.toThrow( + 'release recovery acknowledgement also lost', + ) + expect(recoveryCalls).toBe(1) + expect((await recoverPayment(marker.id, config, { force: true }))?.state).toBe('reconciled') const recovered = await app.request('/v1/agents/a2a-races', { method: 'POST', @@ -968,7 +1085,7 @@ describe('A2A payment ownership races', () => { }), }) expect(recovered.status).toBe(200) - expect((await taskStore.get('task-nonce-release-recovery'))?.metadata?.gatewayPaymentRelease) + expect((await taskStore.get('task-nonce-release-recovery'))?.metadata?.gatewayPaymentRecovery) .toBeUndefined() expect(operations.get(`x402:${commitment}:86`)?.state).toBe('released') expect(recoveryCalls).toBe(2) @@ -1011,7 +1128,7 @@ describe('A2A payment ownership races', () => { a2a: { taskStore, authorizeTaskAccess: async () => true }, } const app = new Hono() - app.route('/v1/agents', createAgentGateway(config)) + app.route('/v1/agents', createTestGateway(config)) const send = await app.request('/v1/agents/a2a-races', { method: 'POST', @@ -1067,4 +1184,224 @@ describe('A2A payment ownership races', () => { expect((await taskStore.get('task-usage-recovery'))?.metadata?.gatewayFinalizing).toBeUndefined() expect(operations.get(`x402:${commitment}:85`)?.state).toBe('settled') }) + + it('cancels during generic MPP confirmation without executing or stranding the charge', async () => { + const taskStore = new InMemoryTaskStore() + const lifecycle = new A2AChargeLifecycle() + let confirmationStarted!: () => void + const confirmationReady = new Promise((resolve) => { confirmationStarted = resolve }) + lifecycle.confirmationStarted = confirmationStarted + let releaseConfirmation!: () => void + lifecycle.confirmationGate = new Promise((resolve) => { releaseConfirmation = resolve }) + let runs = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const send = app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: stripePaymentHeader('spt_cancel_confirm'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-stripe-confirm-cancel') }, + }), + }) + await confirmationReady + expect((await taskStore.get('task-stripe-confirm-cancel'))?.metadata?.gatewayPaymentRecovery) + .toBeDefined() + await taskStore.delete('task-stripe-confirm-cancel') + expect(await taskStore.get('task-stripe-confirm-cancel')).toBeDefined() + const cancel = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'task-stripe-confirm-cancel' }, + }), + }) + expect(cancel.status).toBe(200) + releaseConfirmation() + const response = await send + const body = await response.json() as { result?: { status?: { state?: string } } } + + expect(body.result?.status?.state).toBe('canceled') + expect(runs).toBe(0) + expect(lifecycle.confirmations).toBe(1) + expect(lifecycle.releases).toBe(1) + expect([...lifecycle.operations.values()][0]?.state).toBe('released') + expect((await taskStore.get('task-stripe-confirm-cancel'))?.metadata?.gatewayPaymentRecovery) + .toBeUndefined() + }) + + it('returns the confirmed generic MPP receipt on A2A send and stream responses', async () => { + const taskStore = new InMemoryTaskStore() + const lifecycle = new A2AChargeLifecycle() + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const request = (method: 'message/send' | 'message/stream', taskId: string, token: string) => + app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: stripePaymentHeader(token), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: taskId, + method, + params: { message: message('run', taskId) }, + }), + }) + + const sent = await request('message/send', 'task-stripe-send-receipt', 'spt_send_receipt') + await sent.text() + const streamed = await request('message/stream', 'task-stripe-stream-receipt', 'spt_stream_receipt') + await streamed.text() + + expect(sent.status).toBe(200) + expect(sent.headers.get('Payment-Receipt')).toContain('stripe-receipt=') + expect(sent.headers.get('X-Payment-Operation-Id')).toMatch(/^mpp:stripe:/) + expect(streamed.status).toBe(200) + expect(streamed.headers.get('Payment-Receipt')).toContain('stripe-receipt=') + expect(streamed.headers.get('X-Payment-Operation-Id')).toMatch(/^mpp:stripe:/) + expect(lifecycle.confirmations).toBe(2) + }) + + it('recovers generic MPP usage acknowledgement loss from the durable task receipt', async () => { + const taskStore = new InMemoryTaskStore() + const recoveryStore = new MemoryPaymentRecoveryStore() + const lifecycle = new A2AChargeLifecycle() + const requestIds = new Set() + let recordCalls = 0 + let firstAcknowledgement = true + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid' } } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async (event) => { + recordCalls += 1 + requestIds.add(event.requestId) + if (firstAcknowledgement) { + firstAcknowledgement = false + throw new Error('usage acknowledgement lost after insert') + } + }, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const sent = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: stripePaymentHeader('spt_usage_ack'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { message: message('run', 'task-stripe-usage-recovery') }, + }), + }) + const failed = await sent.json() as { error?: { code?: number } } + expect(failed.error?.code).toBe(-32603) + const retained = await taskStore.get('task-stripe-usage-recovery') + const finalization = retained?.metadata?.gatewayFinalizing as { + lease: { id: string; expiresAt: number } + } + expect(retained?.metadata?.gatewayPaymentRecovery).toBeDefined() + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayFinalizing: { + ...finalization, + lease: { ...finalization.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + + const recovered = await app.request('/v1/agents/a2a-races', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/get', + params: { id: 'task-stripe-usage-recovery' }, + }), + }) + const body = await recovered.json() as { result?: { status?: { state?: string } } } + + expect(body.result?.status?.state).toBe('completed') + expect(recordCalls).toBe(2) + expect(requestIds.size).toBe(1) + expect(lifecycle.confirmations).toBe(1) + expect(lifecycle.releases).toBe(0) + expect((await taskStore.get('task-stripe-usage-recovery'))?.metadata).toBeUndefined() + }) }) diff --git a/tests/a2a.test.ts b/tests/a2a.test.ts index 0589ae9..f978df3 100644 --- a/tests/a2a.test.ts +++ b/tests/a2a.test.ts @@ -194,7 +194,10 @@ describe('A2A — AgentCard discovery', () => { mpp: { realm: 'agents.tangle.tools', method: 'blueprintevm', - verifySigner: async () => 'mpp-signer', + authenticateCredential: async () => ({ + consumerId: 'mpp-signer', + paymentIdentity: 'mpp-signer-payment', + }), }, }) const res = await app.request('/v1/agents/test-agent/.well-known/agent.json') diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index a6a3568..c875663 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -19,10 +19,39 @@ import type { import { MemoryNonceStore, type NonceStore } from '../src/nonce-store' import { MemoryRateLimitStore } from '../src/rate-limit' import { MemoryPaymentOperations } from '../src/payment-operations' +import type { MppChargeLifecycle } from '../src/mpp-payment' const operatorAddress = '0x1111111111111111111111111111111111111111' const fundedRequestAmount = '1000000' +function mppChargeLifecycle(onConfirm?: (credential: string) => void): MppChargeLifecycle { + const operations = new Map>>() + return { + protocolVersion: 1, + async confirmPayment(request) { + onConfirm?.(request.credential) + const operation = { + protocolVersion: 1 as const, + operationId: request.operationId, + acquiredByRequestId: request.requestId, + method: request.method, + receipt: `receipt=${request.operationId}`, + state: 'confirmed' as const, + } + operations.set(operation.operationId, operation) + return operation + }, + async releasePayment(operation) { + const released = { ...operation, state: 'released' as const } + operations.set(operation.operationId, released) + return released + }, + async recoverPayment(operationId) { + return operations.get(operationId) ?? { operationId, state: 'not-found' as const } + }, + } +} + /** Sandbox that emits a fixed reply, captures the prompt + opts for assertion */ class StubSandbox implements SandboxBox { receivedPrompt: string | null = null @@ -574,7 +603,10 @@ describe('POST /:slug/chat/completions — auth paths', () => { mpp: { realm: 'agents.tangle.tools', method: 'blueprintevm', - verifySigner: async () => 'mpp:consumer', + authenticateCredential: async (payload) => ({ + consumerId: 'mpp:consumer', + paymentIdentity: `${String(payload.commitment)}:${String(payload.nonce)}`, + }), }, x402: { operatorAddress, @@ -615,7 +647,12 @@ describe('POST /:slug/chat/completions — auth paths', () => { const { app, settlements, usage } = buildHarness({ mpp: { realm: 'agents.tangle.tools', - verifySigner: async () => 'mpp:stripe-customer', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'mpp:stripe-customer', + paymentIdentity: String(payload.receiptId), + }), + charge: mppChargeLifecycle(), }, x402: { operatorAddress, @@ -639,8 +676,8 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(streamed.combinedText).toBe('Hello, world!') expect(streamed.done).toBe(true) expect(operations.get('x402:stripe-customer:902')).toBeUndefined() - expect(settlements).toHaveLength(1) - expect(settlements[0]?.method).toBe('mpp') + expect(settlements).toHaveLength(0) + expect(response.headers.get('Payment-Receipt')).toContain('receipt=') expect(usage).toHaveLength(1) }) @@ -671,7 +708,11 @@ describe('POST /:slug/chat/completions — auth paths', () => { mpp: { realm: 'agents.tangle.tools', method: 'stripe', - verifySigner: async () => 'mpp:consumer', + authenticateCredential: async (payload) => ({ + consumerId: 'mpp:consumer', + paymentIdentity: String(payload.receiptId), + }), + charge: mppChargeLifecycle(), }, }) const request = () => app.request('/v1/agents/test-agent/chat/completions', { @@ -686,12 +727,14 @@ describe('POST /:slug/chat/completions — auth paths', () => { const responses = await Promise.all([request(), request()]) await Promise.all(responses.map((response) => response.text())) - expect(responses.map((response) => response.status).sort()).toEqual([200, 402]) + const statuses = responses.map((response) => response.status) + expect(statuses.filter((status) => status === 200)).toHaveLength(1) + expect(statuses.filter((status) => status === 401 || status === 402)).toHaveLength(1) expect(executions).toBe(1) - expect(settlements).toHaveLength(1) + expect(settlements).toHaveLength(0) }) - it('requires complete receipts only for requests with durable payment ownership', async () => { + it('requires complete receipts for every durable x402 or generic MPP payment', async () => { const operations = new MemoryPaymentOperations() const { app } = buildHarness({ getSandbox: async () => ({ @@ -703,7 +746,11 @@ describe('POST /:slug/chat/completions — auth paths', () => { mpp: { realm: 'agents.tangle.tools', method: 'stripe', - verifySigner: async () => 'mpp:consumer', + authenticateCredential: async (payload) => ({ + consumerId: 'mpp:consumer', + paymentIdentity: String(payload.receiptId ?? 'empty-receipt'), + }), + charge: mppChargeLifecycle(), }, x402: { operatorAddress, @@ -740,7 +787,7 @@ describe('POST /:slug/chat/completions — auth paths', () => { expect(apiKeyStream.combinedText).toBe('legacy') expect(mppStream.combinedText).toBe('legacy') expect(apiKeyStream.done).toBe(true) - expect(mppStream.done).toBe(true) + expect(mppStream.done).toBe(false) }) it('threads a unique requestId per concurrent request — regression: two same-consumer requests get distinct ids', async () => { @@ -1103,6 +1150,21 @@ describe('createAgentGateway — production-config guard', () => { })).toThrow(/paymentProtocolVersion must be explicit/) }) + it('requires a durable recovery outbox for production payment protocol version 2', () => { + expect(() => createAgentGateway({ + resolveAgent: async () => null, + getSandbox: async () => ({ async *streamPrompt() { /* unused */ } }), + recordUsage: async () => { /* unused */ }, + x402: { + operatorAddress, + chainId: 3799, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations(), + }, + })).toThrow(/durable payment recovery is required in production/) + }) + it('keeps older custom A2A task stores source-compatible', () => { expect(() => createAgentGateway({ resolveAgent: async () => null, diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index 0afcfb1..b8b0e41 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -122,6 +122,7 @@ describe('version 2 payment operations', () => { const settled = await operations.settlePayment(owner, { amount: 200n, totalCostUsd: 0.2, + basis: 'usage-receipt', usage: { inputTokens: 1, outputTokens: 1, @@ -138,6 +139,7 @@ describe('version 2 payment operations', () => { await expect(operations.settlePayment(owner, { amount: 1001n, totalCostUsd: 1, + basis: 'usage-receipt', usage: { ...settledUsage(), budgetEnforced: true }, })).rejects.toThrow() }) @@ -174,6 +176,7 @@ describe('version 2 payment operations', () => { const settled = await operations.settlePayment(retained, { amount: 200n, totalCostUsd: 0.2, + basis: 'usage-receipt', usage: settledUsage(), }) expect(settled.state).toBe('settled') @@ -242,6 +245,7 @@ describe('version 2 payment operations', () => { const input = { amount: 200n, totalCostUsd: 0.2, + basis: 'usage-receipt' as const, usage: settledUsage(), } const settled = await Promise.all([ @@ -289,7 +293,7 @@ describe('version 2 payment operations', () => { }, }) const owner = await operations.claimPayment(payload('1000', '17'), context()) - const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage() } + const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage(), basis: 'usage-receipt' as const } await expect(operations.settlePayment(owner, input)).rejects.toThrow('worker crashed') const first = operations.reclaimPayment(owner.operationId) while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) @@ -313,7 +317,7 @@ describe('version 2 payment operations', () => { }, }) const owner = await operations.claimPayment(payload('1000', '18'), context()) - const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage() } + const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage(), basis: 'usage-receipt' as const } await expect(operations.settlePayment(owner, input)).rejects.toThrow('worker crashed') const first = operations.settlePayment(owner, input) while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) @@ -337,7 +341,7 @@ describe('version 2 payment operations', () => { }, }) const owner = await operations.claimPayment(payload('1000', '19'), context()) - const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage() } + const input = { amount: 200n, totalCostUsd: 0.2, usage: settledUsage(), basis: 'usage-receipt' as const } await expect(operations.settlePayment(owner, input)).rejects.toThrow('worker crashed') const reclaim = operations.reclaimPayment(owner.operationId) while (entered === 0) await new Promise((resolve) => setTimeout(resolve, 0)) @@ -372,12 +376,6 @@ describe('version 2 payment operations', () => { expect(recovered[1].state).toBe('released') }) - it('rejects a release callback without a recovery proof', () => { - expect(() => new MemoryPaymentOperations({ - onRelease: async () => { throw new Error('acknowledgement lost') }, - })).toThrow('onReclaim is required') - }) - it('reclaims a claim that crashed after durable ownership but before completion', async () => { let now = 100 let recoveries = 0 diff --git a/tests/payment-recovery-sql.test.ts b/tests/payment-recovery-sql.test.ts new file mode 100644 index 0000000..0d9a621 --- /dev/null +++ b/tests/payment-recovery-sql.test.ts @@ -0,0 +1,196 @@ +import { describe, expect, it } from 'vitest' + +import type { SqlAdapter } from '../src/a2a/task-store-sql' +import { + PaymentRecoveryFenceError, + updateOwnedPaymentRecovery, + type PaymentRecoveryRecord, +} from '../src/payment-recovery' +import { SqlPaymentRecoveryStore } from '../src/payment-recovery-sql' + +interface Row { + id: string + state: string + next_attempt_at: number + revision: number + payload: string + updated_at: number +} + +function adapter(): SqlAdapter & { rows: Map } { + const rows = new Map() + return { + rows, + async exec(sql, params = []) { + const statement = sql.trim() + if (statement.startsWith('CREATE TABLE') || statement.startsWith('CREATE INDEX')) { + return { rowsAffected: 0 } + } + if (statement.startsWith('INSERT INTO')) { + const [id, state, nextAttemptAt, revision, payload, updatedAt] = params as [ + string, + string, + number, + number, + string, + number, + ] + if (rows.has(id)) throw new Error('duplicate primary key') + rows.set(id, { + id, + state, + next_attempt_at: nextAttemptAt, + revision, + payload, + updated_at: updatedAt, + }) + return { rowsAffected: 1 } + } + if (statement.startsWith('UPDATE')) { + const [state, nextAttemptAt, revision, payload, updatedAt, id, expectedRevision] = params as [ + string, + number, + number, + string, + number, + string, + number, + ] + const row = rows.get(id) + if (!row || row.revision !== expectedRevision) return { rowsAffected: 0 } + rows.set(id, { + id, + state, + next_attempt_at: nextAttemptAt, + revision, + payload, + updated_at: updatedAt, + }) + return { rowsAffected: 1 } + } + throw new Error(`unrecognized SQL: ${statement}`) + }, + async query(sql: string, params: readonly unknown[] = []): Promise { + const statement = sql.trim() + if (statement.includes('WHERE id =')) { + const row = rows.get(params[0] as string) + return (row ? [{ payload: row.payload }] : []) as TRow[] + } + if (statement.includes('next_attempt_at <=')) { + const [, now, limit] = params as [string, number, number] + return [...rows.values()] + .filter((row) => row.state !== 'reconciled' && row.next_attempt_at <= now) + .sort((left, right) => left.next_attempt_at - right.next_attempt_at) + .slice(0, limit) + .map((row) => ({ payload: row.payload })) as TRow[] + } + throw new Error(`unrecognized SQL: ${statement}`) + }, + } +} + +function record(id = 'payment-1'): PaymentRecoveryRecord { + return { + version: 1, + id, + revision: 0, + state: 'claimed', + payment: { kind: 'x402', operationId: id }, + attribution: { + requestId: 'request-1', + agentId: 'agent-1', + agentSlug: 'agent', + consumerId: 'consumer-1', + paymentMethod: 'x402', + startMs: 1, + pricePerTokenUsd: 0.000001, + platformFeePercent: 0.2, + requiredAmount: '100', + currencyDecimals: 6, + maxOutputTokens: 1, + executionBudget: { + maxInputTokens: 1, + maxOutputTokens: 1, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 0.0001, + }, + }, + workStarted: false, + usageRecorded: false, + attempts: 0, + nextAttemptAt: 10, + createdAt: 1, + updatedAt: 1, + } +} + +describe('SqlPaymentRecoveryStore', () => { + it('persists rows, rejects duplicate payment identities, and uses revision CAS', async () => { + const db = adapter() + const store = new SqlPaymentRecoveryStore(db) + await store.migrate() + const initial = record() + expect(await store.createIfAbsent(initial)).toBe(true) + expect(await store.createIfAbsent(initial)).toBe(false) + expect(await store.get(initial.id)).toEqual(initial) + + const next = { ...initial, revision: 1, state: 'executing' as const, updatedAt: 2 } + expect(await store.compareAndSet(initial, next)).toBe(true) + expect(await store.compareAndSet(initial, { ...next, revision: 2 })).toBe(false) + expect((await store.get(initial.id))?.state).toBe('executing') + }) + + it('scans only due unresolved rows and keeps reconciled tombstones', async () => { + const store = new SqlPaymentRecoveryStore(adapter()) + const due = record('due') + const later = { ...record('later'), nextAttemptAt: 100 } + const reconciled = { + ...record('done'), + state: 'reconciled' as const, + reconciledAt: 5, + nextAttemptAt: Number.MAX_SAFE_INTEGER, + } + await store.createIfAbsent(due) + await store.createIfAbsent(later) + await store.createIfAbsent(reconciled) + + expect((await store.listDue(10, 10)).map((value) => value.id)).toEqual(['due']) + expect(await store.get('done')).toEqual(reconciled) + }) + + it('rejects an unsafe interpolated table name', () => { + expect(() => new SqlPaymentRecoveryStore(adapter(), { table: 'payments; DROP TABLE users' })) + .toThrow('table name is invalid') + }) + + it('rejects a stale owner after another lease takes the durable row', async () => { + const store = new SqlPaymentRecoveryStore(adapter()) + const first = { + ...record('fenced'), + lease: { id: 'request-fence', expiresAt: 10 }, + } + await store.createIfAbsent(first) + const replacement: PaymentRecoveryRecord = { + ...first, + revision: 1, + lease: { id: 'worker-fence', expiresAt: 20 }, + updatedAt: 11, + } + expect(await store.compareAndSet(first, replacement)).toBe(true) + + await expect(updateOwnedPaymentRecovery( + store, + first.id, + 'request-fence', + (current) => ({ ...current, state: 'executing' }), + )).rejects.toBeInstanceOf(PaymentRecoveryFenceError) + expect((await updateOwnedPaymentRecovery( + store, + first.id, + 'worker-fence', + (current) => ({ ...current, state: 'executing' }), + )).state).toBe('executing') + }) +}) diff --git a/tests/payment-recovery.test.ts b/tests/payment-recovery.test.ts new file mode 100644 index 0000000..43412f5 --- /dev/null +++ b/tests/payment-recovery.test.ts @@ -0,0 +1,789 @@ +import { Hono } from 'hono' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { InMemoryTaskStore } from '../src/a2a/task-store' +import { createAgentGateway } from '../src/middleware' +import type { + MppChargeLifecycle, + MppChargeOperation, +} from '../src/mpp-payment' +import { MemoryNonceStore } from '../src/nonce-store' +import { MemoryPaymentOperations, type PaymentSettlementInput } from '../src/payment-operations' +import { MemoryPaymentRecoveryStore } from '../src/payment-recovery' +import { recoverPayment, recoverPayments } from '../src/payment-recovery-worker' +import type { + AgentMeta, + GatewayConfig, + GatewayUsageEvent, + SandboxBox, + SandboxUsageReceipt, +} from '../src/types' + +const operatorAddress = '0x1111111111111111111111111111111111111111' +const commitment = `0x${'ab'.repeat(32)}` + +const agent: AgentMeta = { + id: 'agent-recovery', + ownerId: 'owner-recovery', + slug: 'recovery', + systemPrompt: '', + pricePerTokenUsd: 0.000001, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +afterEach(() => { + vi.useRealTimers() +}) + +function spendAuth(nonce: string, amount = '1000000000'): string { + return JSON.stringify({ + commitment, + signature: '0xsig', + operator: operatorAddress, + amount, + nonce, + expiry: String(Math.floor(Date.now() / 1000) + 600), + }) +} + +function mppCredential(value: Record): string { + return Buffer.from(JSON.stringify(value)).toString('base64url') +} + +function rawMppCredential(value: string): string { + return Buffer.from(value).toString('base64url') +} + +function receipt(): SandboxUsageReceipt { + return { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.000002, + budgetEnforced: true, + } +} + +function mount(config: GatewayConfig): Hono { + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + return app +} + +function requestBody(message = 'run') { + return JSON.stringify({ max_tokens: 4, messages: [{ role: 'user', content: message }] }) +} + +class HostileChargeLifecycle implements MppChargeLifecycle { + readonly protocolVersion = 1 as const + readonly operations = new Map() + confirmations = 0 + refunds = 0 + recovered = 0 + confirmGate?: Promise + confirmStarted?: () => void + loseConfirmAcknowledgement = false + loseReleaseAcknowledgement = false + receivedCredential?: string + receiptValue?: string + + async confirmPayment(request: Parameters[0]) { + this.confirmations += 1 + this.receivedCredential = request.credential + this.confirmStarted?.() + await this.confirmGate + const operation: MppChargeOperation = { + protocolVersion: 1, + operationId: request.operationId, + acquiredByRequestId: request.requestId, + method: request.method, + receipt: this.receiptValue ?? `stripe-receipt=${request.operationId}`, + state: 'confirmed', + } + this.operations.set(operation.operationId, operation) + if (this.loseConfirmAcknowledgement) throw new Error('Stripe confirmation acknowledgement lost') + return operation + } + + async releasePayment(operation: MppChargeOperation) { + const existing = this.operations.get(operation.operationId) + if (existing?.state === 'released') return existing + this.refunds += 1 + const released: MppChargeOperation = { ...operation, state: 'released' } + this.operations.set(operation.operationId, released) + if (this.loseReleaseAcknowledgement) { + this.loseReleaseAcknowledgement = false + throw new Error('Stripe refund acknowledgement lost') + } + return released + } + + async recoverPayment(operationId: string) { + this.recovered += 1 + return this.operations.get(operationId) ?? { operationId, state: 'not-found' as const } + } +} + +describe('generic MPP charge lifecycle', () => { + it('confirms only after all denials and before any response, body, or sandbox work', async () => { + const deniedLifecycle = new HostileChargeLifecycle() + let deniedRuns = 0 + const denied = mount({ + resolveAgent: async () => agent, + getSandbox: async () => ({ async *streamPrompt() { deniedRuns += 1 } }), + recordUsage: async () => undefined, + authorizeConsumer: async () => ({ allow: false, reason: 'not a member', code: 'denied' }), + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: deniedLifecycle, + }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + }) + const credential = mppCredential({ sharedPaymentToken: 'spt_secret' }) + const deniedResponse = await denied.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${credential}`, + }, + body: requestBody(), + }) + expect(deniedResponse.status).toBe(403) + expect(deniedLifecycle.confirmations).toBe(0) + expect(deniedRuns).toBe(0) + + const lifecycle = new HostileChargeLifecycle() + let confirmStarted!: () => void + const confirmationStarted = new Promise((resolve) => { confirmStarted = resolve }) + lifecycle.confirmStarted = confirmStarted + let allowConfirmation!: () => void + lifecycle.confirmGate = new Promise((resolve) => { allowConfirmation = resolve }) + let runs = 0 + let legacySettlements = 0 + const app = mount({ + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'paid' } } + yield { type: 'sandbox.usage', data: { usage: receipt() } } + }, + }), + recordUsage: async () => undefined, + settlePayment: async () => { legacySettlements += 1 }, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + }) + let responseResolved = false + const responsePromise = app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${credential}`, + }, + body: requestBody(), + }).then((response) => { + responseResolved = true + return response + }) + await confirmationStarted + await Promise.resolve() + expect(responseResolved).toBe(false) + expect(runs).toBe(0) + allowConfirmation() + const response = await responsePromise + expect(response.status).toBe(200) + expect(response.headers.get('Payment-Receipt')).toContain('stripe-receipt=') + expect(lifecycle.receivedCredential).toContain('spt_secret') + expect(await response.text()).toContain('paid') + expect(runs).toBe(1) + expect(lifecycle.confirmations).toBe(1) + expect(legacySettlements).toBe(0) + }) + + it('recovers lost confirmation and refund acknowledgements without work or a second charge', async () => { + const lifecycle = new HostileChargeLifecycle() + lifecycle.loseConfirmAcknowledgement = true + lifecycle.loseReleaseAcknowledgement = true + const recoveryStore = new MemoryPaymentRecoveryStore() + let runs = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ async *streamPrompt() { runs += 1 } }), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore, retryDelayMs: 1 }, + } + const app = mount(config) + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${mppCredential({ sharedPaymentToken: 'spt_ack_loss' })}`, + }, + body: requestBody(), + }) + expect(response.status).toBe(402) + expect(runs).toBe(0) + expect(lifecycle.confirmations).toBe(1) + + const [record] = await recoveryStore.listDue(Number.MAX_SAFE_INTEGER, 10) + expect(record?.state).toBe('claiming') + await expect(recoverPayment(record.id, config, { force: true })).rejects.toThrow( + 'Stripe refund acknowledgement lost', + ) + expect((await recoveryStore.get(record.id))?.state).toBe('releasing') + const recovered = await recoverPayment(record.id, config, { force: true }) + expect(recovered?.state).toBe('reconciled') + expect(lifecycle.confirmations).toBe(1) + expect(lifecycle.refunds).toBe(1) + expect(runs).toBe(0) + }) + + it('fences a live request when recovery releases its expired claim lease', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-14T00:00:00.000Z')) + const lifecycle = new HostileChargeLifecycle() + const recoveryStore = new MemoryPaymentRecoveryStore() + let observerEntered!: () => void + const observerReady = new Promise((resolve) => { observerEntered = resolve }) + let releaseObserver!: () => void + const observerGate = new Promise((resolve) => { releaseObserver = resolve }) + let runs = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield { type: 'sandbox.usage', data: { usage: receipt() } } + }, + }), + recordUsage: async () => undefined, + observer: { + async onPaymentVerified() { + observerEntered() + await observerGate + }, + }, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore, staleRequestMs: 5 }, + } + const app = mount(config) + const pending = app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${mppCredential({ sharedPaymentToken: 'spt_lease' })}`, + }, + body: requestBody(), + }) + await observerReady + vi.advanceTimersByTime(6) + + expect(await recoverPayments(config, { now: Date.now() })).toMatchObject({ + scanned: 1, + reconciled: 1, + }) + releaseObserver() + const response = await pending + + expect(response.status).toBe(402) + expect(runs).toBe(0) + expect(lifecycle.confirmations).toBe(1) + expect(lifecycle.refunds).toBe(1) + expect([...lifecycle.operations.values()][0]?.state).toBe('released') + }) + + it('uses the adapter payment identity for replay and never persists the Stripe credential', async () => { + const lifecycle = new HostileChargeLifecycle() + const recoveryStore = new MemoryPaymentRecoveryStore() + const claimedKeys = new Set() + let runs = 0 + const secret = 'spt_same_secret' + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + runs += 1 + yield { type: 'sandbox.usage', data: { usage: receipt() } } + }, + }), + recordUsage: async () => undefined, + nonceStore: { + hasSeen: async (key) => claimedKeys.has(key), + claim: async (key) => { + if (claimedKeys.has(key)) return false + claimedKeys.add(key) + return true + }, + }, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore }, + } + const app = mount(config) + const request = (json: string) => app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${rawMppCredential(json)}`, + }, + body: requestBody(), + }) + + const first = await request(`{"sharedPaymentToken":"${secret}"}`) + await first.text() + const replay = await request(`{ "sharedPaymentToken": "${secret}" }`) + await replay.text() + + expect(first.status).toBe(200) + expect(replay.status).toBe(401) + expect(lifecycle.confirmations).toBe(1) + expect(runs).toBe(1) + const [key] = [...claimedKeys] + const record = await recoveryStore.get(key) + expect(key).toMatch(/^mpp:stripe:[a-f0-9]{64}$/) + expect(key).not.toContain(secret) + expect(JSON.stringify(record)).not.toContain(secret) + }) + + it('retains a synchronous sandbox start failure and settles it once at the bounded ceiling', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-14T00:00:00.000Z')) + const lifecycle = new HostileChargeLifecycle() + const recoveryStore = new MemoryPaymentRecoveryStore() + const usage: GatewayUsageEvent[] = [] + let starts = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + streamPrompt() { + starts += 1 + throw new Error('provider started then lost the stream') + }, + }), + recordUsage: async (event) => { usage.push(event) }, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore, receiptTimeoutMs: 10 }, + } + const app = mount(config) + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${mppCredential({ sharedPaymentToken: 'spt_sync_start' })}`, + }, + body: requestBody(), + }) + expect(await response.text()).toContain('provider started then lost the stream') + const [pending] = await recoveryStore.listDue(Number.MAX_SAFE_INTEGER, 10) + + expect(starts).toBe(1) + expect(lifecycle.refunds).toBe(0) + expect(pending.state).toBe('retained') + vi.advanceTimersByTime(11) + const recovered = await recoverPayment(pending.id, config) + expect(recovered?.state).toBe('reconciled') + expect(lifecycle.confirmations).toBe(1) + expect(lifecycle.refunds).toBe(0) + expect(usage).toHaveLength(1) + expect(usage[0]?.settlementBasis).toBe('quoted-ceiling') + }) + + it('refunds a confirmed charge when its receipt is unsafe for an HTTP header', async () => { + const lifecycle = new HostileChargeLifecycle() + lifecycle.receiptValue = 'stripe-receipt=\u2603' + const recoveryStore = new MemoryPaymentRecoveryStore() + let runs = 0 + const app = mount({ + resolveAgent: async () => agent, + getSandbox: async () => ({ async *streamPrompt() { runs += 1 } }), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async (payload) => ({ + consumerId: 'stripe:customer', + paymentIdentity: String(payload.sharedPaymentToken), + }), + charge: lifecycle, + }, + paymentRecovery: { store: recoveryStore }, + }) + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Payment stripe ${mppCredential({ sharedPaymentToken: 'spt_bad_receipt' })}`, + }, + body: requestBody(), + }) + const [operation] = [...lifecycle.operations.values()] + + expect(response.status).toBe(402) + expect(runs).toBe(0) + expect(lifecycle.confirmations).toBe(1) + expect(lifecycle.refunds).toBe(1) + expect(operation?.state).toBe('released') + expect((await recoveryStore.get(operation!.operationId))?.state).toBe('reconciled') + }) +}) + +describe('durable OpenAI recovery', () => { + it('persists operation and attribution before output, then heals settlement acknowledgement loss', async () => { + const recoveryStore = new MemoryPaymentRecoveryStore() + let settleCalls = 0 + let recoveryCalls = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { + settleCalls += 1 + throw new Error('settlement acknowledgement lost') + }, + onReclaim: async () => { recoveryCalls += 1 }, + }) + const usage = new Map() + let recordAtSandboxStart: Awaited> + const operationId = `x402:${commitment}:501` + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + recordAtSandboxStart = await recoveryStore.get(operationId) + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'valuable' } } + yield { type: 'sandbox.usage', data: { usage: receipt() } } + }, + }), + recordUsage: async (event) => { usage.set(event.requestId, event) }, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore, retryDelayMs: 1 }, + maxOutputTokens: 4, + defaultOutputTokens: 4, + } + const app = mount(config) + const payment = spendAuth('501') + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': payment }, + body: requestBody(), + }) + const wire = await response.text() + expect(response.status).toBe(200) + expect(wire).toContain('valuable') + expect(wire).not.toContain('[DONE]') + expect(recordAtSandboxStart?.state).toBe('executing') + expect(recordAtSandboxStart?.workStarted).toBe(true) + expect(recordAtSandboxStart?.attribution.agentId).toBe(agent.id) + expect(recordAtSandboxStart?.attribution.consumerId).toBe(commitment) + expect(recordAtSandboxStart?.usage).toBeUndefined() + + const pending = await recoveryStore.get(operationId) + expect(pending?.state).toBe('settling') + expect(pending?.usage).toEqual(receipt()) + expect(pending?.usageRecorded).toBe(false) + expect(operations.get(operationId)?.state).toBe('settling') + expect(settleCalls).toBe(1) + expect(usage.size).toBe(0) + + const recovered = await recoverPayment(operationId, config, { force: true }) + expect(recovered?.state).toBe('reconciled') + expect(recovered?.usageRecorded).toBe(true) + expect(recoveryCalls).toBe(1) + expect(usage.size).toBe(1) + expect(operations.get(operationId)?.state).toBe('settled') + await recoverPayment(operationId, config, { force: true }) + expect(recoveryCalls).toBe(1) + expect(usage.size).toBe(1) + expect(await recoveryStore.get(operationId)).toBeDefined() + + const replay = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': payment }, + body: requestBody('replay'), + }) + expect(replay.status).toBe(402) + }) + + it('settles a canceled no-receipt run once at the quoted ceiling after the bounded timeout', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-14T00:00:00.000Z')) + const recoveryStore = new MemoryPaymentRecoveryStore() + const settlements: PaymentSettlementInput[] = [] + let settlementStarted!: () => void + const settlementReady = new Promise((resolve) => { settlementStarted = resolve }) + let finishSettlement!: () => void + const settlementReleased = new Promise((resolve) => { finishSettlement = resolve }) + const operations = new MemoryPaymentOperations({ + onSettle: async (_operation, input) => { + settlements.push(input) + settlementStarted() + await settlementReleased + }, + onReclaim: async () => undefined, + }) + const usage = new Map() + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt(_message, options) { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'partial' } } + await new Promise((resolve) => { + if (options?.signal?.aborted) resolve() + else options?.signal?.addEventListener('abort', () => resolve(), { once: true }) + }) + }, + }), + recordUsage: async (event) => { usage.set(event.requestId, event) }, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore, receiptTimeoutMs: 10, retryDelayMs: 1 }, + maxOutputTokens: 4, + defaultOutputTokens: 4, + } + const app = mount(config) + const payment = spendAuth('502') + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': payment }, + body: requestBody(), + }) + const reader = response.body!.getReader() + expect(new TextDecoder().decode((await reader.read()).value)).toContain('partial') + await reader.cancel() + for (let attempt = 0; attempt < 20; attempt += 1) await Promise.resolve() + + const operationId = `x402:${commitment}:502` + const retained = await recoveryStore.get(operationId) + expect(retained?.state).toBe('retained') + expect(operations.get(operationId)?.state).toBe('retained') + expect(settlements).toHaveLength(0) + expect((await recoverPayments(config, { now: Date.now() })).scanned).toBe(0) + + vi.setSystemTime(Date.now() + 11) + const firstRecovery = recoverPayment(operationId, config) + await settlementReady + const secondRecovery = recoverPayment(operationId, config) + finishSettlement() + await Promise.all([firstRecovery, secondRecovery]) + expect((await recoveryStore.get(operationId))?.state).toBe('reconciled') + expect(settlements).toHaveLength(1) + expect(settlements[0]?.basis).toBe('quoted-ceiling') + expect(settlements[0]?.amount).toBe(BigInt(retained!.attribution.requiredAmount)) + expect(settlements[0]?.amount).toBeLessThan(1000000000n) + expect(operations.get(operationId)?.state).toBe('settled') + expect(usage.size).toBe(1) + expect([...usage.values()][0]?.settlementBasis).toBe('quoted-ceiling') + expect((await recoverPayments(config, { now: Date.now() })).scanned).toBe(0) + expect(settlements).toHaveLength(1) + + const replay = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': payment }, + body: requestBody('replay'), + }) + expect(replay.status).toBe(402) + }) + + it('never broadcasts one recovered receipt across the batch worker', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-14T00:00:00.000Z')) + const recoveryStore = new MemoryPaymentRecoveryStore() + const settlements: PaymentSettlementInput[] = [] + const operations = new MemoryPaymentOperations({ + onSettle: async (_operation, input) => { settlements.push(input) }, + onReclaim: async () => undefined, + }) + const usage: GatewayUsageEvent[] = [] + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'partial' } } + }, + }), + recordUsage: async (event) => { usage.push(event) }, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore, receiptTimeoutMs: 10 }, + maxOutputTokens: 4, + defaultOutputTokens: 4, + } + const app = mount(config) + for (const nonce of ['504', '505']) { + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': spendAuth(nonce) }, + body: requestBody(), + }) + await response.text() + expect((await recoveryStore.get(`x402:${commitment}:${nonce}`))?.state).toBe('retained') + } + + vi.advanceTimersByTime(11) + const hostileOptions = { + now: Date.now(), + usage: receipt(), + } as unknown as Parameters[1] + expect(await recoverPayments(config, hostileOptions)).toMatchObject({ + scanned: 2, + reconciled: 2, + failed: 0, + }) + + expect(settlements).toHaveLength(2) + expect(settlements.every((settlement) => settlement.basis === 'quoted-ceiling')).toBe(true) + expect(usage).toHaveLength(2) + expect(usage.every((event) => event.settlementBasis === 'quoted-ceiling')).toBe(true) + expect(usage.every((event) => event.inputTokens === 0 && event.outputTokens === 0)).toBe(true) + }) +}) + +describe('A2A recovery retention', () => { + it('keeps an expired task until its no-receipt payment reconciles, then restores normal TTL', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-14T00:00:00.000Z')) + const taskStore = new InMemoryTaskStore(20) + const recoveryStore = new MemoryPaymentRecoveryStore() + const settlements: PaymentSettlementInput[] = [] + const operations = new MemoryPaymentOperations({ + onSettle: async (_operation, input) => { settlements.push(input) }, + onReclaim: async () => undefined, + }) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'partial task' } } + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore, receiptTimeoutMs: 10 }, + a2a: { taskStore, authorizeTaskAccess: async () => true }, + maxOutputTokens: 4, + defaultOutputTokens: 4, + } + const app = mount(config) + const send = await app.request('/v1/agents/recovery', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': spendAuth('503') }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + messageId: 'message-503', + taskId: 'task-503', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + expect(send.status).toBe(200) + expect((await taskStore.get('task-503'))?.metadata?.gatewayPaymentRecovery).toBeDefined() + await taskStore.delete('task-503') + expect(await taskStore.get('task-503')).toBeDefined() + expect(operations.get(`x402:${commitment}:503`)?.state).toBe('retained') + + vi.setSystemTime(Date.now() + 21) + expect(await taskStore.get('task-503')).toBeDefined() + const recovered = await app.request('/v1/agents/recovery', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ jsonrpc: '2.0', id: 2, method: 'tasks/get', params: { id: 'task-503' } }), + }) + const recoveredBody = await recovered.json() as { result?: { metadata?: Record } } + expect(recoveredBody.result?.metadata?.gatewayPaymentRecovery).toBeUndefined() + expect(settlements).toHaveLength(1) + expect(settlements[0]?.basis).toBe('quoted-ceiling') + expect(operations.get(`x402:${commitment}:503`)?.state).toBe('settled') + + vi.setSystemTime(Date.now() + 21) + expect(await taskStore.get('task-503')).toBeUndefined() + }) +}) diff --git a/tests/protocol-guards.test.ts b/tests/protocol-guards.test.ts index f4f3469..98b25e4 100644 --- a/tests/protocol-guards.test.ts +++ b/tests/protocol-guards.test.ts @@ -384,10 +384,6 @@ describe('final payment boundary protocol guards', () => { onRelease: async () => { releases += 1 }, onReclaim: async () => undefined, }) - let operationPromise: ReturnType | undefined - let authorizations = 0 - let releaseAuthorizations!: () => void - const authorizationsReleased = new Promise((resolve) => { releaseAuthorizations = resolve }) let runStarted!: () => void const sandboxStarted = new Promise((resolve) => { runStarted = resolve }) let finishRun!: () => void @@ -412,14 +408,7 @@ describe('final payment boundary protocol guards', () => { demoMode: true, paymentProtocolVersion: 2, paymentOperations: operations, - authorizePayment: async (payload, context) => { - operationPromise ??= operations.claimPayment(payload, context) - const operation = await operationPromise - authorizations += 1 - if (authorizations === 2) releaseAuthorizations() - await authorizationsReleased - return operation - }, + authorizePayment: (payload, context) => operations.claimPayment(payload, context), }, })) const request = () => app.request('/v1/agents/guards/chat/completions', { diff --git a/tests/verify.test.ts b/tests/verify.test.ts index baa2780..0231bc5 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -186,7 +186,10 @@ describe('verifyMpp', () => { it('parses a valid Payment header and returns the signer', async () => { const header = buildCredential({ commitment: '0xAlice', operator: operatorAddress, amount: '1000', nonce: '5' }) - expect(await verifyMpp(header, mppConfig, baseConfig)).toBe('0xAlice') + expect(await verifyMpp(header, mppConfig, baseConfig)).toMatchObject({ + consumerId: '0xAlice', + replayKey: '0xalice:5', + }) }) it('requires and calls a production MPP verifier, then rejects nonce replay', async () => { @@ -200,15 +203,18 @@ describe('verifyMpp', () => { const seen: Array<{ method: string; credential: string }> = [] const config: MppConfig = { ...mppConfig, - verifySigner: async (_payload, context) => { + authenticateCredential: async (_payload, context) => { seen.push(context) - return 'mpp:alice' + return { consumerId: 'mpp:alice', paymentIdentity: 'payment:alice:6' } }, } const productionX402: X402Config = { ...baseConfig, demoMode: false } const nonceStore = new MemoryNonceStore() - expect(await verifyMpp(header, config, productionX402, nonceStore)).toBe('mpp:alice') + expect(await verifyMpp(header, config, productionX402, nonceStore)).toMatchObject({ + consumerId: 'mpp:alice', + replayKey: '0xalice:6', + }) expect(await verifyMpp(header, config, productionX402, nonceStore)).toBeNull() expect(seen).toHaveLength(1) expect(seen[0].method).toBe('blueprintevm') @@ -245,9 +251,9 @@ describe('verifyMpp', () => { }) const config: MppConfig = { ...mppConfig, - verifySigner: async () => { + authenticateCredential: async () => { calls += 1 - return 'mpp:alice' + return { consumerId: 'mpp:alice', paymentIdentity: 'underfunded' } }, } @@ -277,7 +283,7 @@ describe('verifyMpp', () => { baseConfig, undefined, 1n, - )).toBe('0xAlice') + )).toMatchObject({ consumerId: '0xAlice', replayKey: '0xalice:9' }) }) it('rejects MPP in production when no method verifier is configured', async () => { @@ -299,7 +305,9 @@ describe('verifyMpp', () => { it('falls back to the `from` field when no `commitment` present — regression: EIP-3009 wallets expose `from` only', async () => { const header = buildCredential({ from: '0xWallet', to: operatorAddress, value: '1000' }) - expect(await verifyMpp(header, mppConfig, baseConfig)).toBe('0xWallet') + expect(await verifyMpp(header, mppConfig, baseConfig)).toMatchObject({ + consumerId: '0xWallet', + }) }) it('rejects malformed Payment header shape', async () => { From 5b3ff3617a86be2483ad2b8d848cb4259cb75db7 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 23:35:37 -0600 Subject: [PATCH 20/32] fix(gateway): harden bounded payment recovery --- src/dispatch.ts | 32 ++-- src/payment-operations.ts | 48 ++++-- src/payment-recovery-worker.ts | 22 ++- tests/payment-operations.test.ts | 31 ++++ tests/payment-recovery.test.ts | 260 ++++++++++++++++++++++++++++++- 5 files changed, 366 insertions(+), 27 deletions(-) diff --git a/src/dispatch.ts b/src/dispatch.ts index 47b9760..ffef0a6 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -564,10 +564,6 @@ export async function claimPayment( if (!result) throw new Error('payment authorization was rejected') if (typeof result !== 'boolean') { operation = result - // Retain the durable owner even if a later compatibility check fails. - // The caller can then release or reclaim the operation instead of - // losing its recovery handle. - authz.paymentOperation = operation } else if (config.x402.paymentProtocolVersion === 2) { throw new Error('version 2 payment authorization did not return an operation') @@ -584,6 +580,13 @@ export async function claimPayment( if (operation && operation.protocolVersion !== 2) { throw new Error('payment operation protocol version mismatch') } + if ( + operation && + config.x402.paymentProtocolVersion === 2 && + operation.operationId !== authz.paymentRecoveryId + ) { + throw new Error('x402 payment operation identity mismatch') + } if (operation && !config.x402.paymentOperations) { throw new Error('durable payment operations are required to settle a claimed operation') } @@ -632,6 +635,9 @@ export async function claimPayment( }, hooks) const operation = await config.x402.paymentOperations.claimPayment(durablePayload, context) if (operation.protocolVersion !== 2) throw new Error('payment operation protocol version mismatch') + if (operation.operationId !== authz.paymentRecoveryId) { + throw new Error('x402 payment operation identity mismatch') + } if (operation.acquiredByRequestId !== context.requestId) { throw new Error('payment operation was already claimed') } @@ -851,8 +857,9 @@ export async function releasePayment( await relinquishPaymentRecovery(authz, config, Date.now()) return } - await markRecoveryReleasing(authz, config, reason) + let reconciled = false try { + await markRecoveryReleasing(authz, config, reason) if (ownsX402) { authz.paymentOperation = await config.x402.paymentOperations!.releasePayment( authz.paymentOperation!, @@ -872,11 +879,18 @@ export async function releasePayment( ) authz.mppChargeOperation = operation } - } catch (error) { - await relinquishPaymentRecovery(authz, config, Date.now()) - throw error + await markRecoveryReconciled(authz, config) + reconciled = true + } finally { + if (!reconciled) { + try { + await relinquishPaymentRecovery(authz, config, Date.now()) + } catch { + // Preserve the original provider or metadata error. A later worker + // retry still has the durable row when cleanup itself is unavailable. + } + } } - await markRecoveryReconciled(authz, config) } /** Mark a durable reservation active immediately before sandbox execution. */ diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 77f10ce..e718fe8 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -94,6 +94,8 @@ export interface MemoryPaymentOperationsOptions { export class MemoryPaymentOperations implements PaymentOperations { readonly protocolVersion = PAYMENT_PROTOCOL_VERSION private readonly operations = new Map() + private readonly claimFlights = new Map>() + private readonly claimTokens = new Map() private readonly settleFlights = new Map>() private readonly releaseFlights = new Map>() private readonly reclaimFlights = new Map>() @@ -141,16 +143,39 @@ export class MemoryPaymentOperations implements PaymentOperations { // The map write is synchronous. No second caller can observe a free nonce // between the uniqueness check and ownership write. this.operations.set(operationId, operation) + const claimToken = globalThis.crypto.randomUUID() + this.claimTokens.set(operationId, claimToken) + const claim = (async () => { + try { + await this.options.onClaim?.(operation) + if (this.claimTokens.get(operationId) !== claimToken) { + throw new Error('payment claim ownership was reclaimed') + } + const current = this.operations.get(operationId) + if (!current || current.state !== 'claiming') { + throw new Error('payment claim ownership was lost') + } + const claimed = { ...current, state: 'claimed' as const } + this.operations.set(operationId, claimed) + return claimed + } catch (error) { + // Keep the durable `claiming` row. The external authorization may have + // committed before its acknowledgement was lost, so expiry recovery + // must own the decision to reclaim it. + throw error + } finally { + if (this.claimTokens.get(operationId) === claimToken) { + this.claimTokens.delete(operationId) + } + } + })() + this.claimFlights.set(operationId, claim) try { - await this.options.onClaim?.(operation) - const claimed = { ...operation, state: 'claimed' as const } - this.operations.set(operationId, claimed) - return claimed - } catch (error) { - // Keep the durable `claiming` row. The external authorization may have - // committed before its acknowledgement was lost, so expiry recovery - // must own the decision to reclaim it. - throw error + return await claim + } finally { + if (this.claimFlights.get(operationId) === claim) { + this.claimFlights.delete(operationId) + } } } @@ -301,6 +326,11 @@ export class MemoryPaymentOperations implements PaymentOperations { const flight = this.reclaimFlights.get(operationId) if (flight) return flight if (current.expiresAt > this.now()) throw new Error('payment operation has not expired') + if (current.state === 'claiming') { + // Invalidate the live claim before starting recovery. Its callback may + // still finish, but it can no longer promote the operation to claimed. + this.claimTokens.delete(operationId) + } const reclaimable = { ...current, state: 'reclaimable' as const } this.operations.set(operationId, reclaimable) const recovery = this.runReclaim(reclaimable, 'reclaimed') diff --git a/src/payment-recovery-worker.ts b/src/payment-recovery-worker.ts index 1496881..32e6ebd 100644 --- a/src/payment-recovery-worker.ts +++ b/src/payment-recovery-worker.ts @@ -15,6 +15,8 @@ import type { AgentMeta, GatewayConfig, SandboxUsageReceipt } from './types' interface RecoveryWorkerOptions { now?: number + /** Fresh wall-clock source for each row's lease and retry timestamps. */ + clock?: () => number workerId?: string } @@ -44,18 +46,18 @@ export async function recoverPayments( ): Promise { const recovery = config.paymentRecovery if (!recovery) throw new Error('durable payment recovery is not configured') - const now = options.now ?? Date.now() + const scanNow = options.now ?? recoveryNow(options) const limit = options.limit ?? 100 if (!Number.isSafeInteger(limit) || limit <= 0) { throw new Error('payment recovery limit must be a positive safe integer') } - const due = await recovery.store.listDue(now, limit) + const due = await recovery.store.listDue(scanNow, limit) const run: PaymentRecoveryRun = { scanned: due.length, reconciled: 0, deferred: 0, failed: 0 } for (const candidate of due) { try { const result = await recoverPayment(candidate.id, config, { - now, force: false, + ...(options.clock ? { clock: options.clock } : {}), ...(options.workerId ? { workerId: options.workerId } : {}), }) if (result?.state === 'reconciled') run.reconciled += 1 @@ -75,10 +77,11 @@ export async function recoverPayment( ): Promise { const recovery = config.paymentRecovery if (!recovery) throw new Error('durable payment recovery is not configured') - const now = options.now ?? Date.now() + const scanNow = options.now ?? recoveryNow(options) const current = await recovery.store.get(recoveryId) if (!current || current.state === 'reconciled') return current - if (!options.force && current.nextAttemptAt > now) return current + if (!options.force && current.nextAttemptAt > scanNow) return current + const now = recoveryNow(options) const leased = await acquireLease( current.id, @@ -96,14 +99,15 @@ export async function recoverPayment( await reconcileLeased(ready, config, now) } catch (error) { const message = error instanceof Error ? error.message : String(error) + const failedAt = recoveryNow(options) try { await updateOwnedPaymentRecovery(recovery.store, recoveryId, fenceId, (record) => ({ ...record, attempts: record.attempts + 1, lastError: message, lease: undefined, - nextAttemptAt: now + recoveryTiming(recovery).retryDelayMs, - }), now) + nextAttemptAt: failedAt + recoveryTiming(recovery).retryDelayMs, + }), failedAt) } catch (updateError) { if (!(updateError instanceof PaymentRecoveryFenceError)) throw updateError } @@ -112,6 +116,10 @@ export async function recoverPayment( return recovery.store.get(recoveryId) } +function recoveryNow(options: RecoveryWorkerOptions): number { + return options.clock?.() ?? Date.now() +} + async function reconcileLeased( record: PaymentRecoveryRecord, config: GatewayConfig, diff --git a/tests/payment-operations.test.ts b/tests/payment-operations.test.ts index b8b0e41..eacc2ee 100644 --- a/tests/payment-operations.test.ts +++ b/tests/payment-operations.test.ts @@ -159,6 +159,37 @@ describe('version 2 payment operations', () => { expect(recoveredCalls).toBe(1) }) + it('fences a delayed claim before expiry recovery can reclaim it', async () => { + let now = 100 + let claimStarted!: () => void + let finishClaim!: () => void + const claimReady = new Promise((resolve) => { claimStarted = resolve }) + const claimFinished = new Promise((resolve) => { finishClaim = resolve }) + const events: string[] = [] + const operations = new MemoryPaymentOperations({ + now: () => now, + onClaim: async () => { + events.push('claim-start') + claimStarted() + await claimFinished + events.push('claim-finished') + }, + onReclaim: async () => { events.push('reclaim') }, + }) + const operationId = 'x402:0x' + 'ab'.repeat(32) + ':23' + const claim = operations.claimPayment(payload('1000', '23', '200'), context()) + await claimReady + + now = 201 + const reclaimed = await operations.reclaimPayment(operationId) + expect(reclaimed.state).toBe('reclaimed') + finishClaim() + + await expect(claim).rejects.toThrow('payment claim ownership was reclaimed') + expect(operations.get(operationId)?.state).toBe('reclaimed') + expect(events).toEqual(['claim-start', 'reclaim', 'claim-finished']) + }) + it('does not refund active or retained work after authorization expiry', async () => { let now = 100 const operations = new MemoryPaymentOperations({ now: () => now }) diff --git a/tests/payment-recovery.test.ts b/tests/payment-recovery.test.ts index 43412f5..745a0cb 100644 --- a/tests/payment-recovery.test.ts +++ b/tests/payment-recovery.test.ts @@ -8,8 +8,15 @@ import type { MppChargeOperation, } from '../src/mpp-payment' import { MemoryNonceStore } from '../src/nonce-store' -import { MemoryPaymentOperations, type PaymentSettlementInput } from '../src/payment-operations' -import { MemoryPaymentRecoveryStore } from '../src/payment-recovery' +import { + MemoryPaymentOperations, + type PaymentOperation, + type PaymentSettlementInput, +} from '../src/payment-operations' +import { + MemoryPaymentRecoveryStore, + type PaymentRecoveryRecord, +} from '../src/payment-recovery' import { recoverPayment, recoverPayments } from '../src/payment-recovery-worker' import type { AgentMeta, @@ -80,6 +87,67 @@ function requestBody(message = 'run') { return JSON.stringify({ max_tokens: 4, messages: [{ role: 'user', content: message }] }) } +function pendingMppRecovery( + id: string, + now: number, + nextAttemptAt = now, +): PaymentRecoveryRecord { + return { + version: 1, + id, + revision: 0, + state: 'claiming', + payment: { kind: 'mpp-charge', method: 'stripe', operationId: id }, + attribution: { + requestId: `request-${id}`, + agentId: agent.id, + agentSlug: agent.slug, + consumerId: 'stripe:customer', + paymentMethod: 'mpp', + startMs: now, + pricePerTokenUsd: agent.pricePerTokenUsd, + platformFeePercent: agent.platformFeePercent, + requiredAmount: '1', + currencyDecimals: 6, + maxOutputTokens: 1, + executionBudget: { + maxInputTokens: 1, + maxOutputTokens: 1, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 0.000001, + }, + }, + workStarted: false, + usageRecorded: false, + attempts: 0, + nextAttemptAt, + createdAt: now, + updatedAt: now, + } +} + +function mppRecoveryConfig( + store: MemoryPaymentRecoveryStore, + charge: MppChargeLifecycle, + timing: { leaseMs?: number; retryDelayMs?: number } = {}, +): GatewayConfig { + return { + resolveAgent: async () => agent, + getSandbox: async () => ({ async *streamPrompt() {} }), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + mpp: { + realm: 'gateway.test', + method: 'stripe', + authenticateCredential: async () => ({ consumerId: 'unused', paymentIdentity: 'unused' }), + charge, + }, + paymentRecovery: { store, ...timing }, + } +} + class HostileChargeLifecycle implements MppChargeLifecycle { readonly protocolVersion = 1 as const readonly operations = new Map() @@ -130,6 +198,21 @@ class HostileChargeLifecycle implements MppChargeLifecycle { } } +class FailOnceRecoveryStore extends MemoryPaymentRecoveryStore { + private failed = false + + override async compareAndSet( + expected: PaymentRecoveryRecord, + next: PaymentRecoveryRecord, + ): Promise { + if (!this.failed && next.state === 'releasing') { + this.failed = true + throw new Error('payment metadata unavailable') + } + return super.compareAndSet(expected, next) + } +} + describe('generic MPP charge lifecycle', () => { it('confirms only after all denials and before any response, body, or sandbox work', async () => { const deniedLifecycle = new HostileChargeLifecycle() @@ -270,6 +353,90 @@ describe('generic MPP charge lifecycle', () => { expect(runs).toBe(0) }) + it('starts each batch lease from the time that row begins recovery', async () => { + vi.useFakeTimers() + const startedAt = new Date('2026-08-14T00:00:00.000Z').getTime() + vi.setSystemTime(startedAt) + const recoveryStore = new MemoryPaymentRecoveryStore() + const leaseRemaining: number[] = [] + const lifecycle: MppChargeLifecycle = { + protocolVersion: 1, + async confirmPayment() { + throw new Error('confirmation is not used during recovery') + }, + async releasePayment(operation) { + return operation + }, + async recoverPayment(operationId) { + if (operationId === 'slow') vi.setSystemTime(startedAt + 11) + if (operationId === 'next') { + const current = await recoveryStore.get(operationId) + leaseRemaining.push((current?.lease?.expiresAt ?? 0) - Date.now()) + } + return { operationId, state: 'not-found' } + }, + } + const config = mppRecoveryConfig(recoveryStore, lifecycle, { leaseMs: 10 }) + await recoveryStore.createIfAbsent(pendingMppRecovery('slow', startedAt, startedAt - 2)) + await recoveryStore.createIfAbsent(pendingMppRecovery('next', startedAt, startedAt - 1)) + + expect(await recoverPayments(config)).toMatchObject({ scanned: 2, reconciled: 2 }) + expect(leaseRemaining).toEqual([10]) + }) + + it('starts the retry delay after a failed provider recovery call', async () => { + vi.useFakeTimers() + const startedAt = new Date('2026-08-14T00:00:00.000Z').getTime() + vi.setSystemTime(startedAt) + const recoveryStore = new MemoryPaymentRecoveryStore() + const lifecycle: MppChargeLifecycle = { + protocolVersion: 1, + async confirmPayment() { + throw new Error('confirmation is not used during recovery') + }, + async releasePayment(operation) { + return operation + }, + async recoverPayment() { + vi.setSystemTime(startedAt + 11) + throw new Error('provider unavailable') + }, + } + const config = mppRecoveryConfig(recoveryStore, lifecycle, { retryDelayMs: 10 }) + await recoveryStore.createIfAbsent(pendingMppRecovery('failed', startedAt)) + + await expect(recoverPayment('failed', config)).rejects.toThrow('provider unavailable') + expect((await recoveryStore.get('failed'))?.nextAttemptAt).toBe(startedAt + 21) + }) + + it('uses a fresh row clock after a supplied batch scan time', async () => { + const suppliedNow = 1_000 + let current = suppliedNow + const clock = () => current + const recoveryStore = new MemoryPaymentRecoveryStore() + const lifecycle: MppChargeLifecycle = { + protocolVersion: 1, + async confirmPayment() { + throw new Error('confirmation is not used during recovery') + }, + async releasePayment(operation) { + return operation + }, + async recoverPayment() { + current += 11 + throw new Error('provider unavailable') + }, + } + const config = mppRecoveryConfig(recoveryStore, lifecycle, { retryDelayMs: 10 }) + await recoveryStore.createIfAbsent(pendingMppRecovery('supplied-now', suppliedNow)) + + await expect(recoverPayments(config, { now: suppliedNow, clock })).resolves.toMatchObject({ + scanned: 1, + failed: 1, + }) + expect((await recoveryStore.get('supplied-now'))?.nextAttemptAt).toBe(1_021) + }) + it('fences a live request when recovery releases its expired claim lease', async () => { vi.useFakeTimers() vi.setSystemTime(new Date('2026-08-14T00:00:00.000Z')) @@ -487,6 +654,95 @@ describe('generic MPP charge lifecycle', () => { }) describe('durable OpenAI recovery', () => { + it('rejects a v2 adapter operation with a mismatched identity and clears its claim lease', async () => { + const recoveryStore = new MemoryPaymentRecoveryStore() + let sandboxCalls = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + sandboxCalls += 1 + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations(), + authorizePayment: async (_payload, context): Promise => ({ + protocolVersion: 2, + operationId: 'adapter-operation-id', + acquiredByRequestId: context.requestId, + nonceKey: 'adapter-nonce', + authorizationId: 'adapter-authorization', + reservedAmount: 1_000_000_000n, + settledAmount: 0n, + refundAmount: 1_000_000_000n, + expiresAt: Math.floor(Date.now() / 1000) + 600, + state: 'claimed', + }), + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore }, + } + const app = mount(config) + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json', 'X-Payment-Signature': spendAuth('906') }, + body: requestBody(), + }) + await response.text() + + const record = await recoveryStore.get(`x402:${commitment}:906`) + expect(response.status).toBe(402) + expect(record?.state).toBe('claiming') + expect(record?.lease).toBeUndefined() + expect(record?.payment.operationId).toBe(`x402:${commitment}:906`) + expect(sandboxCalls).toBe(0) + }) + + it('clears the active recovery lease when release metadata fails', async () => { + const recoveryStore = new FailOnceRecoveryStore() + const operationId = `x402:${commitment}:907` + class FailingExecutionOperations extends MemoryPaymentOperations { + override async beginPaymentExecution(_operation: PaymentOperation): Promise { + throw new Error('execution unavailable') + } + } + const operations = new FailingExecutionOperations() + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ async *streamPrompt() {} }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore }, + } + const app = mount(config) + const response = await app.request('/v1/agents/recovery/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': spendAuth('907'), + }, + body: requestBody(), + }) + await response.text() + + const pending = await recoveryStore.get(operationId) + expect(pending?.state).toBe('executing') + expect(pending?.lease).toBeUndefined() + expect(operations.get(operationId)?.state).toBe('claimed') + }) + it('persists operation and attribution before output, then heals settlement acknowledgement loss', async () => { const recoveryStore = new MemoryPaymentRecoveryStore() let settleCalls = 0 From 45468d1b667159c187017f748a5602de42e8436a Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 23:58:57 -0600 Subject: [PATCH 21/32] fix(a2a): harden task lifecycle recovery --- src/a2a/handler.ts | 700 +++++++++++++++++++++----------- src/a2a/task-store.ts | 1 + tests/a2a-atomicity.test.ts | 2 + tests/a2a-lifecycle.test.ts | 399 ++++++++++++++++++ tests/a2a-payment-races.test.ts | 5 +- 5 files changed, 866 insertions(+), 241 deletions(-) create mode 100644 tests/a2a-lifecycle.test.ts diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 157f3dd..cd69b68 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -44,7 +44,7 @@ import { type PushNotificationStore, type TaskPushNotificationConfig, } from './push-notifications' -import type { TaskStore } from './task-store' +import { hasPendingPaymentRecovery, type TaskStore } from './task-store' import { extractTextFromMessage, responseTextToArtifact } from './translate' import { A2A_ERROR_CODES, @@ -108,6 +108,21 @@ interface TaskPaymentRecoveryMarker { id: string } +interface TaskOriginBinding { + version: 1 + agentId: string + agentSlug: string +} + +interface TaskSubmissionRecord { + version: 1 + lease: { id: string; expiresAt: number } + agentId: string + agentSlug: string + requestId: string + consumerId: string +} + /** Terminal task states — fire-once push delivery occurs on these transitions. */ const TERMINAL_STATES: ReadonlySet = new Set([ 'completed', @@ -116,6 +131,11 @@ const TERMINAL_STATES: ReadonlySet = new Set([ 'rejected', ]) +const TASK_ORIGIN_METADATA_KEY = 'gatewayOrigin' +const TASK_SUBMISSION_METADATA_KEY = 'gatewaySubmission' +const TASK_SUBMISSION_RECOVERY_METADATA_KEY = 'gatewaySubmissionRecovery' +const TASK_SUBMISSION_LEASE_MS = 5 * 60 * 1000 + /** * Per-gateway in-process registry of cancellable runs. Keyed by task id; * absent = task already terminal or never streamed. Cleared by the streaming @@ -245,9 +265,11 @@ async function handleMessageSend( const { authz, task } = guard setPaymentResponseHeaders(c, authz) const controller = cancels.register(task.id) + const detachRequestAbort = bindRequestAbort(c.req.raw.signal, controller) try { return await executeMessageSend(c, req, deps, authz, task, controller.signal) } finally { + detachRequestAbort() cancels.clear(task.id) } } @@ -260,6 +282,7 @@ async function executeMessageSend( task: Task, signal: AbortSignal, ): Promise { + if (isTerminal(task.status.state)) return c.json(ok(req.id, task)) const workingTask: Task = task.status.state === 'working' ? task : { ...task, status: { state: 'working', timestamp: nowIso() } } @@ -319,12 +342,12 @@ async function executeMessageSend( workObserved || usage !== undefined, ) const currentTask = await deps.taskStore.get(task.id) ?? releasedTask - const failed = isTerminal(currentTask.status.state) + const failed = shouldPreserveTask(currentTask) ? currentTask - : withStatus(currentTask, 'failed') + : withStatus(clearTaskSubmission(currentTask), 'failed') try { - await deps.taskStore.put(failed) - await maybeDeliverPush(failed, deps) + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) + await maybeDeliverPush(persisted, deps) } catch (taskError) { console.error( `[a2a] failed to persist failed task ${task.id}:`, @@ -432,12 +455,12 @@ async function executeMessageSend( return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) } const currentTask = await deps.taskStore.get(task.id) ?? releasedTask - const failed = isTerminal(currentTask.status.state) + const failed = shouldPreserveTask(currentTask) ? currentTask - : withStatus(currentTask, 'failed') + : withStatus(clearTaskSubmission(currentTask), 'failed') try { - await deps.taskStore.put(failed) - await maybeDeliverPush(failed, deps) + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) + await maybeDeliverPush(persisted, deps) } catch (taskError) { console.error( `[a2a] failed to persist failed task ${task.id}:`, @@ -463,6 +486,7 @@ async function handleMessageStream( setPaymentResponseHeaders(c, authz) const controller = cancels.register(task.id) + const detachRequestAbort = bindRequestAbort(c.req.raw.signal, controller) const workingStatus: TaskStatusUpdateEvent = { kind: 'status-update', taskId: task.id, @@ -470,10 +494,16 @@ async function handleMessageStream( status: { state: 'working', timestamp: nowIso() }, final: false, } + if (isTerminal(task.status.state)) { + detachRequestAbort() + cancels.clear(task.id) + return c.json(ok(req.id, task)) + } const workingTask: Task = task.status.state === 'working' ? task : { ...task, status: workingStatus.status } if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { + detachRequestAbort() cancels.clear(task.id) const current = await releaseTaskPayment( authz, @@ -488,6 +518,7 @@ async function handleMessageStream( try { await beginPaymentExecution(authz, deps.config) } catch (error) { + detachRequestAbort() cancels.clear(task.id) await releaseTaskPayment( authz, @@ -499,6 +530,7 @@ async function handleMessageStream( return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment execution authorization failed')) } if (controller.signal.aborted) { + detachRequestAbort() cancels.clear(task.id) const canceled = await completeCanceledTask(authz, workingTask, '', undefined, false, deps) return c.json(ok(req.id, canceled)) @@ -508,100 +540,71 @@ async function handleMessageStream( let workObserved = false const stream = new ReadableStream({ - async start(ctrl) { - const encoder = new TextEncoder() - const send = (event: StreamingEvent) => { - ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) - } - - let inputRequiredPrompt: string | undefined - let inputRequiredSeen = false - let finalizationLeaseId: string | undefined - try { - send(workingStatus) - - for await (const event of dispatchSandboxStreamRich( - authz.agent, - authz.userMessage, - authz.consumerId, - deps.config, - controller.signal, - task.id, - authz.maxOutputTokens, - undefined, - authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, - () => { - workObserved = true - }, - )) { - if (event.kind === 'text') { - responseText += event.delta - workObserved = true - const artifactEvent: TaskArtifactUpdateEvent = { - kind: 'artifact-update', - taskId: task.id, - contextId: task.contextId, - artifact: { - artifactId: `${task.id}-artifact-0`, - name: 'response', - parts: [{ kind: 'text', text: event.delta }], - }, - append: true, - } - send(artifactEvent) - } else if (event.kind === 'activity') { - workObserved = true - } else if (event.kind === 'usage') { - usage = event.usage - } else { - inputRequiredSeen = true - inputRequiredPrompt = event.prompt - workObserved = true + start(ctrl) { + void (async () => { + const encoder = new TextEncoder() + const send = (event: StreamingEvent) => { + if (ctrl.desiredSize === null) return + try { + ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) + } catch { + // The client can cancel between the desiredSize check and enqueue. } } - // Caller aborted via tasks/cancel. Charge a complete receipt if one - // exists; otherwise retain ownership when output or hidden work was - // observed because releasing would make paid work free. - if (controller.signal.aborted) { - const canceled = await completeCanceledTask( - authz, - workingTask, - responseText, - usage, - workObserved, - deps, - ) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: canceled.status, - final: true, - }) - return - } + let inputRequiredPrompt: string | undefined + let inputRequiredSeen = false + let finalizationLeaseId: string | undefined + try { + send(workingStatus) + + for await (const event of dispatchSandboxStreamRich( + authz.agent, + authz.userMessage, + authz.consumerId, + deps.config, + controller.signal, + task.id, + authz.maxOutputTokens, + undefined, + authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, + () => { + workObserved = true + }, + )) { + if (event.kind === 'text') { + responseText += event.delta + workObserved = true + const artifactEvent: TaskArtifactUpdateEvent = { + kind: 'artifact-update', + taskId: task.id, + contextId: task.contextId, + artifact: { + artifactId: `${task.id}-artifact-0`, + name: 'response', + parts: [{ kind: 'text', text: event.delta }], + }, + append: true, + } + send(artifactEvent) + } else if (event.kind === 'activity') { + workObserved = true + } else if (event.kind === 'usage') { + usage = event.usage + } else { + inputRequiredSeen = true + inputRequiredPrompt = event.prompt + workObserved = true + } + } - // Settle once for whatever the sandbox produced (full or partial). - if (!usage) throw new Error('sandbox did not provide a usage receipt') - const finalizationArtifact = responseTextToArtifact(responseText, `${task.id}-artifact-0`) - const finalization = buildFinalizationRecord( - authz, - usage, - finalizationArtifact, - inputRequiredSeen, - inputRequiredPrompt, - ) - // Let the durable task-store CAS decide the cancellation race. Mark - // the local registry only after that CAS wins, so cancel can replace a - // still-pending finalization instead of being rejected prematurely. - const finalizingTask = withFinalizationRecord(workingTask, finalization) - if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask?.status.state === 'canceled') { + // Caller aborted via tasks/cancel. Charge a complete receipt if one + // exists; otherwise retain ownership when output or hidden work was + // observed because releasing would make paid work free. + if (controller.signal.aborted) { const canceled = await completeCanceledTask( authz, - currentTask, + workingTask, responseText, usage, workObserved, @@ -616,35 +619,98 @@ async function handleMessageStream( }) return } - await releaseTaskPayment( + + // Settle once for whatever the sandbox produced (full or partial). + if (!usage) throw new Error('sandbox did not provide a usage receipt') + const finalizationArtifact = responseTextToArtifact(responseText, `${task.id}-artifact-0`) + const finalization = buildFinalizationRecord( authz, - task, - deps, - 'A2A task changed before payment settlement', - workObserved || usage !== undefined, + usage, + finalizationArtifact, + inputRequiredSeen, + inputRequiredPrompt, ) - return - } - finalizationLeaseId = finalization.lease.id - cancels.beginFinalization(task.id) - let usageRecordedTask = finalizingTask - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { - onUsageRecorded: async () => { - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - }, - }) - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + // Let the durable task-store CAS decide the cancellation race. Mark + // the local registry only after that CAS wins, so cancel can replace a + // still-pending finalization instead of being rejected prematurely. + const finalizingTask = withFinalizationRecord(workingTask, finalization) + if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask?.status.state === 'canceled') { + const canceled = await completeCanceledTask( + authz, + currentTask, + responseText, + usage, + workObserved, + deps, + ) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: canceled.status, + final: true, + }) + return + } + await releaseTaskPayment( + authz, + task, + deps, + 'A2A task changed before payment settlement', + workObserved || usage !== undefined, + ) + return + } + finalizationLeaseId = finalization.lease.id + cancels.beginFinalization(task.id) + let usageRecordedTask = finalizingTask + await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - if (inputRequiredSeen) { - const paused = withStatus( - clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), - 'input-required', - inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, - responseText - ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] - : task.artifacts, - ) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, paused)) { + if (inputRequiredSeen) { + const paused = withStatus( + clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), + 'input-required', + inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, + responseText + ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] + : task.artifacts, + ) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, paused)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: currentTask.status, + final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', + }) + } + return + } + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: paused.status, + final: true, + }) + // input-required is non-terminal — do NOT deliver push notifications. + return + } + + // Final: persist the terminal task before emitting terminal events. + const completed = withStatus(clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), 'completed', undefined, [ + responseTextToArtifact(responseText, `${task.id}-artifact-0`), + ]) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, completed)) { const currentTask = await deps.taskStore.get(task.id) if (currentTask) { send({ @@ -657,122 +723,103 @@ async function handleMessageStream( } return } + send({ + kind: 'artifact-update', + taskId: task.id, + contextId: task.contextId, + artifact: { + artifactId: `${task.id}-artifact-0`, + name: 'response', + parts: [{ kind: 'text', text: '' }], + }, + append: true, + lastChunk: true, + }) send({ kind: 'status-update', taskId: task.id, contextId: task.contextId, - status: paused.status, + status: completed.status, final: true, }) - // input-required is non-terminal — do NOT deliver push notifications. - return - } - - // Final: persist the terminal task before emitting terminal events. - const completed = withStatus(clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), 'completed', undefined, [ - responseTextToArtifact(responseText, `${task.id}-artifact-0`), - ]) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, completed)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask) { - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: currentTask.status, - final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', - }) - } - return - } - send({ - kind: 'artifact-update', - taskId: task.id, - contextId: task.contextId, - artifact: { - artifactId: `${task.id}-artifact-0`, - name: 'response', - parts: [{ kind: 'text', text: '' }], - }, - append: true, - lastChunk: true, - }) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: completed.status, - final: true, - }) - await maybeDeliverPush(completed, deps) - } catch (err) { - const releasedTask = await releaseTaskPayment( - authz, - task, - deps, - err instanceof Error ? err.message : String(err), - workObserved || usage !== undefined, - ) - if (finalizationLeaseId) { - const retained = await retainFinalizationForRecovery( - deps.taskStore, - task.id, - finalizationLeaseId, - asError(err), + await maybeDeliverPush(completed, deps) + } catch (err) { + const releasedTask = await releaseTaskPayment( + authz, + task, + deps, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, ) - if (retained) { + if (finalizationLeaseId) { + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalizationLeaseId, + asError(err), + ) + if (retained) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: retained.status, + final: false, + }) + } + return + } + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = shouldPreserveTask(currentTask) + ? currentTask + : withStatus(clearTaskSubmission(currentTask), 'failed') + try { + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) send({ kind: 'status-update', taskId: task.id, contextId: task.contextId, - status: retained.status, - final: false, + status: persisted.status, + final: true, }) + await maybeDeliverPush(persisted, deps) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + try { + await deps.state.obs?.onStreamError?.( + { + requestId: authz.requestId, + agentSlug: authz.agent.slug, + startMs: authz.startMs, + }, + { + consumerId: authz.consumerId, + errorMessage: err instanceof Error ? err.message : String(err), + }, + ) + } catch (observerError) { + console.error( + `[a2a] stream observer failed for ${authz.requestId}:`, + observerError instanceof Error ? observerError.message : String(observerError), + ) + } + } finally { + detachRequestAbort() + cancels.clear(task.id) + try { + if (ctrl.desiredSize !== null) ctrl.close() + } catch { + // The response may already be closed by client cancellation. } - return - } - const currentTask = await deps.taskStore.get(task.id) ?? releasedTask - const failed = isTerminal(currentTask.status.state) - ? currentTask - : withStatus(currentTask, 'failed') - try { - await deps.taskStore.put(failed) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: failed.status, - final: true, - }) - await maybeDeliverPush(failed, deps) - } catch (taskError) { - console.error( - `[a2a] failed to persist failed task ${task.id}:`, - taskError instanceof Error ? taskError.message : String(taskError), - ) - } - try { - await deps.state.obs?.onStreamError?.( - { - requestId: authz.requestId, - agentSlug: authz.agent.slug, - startMs: authz.startMs, - }, - { - consumerId: authz.consumerId, - errorMessage: err instanceof Error ? err.message : String(err), - }, - ) - } catch (observerError) { - console.error( - `[a2a] stream observer failed for ${authz.requestId}:`, - observerError instanceof Error ? observerError.message : String(observerError), - ) } - } finally { - cancels.clear(task.id) - ctrl.close() - } + })() + }, + cancel() { + controller.abort() }, }) @@ -1125,6 +1172,7 @@ async function guardMessageRequest( ...existing, status: { state: 'submitted', timestamp: nowIso() }, history: [...(existing.history ?? []), appendedMessage], + metadata: withTaskSubmission(existing.metadata, authz), } if (!await compareAndSetTask(deps.taskStore, existing, continued)) { return c.json( @@ -1152,6 +1200,7 @@ async function guardMessageRequest( contextId, status: { state: 'submitted', timestamp: nowIso() }, history: [initialMessage], + metadata: withTaskSubmission(withTaskOrigin(undefined, authz.agent), authz), } if (!await createTask(deps.taskStore, task)) { return c.json( @@ -1203,16 +1252,17 @@ async function claimTaskPayment( 'payment authorization failed', false, ) - const releaseRecord = releasedTask.metadata?.[PAYMENT_RELEASE_METADATA_KEY] + const cleanedReleasedTask = clearTaskSubmission(releasedTask) + const releaseRecord = cleanedReleasedTask.metadata?.[PAYMENT_RELEASE_METADATA_KEY] const failed = releaseRecord !== undefined - ? isTerminal(releasedTask.status.state) - ? releasedTask - : withStatus(releasedTask, 'failed') + ? isTerminal(cleanedReleasedTask.status.state) + ? cleanedReleasedTask + : withStatus(cleanedReleasedTask, 'failed') : paymentFailureTask - ? preservePaymentRecoveryMarker(paymentFailureTask, releasedTask) - : isTerminal(releasedTask.status.state) - ? releasedTask - : withStatus(releasedTask, 'failed') + ? preservePaymentRecoveryMarker(paymentFailureTask, cleanedReleasedTask) + : isTerminal(cleanedReleasedTask.status.state) + ? cleanedReleasedTask + : withStatus(cleanedReleasedTask, 'failed') try { if (await compareAndSetTask(deps.taskStore, releasedTask, failed) && isTerminal(failed.status.state)) { await maybeDeliverPush(failed, deps) @@ -1226,6 +1276,8 @@ async function claimTaskPayment( return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment authorization failed')) } + const submissionClear = await clearTaskSubmissionMarker(deps.taskStore, paymentTask) + if (submissionClear.applied) paymentTask = submissionClear.task const current = await deps.taskStore.get(task.id) if (current && JSON.stringify(current) === JSON.stringify(paymentTask)) { return paymentTask @@ -1434,9 +1486,9 @@ async function completeCanceledTask( artifacts: [responseTextToArtifact(responseText, `${task.id}-artifact-0`)], } : canceledBase - await deps.taskStore.put(canceled) - await maybeDeliverPush(canceled, deps) - return canceled + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask ?? task, canceled) + await maybeDeliverPush(persisted, deps) + return persisted } // ── Helpers ─────────────────────────────────────────────────────────────── @@ -1462,6 +1514,19 @@ async function authorizeTaskAccess( task: Task, deps: A2AHandlerDeps, ): Promise { + const requestedAgentSlug = c.req.param('slug') ?? '' + const origin = readTaskOrigin(task) + if (origin) { + if (origin.agentSlug !== requestedAgentSlug) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task belongs to a different agent'), 403) + } + const requestedAgent = await deps.config.resolveAgent(requestedAgentSlug) + if (!requestedAgent || requestedAgent.id !== origin.agentId) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task belongs to a different agent'), 403) + } + } else if (!deps.config.x402.demoMode) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task origin is not recorded'), 403) + } const authorize = deps.config.a2a?.authorizeTaskAccess if (!authorize && deps.config.x402.demoMode) return undefined if (!authorize) { @@ -1474,7 +1539,7 @@ async function authorizeTaskAccess( try { allowed = await authorize(task, { method: req.method, - agentSlug: c.req.param('slug') ?? '', + agentSlug: requestedAgentSlug, authorization: c.req.header('Authorization') ?? '', paymentSignature: c.req.header('X-Payment-Signature') ?? '', }) @@ -1497,6 +1562,134 @@ function isHttpsUrl(value: string): boolean { } } +function bindRequestAbort(requestSignal: AbortSignal, controller: AbortController): () => void { + const abort = () => controller.abort() + if (requestSignal.aborted) abort() + else requestSignal.addEventListener('abort', abort, { once: true }) + return () => requestSignal.removeEventListener('abort', abort) +} + +function withTaskOrigin( + metadata: Record | undefined, + agent: { id: string; slug: string }, +): Record { + return { + ...(metadata ?? {}), + [TASK_ORIGIN_METADATA_KEY]: { + version: 1, + agentId: agent.id, + agentSlug: agent.slug, + } satisfies TaskOriginBinding, + } +} + +function withTaskSubmission( + metadata: Record | undefined, + authz: AuthorizedRequest, +): Record { + const origin = metadata?.[TASK_ORIGIN_METADATA_KEY] + return { + ...(metadata ?? {}), + ...(origin === undefined + ? { + [TASK_ORIGIN_METADATA_KEY]: { + version: 1, + agentId: authz.agent.id, + agentSlug: authz.agent.slug, + } satisfies TaskOriginBinding, + } + : {}), + [TASK_SUBMISSION_METADATA_KEY]: { + version: 1, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + TASK_SUBMISSION_LEASE_MS }, + agentId: authz.agent.id, + agentSlug: authz.agent.slug, + requestId: authz.requestId, + consumerId: authz.consumerId, + } satisfies TaskSubmissionRecord, + } +} + +function readTaskOrigin(task: Task): TaskOriginBinding | undefined { + const raw = task.metadata?.[TASK_ORIGIN_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const origin = raw as Partial + if ( + origin.version !== 1 || + typeof origin.agentId !== 'string' || + origin.agentId.length === 0 || + typeof origin.agentSlug !== 'string' || + origin.agentSlug.length === 0 + ) { + return undefined + } + return origin as TaskOriginBinding +} + +function readTaskSubmission(task: Task): TaskSubmissionRecord | undefined { + const raw = task.metadata?.[TASK_SUBMISSION_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const submission = raw as Partial + if ( + submission.version !== 1 || + !submission.lease || + typeof submission.lease.id !== 'string' || + submission.lease.id.length === 0 || + typeof submission.lease.expiresAt !== 'number' || + !Number.isFinite(submission.lease.expiresAt) || + typeof submission.agentId !== 'string' || + submission.agentId.length === 0 || + typeof submission.agentSlug !== 'string' || + submission.agentSlug.length === 0 || + typeof submission.requestId !== 'string' || + submission.requestId.length === 0 || + typeof submission.consumerId !== 'string' + ) { + return undefined + } + return submission as TaskSubmissionRecord +} + +function clearTaskSubmission(task: Task): Task { + if (!task.metadata || !(TASK_SUBMISSION_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[TASK_SUBMISSION_METADATA_KEY] + return Object.keys(metadata).length > 0 + ? { ...task, metadata } + : (() => { + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata + })() +} + +async function clearTaskSubmissionMarker( + taskStore: TaskStore, + expected: Task, +): Promise<{ task: Task; applied: boolean }> { + const current = await taskStore.get(expected.id) + if (!current || JSON.stringify(current) !== JSON.stringify(expected)) { + return { task: current ?? expected, applied: false } + } + const cleared = clearTaskSubmission(current) + if (cleared === current) return { task: current, applied: true } + if (await compareAndSetTask(taskStore, current, cleared)) return { task: cleared, applied: true } + return { task: await taskStore.get(expected.id) ?? expected, applied: false } +} + +function shouldPreserveTask(task: Task): boolean { + return isTerminal(task.status.state) || hasPendingPaymentRecovery(task) +} + +async function persistTaskIfCurrent( + taskStore: TaskStore, + expected: Task, + next: Task, +): Promise { + if (expected === next || JSON.stringify(expected) === JSON.stringify(next)) return expected + if (await compareAndSetTask(taskStore, expected, next)) return next + return await taskStore.get(expected.id) ?? expected +} + async function createTask(taskStore: TaskStore, task: Task): Promise { if (!taskStore.createIfAbsent) { throw new Error('A2A task store does not provide createIfAbsent') @@ -1932,8 +2125,6 @@ async function recoverFinalizationIfNeeded( throw new Error('A2A payment operation recovery is not configured') } paymentOperation = deserializePaymentOperation(renewed.paymentOperation) - } else if (!deps.config.x402.demoMode) { - throw new Error('A2A production recovery requires a durable payment operation') } const authz: AuthorizedRequest = { @@ -2007,7 +2198,36 @@ async function recoverTaskIfNeeded( ): Promise { const released = await recoverPaymentReleaseIfNeeded(task, deps) const finalized = await recoverFinalizationIfNeeded(released, deps, requestedAgentSlug) - return recoverPaymentMarkerIfNeeded(finalized, deps) + const paymentRecovered = await recoverPaymentMarkerIfNeeded(finalized, deps) + return recoverSubmissionIfNeeded(paymentRecovered, deps) +} + +async function recoverSubmissionIfNeeded(task: Task, deps: A2AHandlerDeps): Promise { + const raw = task.metadata?.[TASK_SUBMISSION_METADATA_KEY] + if (raw === undefined) return task + const submission = readTaskSubmission(task) + if (submission && submission.lease.expiresAt > Date.now()) return task + if (task.status.state !== 'submitted') { + return (await clearTaskSubmissionMarker(deps.taskStore, task)).task + } + + const cleanTask = clearTaskSubmission(task) + const failed: Task = { + ...withStatus(cleanTask, 'failed'), + metadata: { + ...(cleanTask.metadata ?? {}), + [TASK_SUBMISSION_RECOVERY_METADATA_KEY]: { + error: submission + ? 'A2A task submission lease expired before payment authorization completed' + : 'A2A task submission lease is invalid', + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await maybeDeliverPush(failed, deps) + return failed + } + return await deps.taskStore.get(task.id) ?? task } async function recoverPaymentMarkerIfNeeded( diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index c986273..7746875 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -23,6 +23,7 @@ const PAYMENT_RECOVERY_KEYS = [ 'gatewayFinalizing', 'gatewayPaymentRelease', 'gatewayPaymentRecovery', + 'gatewaySubmission', ] as const /** Recovery-bearing tasks must remain readable until reconciliation clears the marker. */ diff --git a/tests/a2a-atomicity.test.ts b/tests/a2a-atomicity.test.ts index 5dbf022..d700a21 100644 --- a/tests/a2a-atomicity.test.ts +++ b/tests/a2a-atomicity.test.ts @@ -243,6 +243,7 @@ describe('A2A task atomicity and restart recovery', () => { status: { state: 'working', timestamp: new Date().toISOString() }, artifacts: [artifact], metadata: { + gatewayOrigin: { version: 1, agentId: agent.id, agentSlug: agent.slug }, gatewayFinalizing: { version: 1, lease: { id: 'crashed-lease', expiresAt: Date.now() - 1 }, @@ -357,6 +358,7 @@ describe('A2A task atomicity and restart recovery', () => { contextId: 'ctx-retry', status: { state: 'working', timestamp: new Date().toISOString() }, metadata: { + gatewayOrigin: { version: 1, agentId: agent.id, agentSlug: agent.slug }, gatewayFinalizing: { version: 1, lease: { id: 'retry-lease', expiresAt: Date.now() - 1 }, diff --git a/tests/a2a-lifecycle.test.ts b/tests/a2a-lifecycle.test.ts new file mode 100644 index 0000000..5b2ff5b --- /dev/null +++ b/tests/a2a-lifecycle.test.ts @@ -0,0 +1,399 @@ +import { Hono } from 'hono' +import { describe, expect, it } from 'vitest' + +import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' +import { createAgentGateway } from '../src/middleware' +import { MemoryNonceStore } from '../src/nonce-store' +import type { AgentMeta, GatewayConfig, SandboxBox } from '../src/types' +import type { Task } from '../src/a2a/types' + +const operatorAddress = '0x1111111111111111111111111111111111111111' +const commitment = `0x${'ef'.repeat(32)}` + +const agentA: AgentMeta = { + id: 'agent-lifecycle-a', + ownerId: 'owner', + slug: 'lifecycle-a', + systemPrompt: '', + pricePerTokenUsd: 0.000001, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +const agentB: AgentMeta = { ...agentA, id: 'agent-lifecycle-b', slug: 'lifecycle-b' } + +const usage = { + inputTokens: 1, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.000002, + budgetEnforced: true, +} + +function paymentHeader(nonce: string): string { + return JSON.stringify({ + commitment, + signature: '0xsig', + operator: operatorAddress, + amount: '1000000000', + nonce, + expiry: String(Math.floor(Date.now() / 1000) + 300), + }) +} + +function message(taskId: string, text = 'hello') { + return { + kind: 'message', + role: 'user', + taskId, + contextId: `context-${taskId}`, + messageId: `message-${taskId}-${text}`, + parts: [{ kind: 'text', text }], + } +} + +function body(method: string, taskId: string, text = 'hello') { + return { + jsonrpc: '2.0', + id: `${method}-${taskId}`, + method, + params: method.startsWith('tasks/') ? { id: taskId } : { message: message(taskId, text) }, + } +} + +function standardSandbox(): SandboxBox { + return { + async *streamPrompt() { + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'done' } } + yield { type: 'sandbox.usage', data: { usage } } + }, + } +} + +function gatewayConfig(taskStore: TaskStore, overrides: Partial = {}): GatewayConfig { + return { + resolveAgent: async (slug) => (slug === agentA.slug ? agentA : null), + getSandbox: async () => standardSandbox(), + recordUsage: async () => undefined, + x402: { operatorAddress, chainId: 1, demoMode: true }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + ...overrides, + } +} + +async function post(app: Hono, slug: string, requestBody: unknown, headers: Record = {}) { + return app.request(`/v1/agents/${slug}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json', ...headers }, + body: JSON.stringify(requestBody), + }) +} + +describe('A2A lifecycle recovery and ownership', () => { + it('retries legacy settlement from the durable task record after a crash', async () => { + const taskStore = new InMemoryTaskStore() + let settlementCalls = 0 + let usageCalls = 0 + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + ...gatewayConfig(taskStore, { + resolveAgent: async (slug) => (slug === agentA.slug ? agentA : null), + recordUsage: async () => { usageCalls += 1 }, + settlePayment: async () => { + settlementCalls += 1 + if (settlementCalls === 1) throw new Error('legacy settlement acknowledgement lost') + }, + x402: { + operatorAddress, + chainId: 1, + paymentProtocolVersion: 1, + verifySigner: async () => true, + }, + }), + })) + + const first = await post( + app, + agentA.slug, + body('message/send', 'task-legacy-recovery'), + { 'X-Payment-Signature': paymentHeader('901') }, + ) + const firstBody = await first.json() as { error?: { code?: number } } + expect(firstBody.error?.code).toBe(-32603) + expect(settlementCalls).toBe(1) + + const retained = await taskStore.get('task-legacy-recovery') + expect(retained?.metadata?.gatewayFinalizing).toBeDefined() + const finalizing = retained?.metadata?.gatewayFinalizing as { + lease: { id: string; expiresAt: number } + } + await taskStore.put({ + ...retained!, + metadata: { + ...retained!.metadata, + gatewayFinalizing: { + ...finalizing, + lease: { ...finalizing.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + + const recovered = await post(app, agentA.slug, body('tasks/get', 'task-legacy-recovery')) + const recoveredBody = await recovered.json() as { result?: Task } + expect(recoveredBody.result?.status.state).toBe('completed') + expect(settlementCalls).toBe(2) + expect(usageCalls).toBe(1) + expect((await taskStore.get('task-legacy-recovery'))?.metadata?.gatewayFinalizing).toBeUndefined() + }) + + it('does not overwrite a newer cancellation or recovery marker on failure', async () => { + class InterleavingFailureStore implements TaskStore { + private readonly inner = new InMemoryTaskStore() + + async get(id: string) { + return this.inner.get(id) + } + + async put(task: Task) { + if (task.status.state === 'failed') { + const current = await this.inner.get(task.id) + if (current) { + await this.inner.put({ + ...current, + status: { state: 'canceled', timestamp: new Date().toISOString() }, + metadata: { + ...(current.metadata ?? {}), + gatewayPaymentRecovery: { version: 1, id: 'recovery-race' }, + }, + }) + } + } + await this.inner.put(task) + } + + async createIfAbsent(task: Task) { + return this.inner.createIfAbsent(task) + } + + async compareAndSet(expected: Task, next: Task) { + if (next.status.state === 'failed') { + const current = await this.inner.get(expected.id) + if (current) { + await this.inner.put({ + ...current, + status: { state: 'canceled', timestamp: new Date().toISOString() }, + metadata: { + ...(current.metadata ?? {}), + gatewayPaymentRecovery: { version: 1, id: 'recovery-race' }, + }, + }) + } + } + return this.inner.compareAndSet(expected, next) + } + + async delete(id: string) { + return this.inner.delete(id) + } + } + + const taskStore = new InterleavingFailureStore() + const app = new Hono() + app.route('/v1/agents', createAgentGateway(gatewayConfig(taskStore, { + getSandbox: async () => ({ + async *streamPrompt() { + throw new Error('sandbox failed before a receipt') + }, + }), + }))) + + const response = await post( + app, + agentA.slug, + body('message/send', 'task-stale-failure'), + { 'X-Payment-Signature': paymentHeader('902') }, + ) + expect((await response.json() as { error?: unknown }).error).toBeDefined() + + const stored = await taskStore.get('task-stale-failure') + expect(stored?.status.state).toBe('canceled') + expect(stored?.metadata?.gatewayPaymentRecovery).toEqual({ version: 1, id: 'recovery-race' }) + }) + + it('rejects task access and continuation through a different originating agent', async () => { + const taskStore = new InMemoryTaskStore() + let resolvedOriginAgent = agentA + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + ...gatewayConfig(taskStore), + resolveAgent: async (slug) => slug === agentA.slug + ? resolvedOriginAgent + : slug === agentB.slug + ? agentB + : null, + })) + + const created = await post( + app, + agentA.slug, + body('message/send', 'task-origin-bound'), + { 'X-Payment-Signature': paymentHeader('903') }, + ) + expect((await created.json() as { result?: Task }).result?.status.state).toBe('completed') + const original = await taskStore.get('task-origin-bound') + expect(original?.metadata?.gatewayOrigin).toEqual({ + version: 1, + agentId: agentA.id, + agentSlug: agentA.slug, + }) + + const getFromOtherAgent = await post(app, agentB.slug, body('tasks/get', 'task-origin-bound')) + expect(getFromOtherAgent.status).toBe(403) + expect((await getFromOtherAgent.json() as { error?: { code?: number } }).error?.code) + .toBe(-32008) + + resolvedOriginAgent = { ...agentA, id: 'agent-lifecycle-replaced' } + const getAfterAgentReplacement = await post(app, agentA.slug, body('tasks/get', 'task-origin-bound')) + expect(getAfterAgentReplacement.status).toBe(403) + expect((await getAfterAgentReplacement.json() as { error?: { code?: number } }).error?.code) + .toBe(-32008) + + const cancelFromOtherAgent = await post(app, agentB.slug, body('tasks/cancel', 'task-origin-bound')) + expect(cancelFromOtherAgent.status).toBe(403) + expect((await cancelFromOtherAgent.json() as { error?: { code?: number } }).error?.code) + .toBe(-32008) + expect((await taskStore.get('task-origin-bound'))?.status.state).toBe('completed') + }) + + it('expires a task created before payment authorization can finish', async () => { + class CrashAfterCreateStore implements TaskStore { + private readonly inner = new InMemoryTaskStore() + private crashed = false + + get(id: string) { return this.inner.get(id) } + put(task: Task) { return this.inner.put(task) } + compareAndSet(expected: Task, next: Task) { return this.inner.compareAndSet(expected, next) } + delete(id: string) { return this.inner.delete(id) } + + async createIfAbsent(task: Task) { + const created = await this.inner.createIfAbsent(task) + if (created && !this.crashed) { + this.crashed = true + throw new Error('process crashed after task creation') + } + return created + } + } + + const taskStore = new CrashAfterCreateStore() + const app = new Hono() + app.route('/v1/agents', createAgentGateway(gatewayConfig(taskStore))) + + await post( + app, + agentA.slug, + body('message/send', 'task-created-before-claim'), + { 'X-Payment-Signature': paymentHeader('904') }, + ) + + const created = await taskStore.get('task-created-before-claim') + expect(created?.status.state).toBe('submitted') + expect(created?.metadata?.gatewaySubmission).toBeDefined() + await taskStore.delete('task-created-before-claim') + expect(await taskStore.get('task-created-before-claim')).toBeDefined() + const submission = created?.metadata?.gatewaySubmission as { + lease: { id: string; expiresAt: number } + } + await taskStore.put({ + ...created!, + metadata: { + ...created!.metadata, + gatewaySubmission: { + ...submission, + lease: { ...submission.lease, expiresAt: Date.now() - 1 }, + }, + }, + }) + + const recovered = await post(app, agentA.slug, body('tasks/get', 'task-created-before-claim')) + const recoveredBody = await recovered.json() as { result?: Task } + expect(recoveredBody.result?.status.state).toBe('failed') + expect(recoveredBody.result?.status.state).not.toBe('submitted') + }) +}) + +describe('A2A client disconnect cancellation', () => { + async function makeDisconnectHarness() { + const taskStore = new InMemoryTaskStore() + let started!: () => void + let aborted!: () => void + const sandboxStarted = new Promise((resolve) => { started = resolve }) + const sandboxAborted = new Promise((resolve) => { aborted = resolve }) + const sandbox: SandboxBox = { + async *streamPrompt(_message, opts) { + const signal = opts?.signal + if (signal?.aborted) { + aborted() + return + } + signal?.addEventListener('abort', () => { + aborted() + }, { once: true }) + started() + yield { type: 'message.part.updated', data: { part: { type: 'text' }, delta: 'started' } } + await new Promise((resolve) => { + signal?.addEventListener('abort', resolve, { once: true }) + }) + }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(gatewayConfig(taskStore, { + getSandbox: async () => sandbox, + }))) + return { app, sandboxStarted, sandboxAborted } + } + + it('aborts sandbox work when the response reader disconnects', async () => { + const { app, sandboxStarted, sandboxAborted } = await makeDisconnectHarness() + const response = await post( + app, + agentA.slug, + body('message/stream', 'task-reader-disconnect'), + { 'X-Payment-Signature': paymentHeader('905') }, + ) + const reader = response.body!.getReader() + const firstRead = reader.read() + await sandboxStarted + await reader.cancel() + await sandboxAborted + await firstRead + }) + + it('aborts sandbox work when the incoming request signal disconnects', async () => { + const { app, sandboxStarted, sandboxAborted } = await makeDisconnectHarness() + const requestAbort = new AbortController() + const request = new Request(`http://localhost/v1/agents/${agentA.slug}`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('906'), + }, + body: JSON.stringify(body('message/stream', 'task-request-disconnect')), + signal: requestAbort.signal, + }) + const response = await app.fetch(request) + const reader = response.body!.getReader() + const firstRead = reader.read() + await sandboxStarted + requestAbort.abort() + await sandboxAborted + await reader.cancel() + await firstRead + }) +}) diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index 9f00225..c1817b1 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -1402,6 +1402,9 @@ describe('A2A payment ownership races', () => { expect(requestIds.size).toBe(1) expect(lifecycle.confirmations).toBe(1) expect(lifecycle.releases).toBe(0) - expect((await taskStore.get('task-stripe-usage-recovery'))?.metadata).toBeUndefined() + const recoveredTask = await taskStore.get('task-stripe-usage-recovery') + expect(recoveredTask?.metadata?.gatewayOrigin).toBeDefined() + expect(recoveredTask?.metadata?.gatewayFinalizing).toBeUndefined() + expect(recoveredTask?.metadata?.gatewayPaymentRecovery).toBeUndefined() }) }) From a2af0c58f918c85b6f63e7f35a7c270c585684bb Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Fri, 14 Aug 2026 23:59:52 -0600 Subject: [PATCH 22/32] fix(gateway): preserve payment integration compatibility --- .agent/skill-runs.jsonl | 3 ++ README.md | 4 ++ src/a2a/task-store-sql.ts | 10 ++--- src/index.ts | 2 + src/middleware.ts | 28 ++++++++++++-- src/nonce-store.ts | 33 ++++++++++++---- src/types.ts | 8 ++++ src/verify.ts | 21 ++++++++++- tests/a2a-durability.test.ts | 33 +++++++++++++++- tests/middleware.test.ts | 73 +++++++++++++++++++++++++++++++++--- tests/nonce-store.test.ts | 18 +++++++-- tests/verify.test.ts | 40 ++++++++++++++++++-- 12 files changed, 239 insertions(+), 34 deletions(-) diff --git a/.agent/skill-runs.jsonl b/.agent/skill-runs.jsonl index 29a68a7..47dc524 100644 --- a/.agent/skill-runs.jsonl +++ b/.agent/skill-runs.jsonl @@ -1,2 +1,5 @@ {"skill":"/deep-clean","ts":"2026-08-15T04:14:22Z","project":"agent-gateway-x402-priced-reservation","target":"PR #11 payment lifecycle and recovery: 26 files","operatorPrompt":"","durationMin":null,"verdict":"60 LOC and 2249 bytes of proven duplicate code removed; all gates green","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} {"skill":"/orchestrate","ts":"2026-08-15T04:14:22Z","project":"agent-gateway-x402-priced-reservation","target":"agent-gateway PR #11 deep-clean","operatorPrompt":"","durationMin":null,"verdict":"six read-only Luna Max audits synthesized; local commit b83f1f2; all gates green","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/harden","ts":"2026-08-15T06:00:39Z","project":"agent-gateway-compat-store","target":"Gateway compatibility, MPP auth, SQL task GC","operatorPrompt":"","durationMin":null,"verdict":"PASS","dispatchedTo":"/critical-audit","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/critical-audit","ts":"2026-08-15T06:00:39Z","project":"agent-gateway-compat-store","target":"Gateway compatibility/store commit n=11 files","operatorPrompt":"","durationMin":null,"verdict":"APPROVE","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/verify","ts":"2026-08-15T06:00:39Z","project":"agent-gateway-compat-store","target":"Gateway compatibility/store commit","operatorPrompt":"","durationMin":null,"verdict":"SHIP_IT","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} diff --git a/README.md b/README.md index f92dc77..2d243f1 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,10 @@ The live credential is passed to `confirmPayment` only on the original request. The nonce and recovery stores persist only the SHA-256 digest of `paymentIdentity`. `Payment-Receipt` values must contain visible ASCII only. +`NonceStore` remains source-compatible with 0.7.1 `hasSeen`/`markSeen` stores. +That legacy path is limited to non-owner payment claims; version 2 and MPP charge lifecycles require an atomic `claim` method. +The 0.7.1 `mpp.verifySigner` callback is also supported; the gateway derives a stable identity until the integration moves to `authenticateCredential`. + The same authentication, authorization, rate-limit, filtering, sandbox, settlement, and usage-recording pipeline is used by the OpenAI-compatible and A2A endpoints. Wire protocol handlers only translate their request and response shapes. diff --git a/src/a2a/task-store-sql.ts b/src/a2a/task-store-sql.ts index 96e3cc2..b4ba33f 100644 --- a/src/a2a/task-store-sql.ts +++ b/src/a2a/task-store-sql.ts @@ -138,13 +138,11 @@ export class SqlTaskStore implements TaskStore { if (!row) return undefined const task = JSON.parse(row.payload) as Task if (Date.now() - row.updated_at > this.ttlMs && !hasPendingPaymentRecovery(task)) { - // Lazy GC. If the delete loses a race with another reader, that reader - // observes either the stale-then-deleted task (returning undefined here) - // or, after this delete commits, observes undefined directly — either - // way callers see consistent "expired" semantics. + // Delete only the version that was observed as stale. A refresh can reuse + // the same payload, so payload equality alone does not protect the row. void this.db.exec( - `DELETE FROM ${this.table} WHERE id = ? AND payload = ?`, - [id, row.payload], + `DELETE FROM ${this.table} WHERE id = ? AND payload = ? AND updated_at = ?`, + [id, row.payload, row.updated_at], ) return undefined } diff --git a/src/index.ts b/src/index.ts index 9bf2506..3573f1d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -75,6 +75,8 @@ export { export { MemoryNonceStore, KvNonceStore, + isAtomicNonceStore, + type AtomicNonceStore, type NonceStore, } from './nonce-store' export { diff --git a/src/middleware.ts b/src/middleware.ts index b7def31..819d220 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -13,7 +13,7 @@ import { releasePaymentAfterFailure, settleAndRecord, } from './dispatch' -import { MemoryNonceStore } from './nonce-store' +import { isAtomicNonceStore, MemoryNonceStore } from './nonce-store' import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' import { MemoryPaymentRecoveryStore, @@ -93,24 +93,44 @@ export function createAgentGateway(inputConfig: GatewayConfig) { throw new Error('createAgentGateway: version 1 cannot be combined with version 2 payment operations') } const mppMethod = (config.mpp?.method ?? 'blueprintevm').toLowerCase() + if (config.mpp?.authenticateCredential !== undefined && + typeof config.mpp.authenticateCredential !== 'function') { + throw new Error('createAgentGateway: mpp.authenticateCredential must be a function') + } + if (config.mpp?.verifySigner !== undefined && typeof config.mpp.verifySigner !== 'function') { + throw new Error('createAgentGateway: mpp.verifySigner must be a function') + } + const mppAuthenticator = typeof config.mpp?.authenticateCredential === 'function' + ? config.mpp.authenticateCredential + : typeof config.mpp?.verifySigner === 'function' + ? config.mpp.verifySigner + : undefined if (config.mpp?.charge && config.mpp.charge.protocolVersion !== 1) { throw new Error('createAgentGateway: unsupported MPP charge lifecycle version') } if ( config.mpp?.charge && mppMethod !== 'blueprintevm' && - !config.mpp.authenticateCredential + !mppAuthenticator ) { - throw new Error('createAgentGateway: generic MPP methods require pure credential authentication') + throw new Error('createAgentGateway: generic MPP methods require credential authentication') } if ( config.mpp && mppMethod !== 'blueprintevm' && - config.mpp.authenticateCredential && + mppAuthenticator && !config.mpp.charge ) { throw new Error('createAgentGateway: generic MPP methods require a charge lifecycle') } + if (config.mpp && mppMethod !== 'blueprintevm' && !mppAuthenticator) { + throw new Error('createAgentGateway: generic MPP methods require credential authentication') + } + const needsAtomicNonce = config.x402.paymentProtocolVersion === 2 || + (mppMethod !== 'blueprintevm' && config.mpp?.charge !== undefined) + if (needsAtomicNonce && config.nonceStore && !isAtomicNonceStore(config.nonceStore)) { + throw new Error('createAgentGateway: durable payment ownership requires an atomic nonce store') + } const needsRecovery = config.x402.paymentProtocolVersion === 2 || (mppMethod !== 'blueprintevm' && config.mpp?.charge !== undefined) if (needsRecovery && !config.paymentRecovery) { diff --git a/src/nonce-store.ts b/src/nonce-store.ts index 415595d..352ac5b 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -4,13 +4,19 @@ */ export interface NonceStore { - /** Optional observation. This method never grants ownership. */ - hasSeen?(nonce: string): Promise + /** Check if nonce has been seen. This method never grants ownership. */ + hasSeen(nonce: string): Promise /** * Atomically claim a nonce. An owner id makes a retry by the same payment - * operation idempotent. Calls without an owner id use first-writer-wins - * semantics for legacy payment authorization. + * operation idempotent. This is optional only for the 0.7.1 check-and-mark + * compatibility contract; durable owner claims require this method. */ + claim?(nonce: string, ttlSeconds: number, ownerId?: string): Promise + /** @deprecated Use claim() for atomic ownership in new stores. */ + markSeen?(nonce: string, ttlSeconds: number): Promise +} + +export interface AtomicNonceStore extends NonceStore { claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise } @@ -135,15 +141,26 @@ export class KvNonceStore implements NonceStore { } } -/** Claim through the single atomic nonce-ownership contract. */ +/** Claim through the atomic contract or the explicit 0.7.1 legacy path. */ export async function claimStoredNonce( store: NonceStore, nonce: string, ttlSeconds: number, ownerId?: string, ): Promise { - if (typeof store.claim !== 'function') { - throw new Error('NonceStore.claim is required for payment replay protection') + if (typeof store.claim === 'function') return store.claim(nonce, ttlSeconds, ownerId) + if (ownerId !== undefined) { + throw new Error('NonceStore.claim is required for atomic payment ownership') + } + if (typeof store.markSeen !== 'function') { + throw new Error('NonceStore.markSeen is required for legacy payment replay protection') } - return store.claim(nonce, ttlSeconds, ownerId) + if (await store.hasSeen(nonce)) return false + await store.markSeen(nonce, ttlSeconds) + return true +} + +/** Durable payment paths must use a store with a single atomic claim operation. */ +export function isAtomicNonceStore(store: NonceStore): store is AtomicNonceStore { + return typeof store.claim === 'function' } diff --git a/src/types.ts b/src/types.ts index 182ba9c..136a04e 100644 --- a/src/types.ts +++ b/src/types.ts @@ -144,6 +144,14 @@ export interface MppConfig { payload: Record, context: { method: string; credential: string }, ) => Promise + /** + * @deprecated Use authenticateCredential and return a stable payment identity. + * This 0.7.1 callback remains supported through an explicit compatibility adapter. + */ + verifySigner?: ( + payload: Record, + context: { method: string; credential: string }, + ) => Promise /** Required immediate-charge lifecycle for every non-BlueprinTEVM method. */ charge?: MppChargeLifecycle } diff --git a/src/verify.ts b/src/verify.ts index aaf6634..bfb3152 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -91,6 +91,16 @@ function stableJson(value: unknown): string { return JSON.stringify(value) } +function legacyMppPaymentIdentity( + method: string, + payload: Record, + credential: string, +): string { + const canonicalPayload = stableJson(payload) + if (canonicalPayload !== '{}') return `legacy:${method}:${canonicalPayload}` + return `legacy:${method}:credential:${Buffer.from(credential).toString('base64url')}` +} + /** Pure capability checks shared by discovery and every request protocol. */ export function isApiKeyAuthEnabled( config: Pick, @@ -104,7 +114,8 @@ export function isMppAuthEnabled( ): boolean { const method = (config.mpp?.method ?? 'blueprintevm').toLowerCase() if (!config.mpp) return false - const authenticated = config.mpp.authenticateCredential !== undefined || + const authenticated = typeof config.mpp.authenticateCredential === 'function' || + typeof config.mpp.verifySigner === 'function' || (method === 'blueprintevm' && config.x402.verifySigner !== undefined) || (method === 'blueprintevm' && config.x402.demoMode === true) if (!authenticated) return false @@ -236,6 +247,14 @@ export async function verifyMpp( let authenticated: MppAuthenticatedCredential | null = null if (config.authenticateCredential) { authenticated = await config.authenticateCredential(payload, { method, credential: decoded }) + } else if (config.verifySigner) { + const consumerId = await config.verifySigner(payload, { method, credential: decoded }) + authenticated = typeof consumerId === 'string' && consumerId.length > 0 + ? { + consumerId, + paymentIdentity: legacyMppPaymentIdentity(method, payload, decoded), + } + : null } else if (method === 'blueprintevm' && x402Config.verifySigner && payload.commitment) { const verified = await x402Config.verifySigner(payload, { protocolVersion: x402Config.paymentProtocolVersion ?? (x402Config.paymentOperations ? 2 : 1), diff --git a/tests/a2a-durability.test.ts b/tests/a2a-durability.test.ts index 8404976..5f52df1 100644 --- a/tests/a2a-durability.test.ts +++ b/tests/a2a-durability.test.ts @@ -32,7 +32,9 @@ interface FakePushRow { authentication: string | null } -function makeTaskAdapter(): SqlAdapter & { rows: Map } { +function makeTaskAdapter( + beforeDelete?: () => Promise, +): SqlAdapter & { rows: Map } { const rows = new Map() return { rows, @@ -69,10 +71,14 @@ function makeTaskAdapter(): SqlAdapter & { rows: Map } { return { rowsAffected: 1 } } if (s.startsWith('DELETE')) { - const [id, expectedPayload] = params as [string, string | undefined] + await beforeDelete?.() + const [id, expectedPayload, expectedUpdatedAt] = params as [string, string | undefined, number | undefined] if (expectedPayload !== undefined && rows.get(id)?.payload !== expectedPayload) { return { rowsAffected: 0 } } + if (expectedUpdatedAt !== undefined && rows.get(id)?.updated_at !== expectedUpdatedAt) { + return { rowsAffected: 0 } + } const had = rows.delete(id) return { rowsAffected: had ? 1 : 0 } } @@ -158,6 +164,29 @@ describe('SqlTaskStore', () => { expect(await store.get('t1')).toBeUndefined() }) + it('does not let lazy GC delete a row refreshed with the same payload', async () => { + let deleteStarted!: () => void + const deleteStartedPromise = new Promise((resolve) => { deleteStarted = resolve }) + let releaseDelete!: () => void + const deleteGate = new Promise((resolve) => { releaseDelete = resolve }) + const db = makeTaskAdapter(async () => { + deleteStarted() + await deleteGate + }) + const store = new SqlTaskStore(db, { ttlMs: 1 }) + const task = makeTask('t1') + await store.put(task) + db.rows.get('t1')!.updated_at = Date.now() - 100 + + const staleRead = store.get('t1') + await deleteStartedPromise + await store.put(task) + releaseDelete() + expect(await staleRead).toBeUndefined() + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(db.rows.has('t1')).toBe(true) + }) + it.each([ 'gatewayFinalizing', 'gatewayPaymentRelease', diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index c875663..3a04b77 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -19,6 +19,7 @@ import type { import { MemoryNonceStore, type NonceStore } from '../src/nonce-store' import { MemoryRateLimitStore } from '../src/rate-limit' import { MemoryPaymentOperations } from '../src/payment-operations' +import { MemoryPaymentRecoveryStore } from '../src/payment-recovery' import type { MppChargeLifecycle } from '../src/mpp-payment' const operatorAddress = '0x1111111111111111111111111111111111111111' @@ -257,8 +258,8 @@ describe('GET /:slug/chat/completions (discovery)', () => { expect(response.status).toBe(401) }) - it('does not advertise an MPP method without a compatible verifier', async () => { - const { app } = buildHarness({ + it('rejects an MPP method without a compatible verifier at gateway construction', () => { + expect(() => buildHarness({ mpp: { realm: 'agents.tangle.tools', method: 'stripe' }, x402: { operatorAddress, @@ -266,10 +267,69 @@ describe('GET /:slug/chat/completions (discovery)', () => { demoMode: false, verifySigner: async () => true, }, - }) - const discovery = await app.request('/v1/agents/test-agent/chat/completions') - const body = await discovery.json() as { payment_methods: Array<{ type: string }> } - expect(body.payment_methods.map((method) => method.type)).not.toContain('mpp') + })).toThrow('credential authentication') + }) + + it('rejects malformed MPP callback values at gateway construction', () => { + expect(() => buildHarness({ + mpp: { + realm: 'agents.tangle.tools', + method: 'stripe', + authenticateCredential: 'not-a-function' as unknown as NonNullable['authenticateCredential'], + }, + })).toThrow('must be a function') + }) + + it('rejects an old generic MPP verifier without a charge lifecycle instead of silently disabling MPP', () => { + expect(() => buildHarness({ + x402: { + operatorAddress, + chainId: 3799, + demoMode: false, + verifySigner: async () => true, + }, + mpp: { + realm: 'agents.tangle.tools', + method: 'stripe', + verifySigner: async () => 'mpp:legacy', + }, + })).toThrow('charge lifecycle') + }) + + it('accepts the old generic MPP verifier when its charge lifecycle is explicit', () => { + expect(() => buildHarness({ + x402: { + operatorAddress, + chainId: 3799, + demoMode: false, + verifySigner: async () => true, + }, + mpp: { + realm: 'agents.tangle.tools', + method: 'stripe', + verifySigner: async () => 'mpp:legacy', + charge: mppChargeLifecycle(), + }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + })).not.toThrow() + }) + + it('rejects a legacy nonce store for version 2 payment ownership', () => { + expect(() => buildHarness({ + nonceStore: { + hasSeen: async () => false, + markSeen: async () => undefined, + } as unknown as NonceStore, + x402: { + operatorAddress, + chainId: 3799, + demoMode: false, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations(), + }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + })).toThrow('atomic nonce') }) it('does not serve an agent whose resolver marks it disabled', async () => { @@ -586,6 +646,7 @@ describe('POST /:slug/chat/completions — auth paths', () => { it('durably claims an x402-compatible MPP receipt and rejects its replay', async () => { const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) const nonceStore: NonceStore = { + hasSeen: async () => false, claim: async () => true, } const credential = Buffer.from(JSON.stringify({ diff --git a/tests/nonce-store.test.ts b/tests/nonce-store.test.ts index 56963af..0a38aaf 100644 --- a/tests/nonce-store.test.ts +++ b/tests/nonce-store.test.ts @@ -67,13 +67,25 @@ describe('MemoryNonceStore', () => { }) describe('claimStoredNonce', () => { - it('fails closed for a store without an atomic claim method', async () => { + it('keeps the 0.7.1 check-and-mark store contract for legacy claims', async () => { + const seen = new Set() + const legacyStore: NonceStore = { + hasSeen: async (nonce) => seen.has(nonce), + markSeen: async (nonce) => { seen.add(nonce) }, + } + + expect(await claimStoredNonce(legacyStore, 'legacy', 60)).toBe(true) + expect(await claimStoredNonce(legacyStore, 'legacy', 60)).toBe(false) + await expect(claimStoredNonce(legacyStore, 'legacy', 60, 'operation-1')) + .rejects.toThrow('atomic payment ownership') + }) + + it('fails closed for a store without either replay contract', async () => { const legacyStore = { hasSeen: async () => false, - markSeen: async () => undefined, } as unknown as NonceStore - await expect(claimStoredNonce(legacyStore, 'legacy', 60)).rejects.toThrow('payment replay') + await expect(claimStoredNonce(legacyStore, 'legacy', 60)).rejects.toThrow('markSeen') }) }) diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 0231bc5..682c54f 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -98,6 +98,7 @@ describe('verifyX402', () => { it('retains a nonce until its signed expiry, beyond the old one-hour cap', async () => { const ttls: number[] = [] const nonceStore: NonceStore = { + hasSeen: async () => false, claim: async (_nonce, ttlSeconds) => { ttls.push(ttlSeconds) return true @@ -115,13 +116,16 @@ describe('verifyX402', () => { expect(ttls[0]).toBeLessThanOrEqual(7200) }) - it('fails closed for a custom nonce store without an atomic claim', async () => { + it('keeps 0.7.1 custom nonce stores working on the legacy path', async () => { + const seen = new Set() const nonceStore = { - hasSeen: async () => false, - markSeen: async () => undefined, + hasSeen: async (nonce: string) => seen.has(nonce), + markSeen: async (nonce: string) => { seen.add(nonce) }, } as unknown as NonceStore - expect(await verifyX402(buildSpendAuth({ nonce: '101' }), baseConfig, nonceStore)).toBeNull() + const payload = buildSpendAuth({ nonce: '101' }) + expect(await verifyX402(payload, baseConfig, nonceStore)).toBe('0xCommitmentAlice') + expect(await verifyX402(payload, baseConfig, nonceStore)).toBeNull() }) it('isolates nonces per commitment — regression: commitment-less nonce tracking lets Alice replay Bob\'s nonce', async () => { @@ -221,6 +225,34 @@ describe('verifyMpp', () => { expect(seen[0].credential).toContain('commitment') }) + it('adapts the 0.7.1 mpp.verifySigner contract to a stable payment identity', async () => { + const header = buildCredential({ + commitment: '0xAlice', + operator: operatorAddress, + amount: '1000', + nonce: '7', + expiry: String(Math.floor(Date.now() / 1000) + 600), + }).replace('Payment blueprintevm ', 'Payment stripe ') + const legacyConfig = { + ...mppConfig, + method: 'stripe', + verifySigner: async (_payload: Record, context: { method: string; credential: string }) => { + expect(context.method).toBe('stripe') + expect(context.credential).toContain('commitment') + return 'mpp:alice' + }, + } + + await expect(verifyMpp( + header, + legacyConfig, + { ...baseConfig, demoMode: false }, + )).resolves.toMatchObject({ + consumerId: 'mpp:alice', + replayKey: expect.stringMatching(/^mpp:stripe:[0-9a-f]{64}$/), + }) + }) + it('shares the x402 nonce authority with equivalent BlueprinTEVM credentials', async () => { const nonceStore = new MemoryNonceStore() const payload = { From 5e5a4c75850d5dea4c59ac31978c64009f23925a Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 02:26:41 -0600 Subject: [PATCH 23/32] fix(gateway): close execution and payment races --- README.md | 14 +- docs/a2a-long-horizon.md | 15 +- package.json | 2 +- src/a2a/execution-fence.ts | 113 +++++++ src/a2a/handler.ts | 120 +++++-- src/dispatch.ts | 168 ++++++++-- src/index.ts | 7 +- src/middleware.ts | 57 ++-- src/nonce-store.ts | 92 ++++-- src/payment-operations.ts | 27 +- src/payment-recovery-worker.ts | 14 + src/types.ts | 13 +- src/verify.ts | 25 +- tests/a2a-atomicity.test.ts | 13 +- tests/kv-stores.test.ts | 29 +- tests/middleware.test.ts | 2 +- tests/nonce-store.test.ts | 10 +- tests/payment-recovery.test.ts | 5 +- tests/pr11-regressions.test.ts | 565 +++++++++++++++++++++++++++++++++ tests/verify.test.ts | 40 +-- 20 files changed, 1157 insertions(+), 174 deletions(-) create mode 100644 src/a2a/execution-fence.ts create mode 100644 tests/pr11-regressions.test.ts diff --git a/README.md b/README.md index 2d243f1..e2a3f58 100644 --- a/README.md +++ b/README.md @@ -59,14 +59,18 @@ A retained operation settles from its receipt when one exists. If the receipt does not arrive before `receiptTimeoutMs`, recovery settles the original quoted ceiling. The fallback never settles the payer's larger authorization amount. Keep version 1 explicitly configured while old and new gateways coexist; shared nonce storage must reject a version 1 claim owned by a version 2 operation. -Before it calls the verifier, the gateway requires the signed amount to cover filtered input plus the requested output limit. +Before it calls the verifier, the gateway requires the signed amount to cover the complete filtered conversation plus the requested output limit. +The default bound includes system text, message roles, and JSON framing. +Set `inputTokenBound` when the provider adds harness, tool, workspace, or other hidden context. The gateway rejects `max_tokens` above `maxOutputTokens` and stops the sandbox stream at the accepted limit. An unpaid request receives `required_amount`, `currency_decimals`, and `max_output_tokens` in the 402 response. Sandbox adapters should emit a complete `sandbox.usage` receipt. Requests with a version 2 operation or generic MPP charge reject missing receipts. API-key requests keep the legacy visible-token estimate path. recordUsage must atomically upsert by event.requestId; recovery may retry an event after its acknowledgement is lost. -Older custom A2A task stores remain source-compatible through a process-safe fallback. +The default gateway still exposes A2A with an in-memory task store. +Older custom A2A task stores remain source-compatible at the type boundary. +The OpenAI surface stays available when such a store is configured, while A2A returns `503` until its owner supplies atomic methods. Use an atomic task store for multi-worker production deployments. MPP is method-specific. @@ -90,7 +94,11 @@ The nonce and recovery stores persist only the SHA-256 digest of `paymentIdentit `Payment-Receipt` values must contain visible ASCII only. `NonceStore` remains source-compatible with 0.7.1 `hasSeen`/`markSeen` stores. -That legacy path is limited to non-owner payment claims; version 2 and MPP charge lifecycles require an atomic `claim` method. +Payment requests now require its atomic `claim` method, including version 1. +This is a deliberate safety boundary: a check followed by a write can accept two concurrent payments. +`KvNonceStore` with plain Cloudflare KV is not atomic and is rejected by `createAgentGateway`. +Provide `KvNonceStore` an `atomicClaim` callback backed by D1, a Durable Object, or another linearizable store. +Payment paths fail closed unless the store also provides one atomic `claim` method. The 0.7.1 `mpp.verifySigner` callback is also supported; the gateway derives a stable identity until the integration moves to `authenticateCredential`. The same authentication, authorization, rate-limit, filtering, sandbox, settlement, and usage-recording pipeline is used by the OpenAI-compatible and A2A endpoints. diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 99094dd..babb2f8 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -9,17 +9,22 @@ The A2A protocol works well for short request-response calls out of the box. For | Resubscribe | `tasks/resubscribe` | A client that lost its SSE stream can re-attach to find out where the task ended up | | Input-required + multi-turn | `input-required` state + follow-up `message/send` with the same `taskId` | The agent can pause and ask the user a question without ending the task | -All four are gated on configuration — they cost nothing for agents that don't need them, and the agent card honestly reflects what each gateway will actually do. +The A2A surface is available with an in-memory task store by default. +Durable storage and push notifications are enabled through `GatewayConfig.a2a`. ## Durable tasks (SqlTaskStore) By default `GatewayConfig.a2a.taskStore` is in-memory: fast, zero-config, fine for tests and single-machine deployments. Production deployments swap in `SqlTaskStore` against any SQL store — D1, postgres, sqlite, libSQL, Turso — via a `SqlAdapter` shim. Custom production task stores must implement both atomic methods, `createIfAbsent` and `compareAndSet`. -The gateway rejects a production task store that lacks either method before it serves A2A requests. -Explicit `x402.demoMode` keeps a read-then-write fallback for local tests only. +The gateway keeps the OpenAI surface available when an older store lacks either method. +It returns `503` for A2A until the store is upgraded, rather than using a cross-worker unsafe fallback. Use `SqlTaskStore` or another atomic adapter for multi-worker production deployments. +Payment nonce stores also require an atomic claim operation on every payment protocol version. +Plain Cloudflare KV cannot provide that operation. +Use D1, a Durable Object, or a custom atomic adapter for the claim callback passed to `KvNonceStore`. + Before a payment provider can mutate state, the gateway stores the recovery identity in task metadata. The record contains the payment operation, usage receipt, output artifact, and a five-minute lease. The task also points to the shared payment recovery outbox. @@ -300,6 +305,10 @@ The gateway then: 4. Does NOT fire push notifications — `input-required` is non-terminal. 5. Returns the task envelope (for `message/send`) or emits a final `status-update` (for `message/stream`). +The gateway records a short-lived execution fence after sandbox acquisition and before the provider call. +Another worker can cancel before that fence is recorded. +After the fence is recorded, a remote cancellation returns `TASK_NOT_CANCELABLE` because it cannot abort the live worker safely. + Sandboxes that never emit input-required see identical behavior to before. ### Client continuation diff --git a/package.json b/package.json index 4bb2b9d..8e53529 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@tangle-network/agent-gateway", - "version": "0.7.2", + "version": "0.8.0", "packageManager": "pnpm@10.28.0", "repository": { "type": "git", diff --git a/src/a2a/execution-fence.ts b/src/a2a/execution-fence.ts new file mode 100644 index 0000000..d12c3fe --- /dev/null +++ b/src/a2a/execution-fence.ts @@ -0,0 +1,113 @@ +import type { Task } from './types' +import type { TaskStore } from './task-store' + +/** Durable marker that prevents cancellation from racing sandbox start. */ +export const TASK_EXECUTION_METADATA_KEY = 'gatewayExecution' + +const TASK_EXECUTION_VERSION = 1 as const +const TASK_EXECUTION_LEASE_MS = 5 * 60 * 1000 + +interface TaskExecutionMarker { + version: typeof TASK_EXECUTION_VERSION + requestId: string + lease: { id: string; expiresAt: number } +} + +export class TaskExecutionCanceledError extends Error { + constructor(taskId: string) { + super(`A2A task '${taskId}' was canceled before sandbox execution`) + this.name = 'TaskExecutionCanceledError' + } +} + +/** Claim the right to start one task after its sandbox has been acquired. */ +export async function claimTaskExecution( + store: TaskStore, + task: Task, + requestId: string, + now = Date.now(), +): Promise { + for (let attempt = 0; attempt < 8; attempt += 1) { + const current = await store.get(task.id) + if (!current || current.status.state !== 'working') { + throw new TaskExecutionCanceledError(task.id) + } + const existing = readTaskExecution(current) + if (existing && existing.lease.expiresAt > now) { + if (existing.requestId === requestId) return current + throw new Error(`A2A task '${task.id}' is already executing`) + } + const next = withTaskExecution(current, requestId, now) + if (store.compareAndSet && await store.compareAndSet(current, next)) return next + } + throw new Error(`A2A task '${task.id}' changed too many times before sandbox execution`) +} + +/** Renew both the task execution fence and its cancellation protection. */ +export async function renewTaskExecution( + store: TaskStore, + taskId: string, + requestId: string, + now = Date.now(), +): Promise { + for (let attempt = 0; attempt < 8; attempt += 1) { + const current = await store.get(taskId) + const marker = current ? readTaskExecution(current) : undefined + if (!current || current.status.state !== 'working' || marker?.requestId !== requestId) { + throw new TaskExecutionCanceledError(taskId) + } + const next = withTaskExecution(current, requestId, now) + if (store.compareAndSet && await store.compareAndSet(current, next)) return next + } + throw new Error(`A2A task '${taskId}' changed too many times while execution was active`) +} + +/** Cancellation is rejected only while a live execution fence is held. */ +export function hasActiveTaskExecution(task: Task, now = Date.now()): boolean { + const marker = readTaskExecution(task) + return marker !== undefined && marker.lease.expiresAt > now +} + +/** Remove the marker when the task reaches a terminal or paused state. */ +export function clearTaskExecution(task: Task): Task { + if (!task.metadata || !(TASK_EXECUTION_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[TASK_EXECUTION_METADATA_KEY] + return Object.keys(metadata).length > 0 + ? { ...task, metadata } + : (() => { + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata + })() +} + +function withTaskExecution(task: Task, requestId: string, now: number): Task { + return { + ...task, + metadata: { + ...(task.metadata ?? {}), + [TASK_EXECUTION_METADATA_KEY]: { + version: TASK_EXECUTION_VERSION, + requestId, + lease: { id: requestId, expiresAt: now + TASK_EXECUTION_LEASE_MS }, + } satisfies TaskExecutionMarker, + }, + } +} + +function readTaskExecution(task: Task): TaskExecutionMarker | undefined { + const raw = task.metadata?.[TASK_EXECUTION_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const marker = raw as Partial + if ( + marker.version !== TASK_EXECUTION_VERSION || + typeof marker.requestId !== 'string' || + marker.requestId.length === 0 || + !marker.lease || + typeof marker.lease.id !== 'string' || + marker.lease.id.length === 0 || + typeof marker.lease.expiresAt !== 'number' || + !Number.isFinite(marker.lease.expiresAt) + ) return undefined + return marker as TaskExecutionMarker +} diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index cd69b68..17f9025 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -18,6 +18,8 @@ import { type GatewayState, authenticateAndGuard, beginPaymentExecution, + markPaymentExecutionStarted, + renewPaymentExecution, claimPayment, dispatchSandboxStreamRich, releasePayment, @@ -32,12 +34,19 @@ import { } from '../payment-recovery' import { recoverPayment as recoverDurablePayment } from '../payment-recovery-worker' import type { + ChatMessage, GatewayConfig, PaymentMethod, SandboxExecutionBudget, SandboxUsageReceipt, } from '../types' import { buildAgentCard } from './agent-card' +import { + claimTaskExecution, + clearTaskExecution, + hasActiveTaskExecution, + renewTaskExecution, +} from './execution-fence' import { fail, ok, parseEnvelope } from './jsonrpc' import { deliverPushNotifications, @@ -168,6 +177,11 @@ class CancelRegistry { return this.finalizing.has(taskId) } + hasController(taskId: string): boolean { + const controller = this.controllers.get(taskId) + return controller !== undefined && !controller.signal.aborted + } + cancel(taskId: string): boolean { if (this.finalizing.has(taskId)) return false const c = this.controllers.get(taskId) @@ -283,7 +297,7 @@ async function executeMessageSend( signal: AbortSignal, ): Promise { if (isTerminal(task.status.state)) return c.json(ok(req.id, task)) - const workingTask: Task = task.status.state === 'working' + let workingTask: Task = task.status.state === 'working' ? task : { ...task, status: { state: 'working', timestamp: nowIso() } } if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { @@ -305,7 +319,6 @@ async function executeMessageSend( let inputRequiredSeen = false let finalizationLeaseId: string | undefined try { - await beginPaymentExecution(authz, deps.config) for await (const event of dispatchSandboxStreamRich( authz.agent, authz.userMessage, @@ -314,10 +327,19 @@ async function executeMessageSend( signal, task.id, authz.maxOutputTokens, - undefined, + async () => { + workingTask = await claimTaskExecution(deps.taskStore, workingTask, authz.requestId) + await beginPaymentExecution(authz, deps.config) + }, authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, - () => { + async () => { workObserved = true + await markPaymentExecutionStarted(authz, deps.config) + }, + authz.executionBudget.maxInputTokens, + async () => { + workingTask = await renewTaskExecution(deps.taskStore, task.id, authz.requestId) + await renewPaymentExecution(authz, deps.config) }, )) { if (event.kind === 'text') { @@ -499,7 +521,7 @@ async function handleMessageStream( cancels.clear(task.id) return c.json(ok(req.id, task)) } - const workingTask: Task = task.status.state === 'working' + let workingTask: Task = task.status.state === 'working' ? task : { ...task, status: workingStatus.status } if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { @@ -515,26 +537,6 @@ async function handleMessageStream( if (current.status.state === 'canceled') return c.json(ok(req.id, current)) return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) } - try { - await beginPaymentExecution(authz, deps.config) - } catch (error) { - detachRequestAbort() - cancels.clear(task.id) - await releaseTaskPayment( - authz, - workingTask, - deps, - error instanceof Error ? error.message : String(error), - false, - ) - return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment execution authorization failed')) - } - if (controller.signal.aborted) { - detachRequestAbort() - cancels.clear(task.id) - const canceled = await completeCanceledTask(authz, workingTask, '', undefined, false, deps) - return c.json(ok(req.id, canceled)) - } let responseText = '' let usage: SandboxUsageReceipt | undefined let workObserved = false @@ -566,10 +568,19 @@ async function handleMessageStream( controller.signal, task.id, authz.maxOutputTokens, - undefined, + async () => { + workingTask = await claimTaskExecution(deps.taskStore, workingTask, authz.requestId) + await beginPaymentExecution(authz, deps.config) + }, authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, - () => { + async () => { workObserved = true + await markPaymentExecutionStarted(authz, deps.config) + }, + authz.executionBudget.maxInputTokens, + async () => { + workingTask = await renewTaskExecution(deps.taskStore, task.id, authz.requestId) + await renewPaymentExecution(authz, deps.config) }, )) { if (event.kind === 'text') { @@ -900,7 +911,11 @@ async function handleTasksCancel( ) } - if (isTaskFinalizing(task) || cancels.isFinalizing(task.id)) { + if ( + isTaskFinalizing(task) || + cancels.isFinalizing(task.id) || + (hasActiveTaskExecution(task) && !cancels.hasController(task.id)) + ) { return c.json( fail( req.id, @@ -1132,15 +1147,34 @@ async function guardMessageRequest( return c.json(fail(req.id, extracted.error.code, extracted.error.message)) } + let billingMessages: ChatMessage[] = [{ role: 'user', content: extracted.text }] + if (typeof params.message.taskId === 'string') { + const storedForQuote = await deps.taskStore.get(params.message.taskId) + if (storedForQuote) { + const accessError = await authorizeTaskAccess(c, req, storedForQuote, deps) + if (accessError) return accessError + const quotedTask = await recoverTaskIfNeeded(storedForQuote, deps, slug) + if (quotedTask.status.state === 'input-required') { + billingMessages = [ + ...taskHistoryAsChatMessages(quotedTask), + { role: 'user', content: extracted.text }, + ] + } + } + } + const guard = await authenticateAndGuard( c, slug, - [{ role: 'user', content: extracted.text }], + billingMessages, deps.config, deps.state, ) if (guard instanceof Response) return guard const authz = guard + // The provider session receives only the new turn. The quote above covers + // every retained message and the configured hidden provider context. + authz.userMessage = extracted.text // Multi-turn continuation: if the caller addressed an existing task that is // currently in `input-required`, append the new message and reserve it as @@ -1496,6 +1530,7 @@ async function completeCanceledTask( const FINALIZING_METADATA_KEY = 'gatewayFinalizing' const PAYMENT_RELEASE_METADATA_KEY = 'gatewayPaymentRelease' const PAYMENT_RECOVERY_METADATA_KEY = 'gatewayPaymentRecovery' +const EXECUTION_RECOVERY_METADATA_KEY = 'gatewayExecutionRecovery' const FINALIZATION_LEASE_MS = 5 * 60 * 1000 const PAYMENT_RELEASE_LEASE_MS = 5 * 60 * 1000 @@ -1569,6 +1604,17 @@ function bindRequestAbort(requestSignal: AbortSignal, controller: AbortControlle return () => requestSignal.removeEventListener('abort', abort) } +function taskHistoryAsChatMessages(task: Task): ChatMessage[] { + return (task.history ?? []).flatMap((message) => { + const extracted = extractTextFromMessage(message) + if ('error' in extracted) return [] + return [{ + role: message.role === 'agent' ? 'assistant' : 'user', + content: extracted.text, + }] + }) +} + function withTaskOrigin( metadata: Record | undefined, agent: { id: string; slug: string }, @@ -2258,6 +2304,19 @@ async function clearReconciledPaymentRecoveryMarker( const record = await deps.config.paymentRecovery.store.get(marker.id) if (record?.state !== 'reconciled') return task const cleared = clearPaymentRecoveryMarker(task) + if (cleared.status.state === 'working') { + const failed: Task = { + ...withStatus(cleared, 'failed'), + metadata: { + ...(cleared.metadata ?? {}), + [EXECUTION_RECOVERY_METADATA_KEY]: { + error: 'payment recovery completed without a task result', + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) return failed + return await deps.taskStore.get(task.id) ?? failed + } if (await compareAndSetTask(deps.taskStore, task, cleared)) return cleared return await deps.taskStore.get(task.id) ?? cleared } @@ -2378,11 +2437,12 @@ function withStatus( message?: Message, artifacts?: Task['artifacts'], ): Task { - return { + const next: Task = { ...task, status: { state, timestamp: nowIso(), ...(message ? { message } : {}) }, ...(artifacts !== undefined ? { artifacts } : {}), } + return isTerminal(state) || state === 'input-required' ? clearTaskExecution(next) : next } /** diff --git a/src/dispatch.ts b/src/dispatch.ts index ffef0a6..767445c 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -19,7 +19,11 @@ import { import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' import { type RateLimitStore, checkRateLimit } from './rate-limit' import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' -import { paymentNonceKey, type PaymentOperation } from './payment-operations' +import { + paymentNonceKey, + type PaymentOperation, + type PaymentOperationRecoveryResult, +} from './payment-operations' import { PAYMENT_RECOVERY_VERSION, PaymentRecoveryFenceError, @@ -47,7 +51,7 @@ import { isMppAuthEnabled, mppPaymentPayload, mppPaymentCredential, - verifyMpp, + verifyMppCredential, verifyX402, } from './verify' @@ -228,7 +232,38 @@ export async function authenticateAndGuard( } let requiredPaymentAmount: bigint - const maxInputTokens = maximumBillableInputTokens(agent, userMessage) + const messageInputBound = maximumBillableInputTokens(agent, filtered) + let maxInputTokens = messageInputBound + if (config.inputTokenBound) { + let configuredBound: number + try { + configuredBound = config.inputTokenBound({ agent, messages: filtered }) + } catch { + return c.json( + { + error: { + message: 'Agent input token bound is unavailable', + type: 'server_error', + code: 'input_token_bound_unavailable', + }, + }, + 503, + ) + } + if (!Number.isSafeInteger(configuredBound) || configuredBound < messageInputBound) { + return c.json( + { + error: { + message: 'Agent input token bound is invalid', + type: 'server_error', + code: 'invalid_input_token_bound', + }, + }, + 503, + ) + } + maxInputTokens = configuredBound + } const maxReasoningTokens = state.maxReasoningTokens const maxToolTokens = state.maxToolTokens const maxToolCalls = state.maxToolCalls @@ -312,7 +347,7 @@ export async function authenticateAndGuard( consumerId = signer paymentMethod = 'x402' } else if (isMppAuthEnabled(config) && authHeader.toLowerCase().startsWith('payment ')) { - const authenticated = await verifyMpp( + const authenticated = await verifyMppCredential( authHeader, config.mpp!, config.x402, @@ -898,12 +933,17 @@ export async function beginPaymentExecution( authz: AuthorizedRequest, config: GatewayConfig, ): Promise { - if (!authz.paymentRecoveryId) { - if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { - authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) - } - return + if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { + authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) } +} + +/** Persist the sandbox handoff immediately before the adapter call. */ +export async function markPaymentExecutionStarted( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + if (!authz.paymentRecoveryId) return const recovery = config.paymentRecovery if (!recovery) throw new Error('durable payment recovery is not configured') const fenceId = requirePaymentRecoveryFence(authz) @@ -918,9 +958,26 @@ export async function beginPaymentExecution( lease: { id: fenceId, expiresAt: fallbackAt }, nextAttemptAt: fallbackAt, }), now) - if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { - authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) - } +} + +/** Renew the live execution lease while a provider stream is still open. */ +export async function renewPaymentExecution( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + if (!authz.paymentRecoveryId) return + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const fenceId = requirePaymentRecoveryFence(authz) + const now = Date.now() + const fallbackAt = now + recoveryTiming(recovery).receiptTimeoutMs + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + state: 'executing', + fallbackAt, + lease: { id: fenceId, expiresAt: fallbackAt }, + nextAttemptAt: fallbackAt, + }), now) } /** @@ -1024,7 +1081,7 @@ async function markRecoveryReconciled( export async function reclaimPayment( operationId: string, config: GatewayConfig, -): Promise { +): Promise { if (!config.x402.paymentOperations) throw new Error('durable payment operations are not configured') return config.x402.paymentOperations.reclaimPayment(operationId) } @@ -1131,7 +1188,9 @@ export async function* dispatchSandboxStreamRich( maxOutputTokens?: number, onExecutionStart?: () => Promise, requiresReceipt = config.x402.paymentOperations !== undefined, - onSandboxStart?: () => void, + onSandboxStart?: () => void | Promise, + maxInputTokens?: number, + onExecutionHeartbeat?: () => Promise, ): AsyncIterable { if (signal?.aborted) return const box = await config.getSandbox(agent) @@ -1153,34 +1212,66 @@ export async function* dispatchSandboxStreamRich( let observedToolTokens = 0 let observedToolCalls = 0 let legacyOutputText = '' + const executionController = new AbortController() + const forwardAbort = () => executionController.abort() + if (signal?.aborted) return + signal?.addEventListener('abort', forwardAbort, { once: true }) const executionBudget: SandboxExecutionBudget = { - maxInputTokens: maximumBillableInputTokens(agent, userMessage), + 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 ?? - (maximumBillableInputTokens(agent, userMessage) + outputLimit + + maxProviderCostUsd: config.executionBudget?.maxProviderCostUsd ?? ( + (maxInputTokens ?? maximumBillableInputTokens(agent, userMessage)) + outputLimit + (config.executionBudget?.maxReasoningTokens ?? outputLimit) + - (config.executionBudget?.maxToolTokens ?? outputLimit)) * agent.pricePerTokenUsd, + (config.executionBudget?.maxToolTokens ?? outputLimit) + ) * agent.pricePerTokenUsd, } - if (signal?.aborted) return + if (executionController.signal.aborted) return await onExecutionStart?.() - if (signal?.aborted) return - // A stream adapter can start paid work synchronously during this call. - onSandboxStart?.() - const promptStream = box.streamPrompt(userMessage, { - sessionId: sessionId ?? `consumer:${consumerId}`, - systemPrompt: agent.systemPrompt, - maxOutputTokens: outputLimit, - executionBudget, - signal, - }) - const iterator = promptStream[Symbol.asyncIterator]() + if (executionController.signal.aborted) return + let heartbeatError: unknown + let heartbeatInFlight: Promise | undefined + let heartbeatTimer: ReturnType | undefined + let iterator: AsyncIterator | undefined try { + // 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, + }) + iterator = promptStream[Symbol.asyncIterator]() + const heartbeatMs = onExecutionHeartbeat + ? Math.max(100, Math.min( + Math.floor((config.paymentRecovery?.receiptTimeoutMs ?? 5 * 60_000) / 3), + 5_000, + )) + : 0 + if (onExecutionHeartbeat) { + heartbeatTimer = setInterval(() => { + if (heartbeatInFlight || heartbeatError !== undefined) return + heartbeatInFlight = onExecutionHeartbeat() + .catch((error: unknown) => { + heartbeatError = error + executionController.abort() + }) + .finally(() => { + heartbeatInFlight = undefined + }) + }, heartbeatMs) + } while (true) { - const next = await readSandboxEvent(iterator, signal) - if (next === ABORTED_SANDBOX_READ) return + const next = await readSandboxEvent(iterator, executionController.signal) + if (next === ABORTED_SANDBOX_READ) { + if (heartbeatError !== undefined) throw heartbeatError + return + } if (next.done) break const event = next.value if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) @@ -1249,7 +1340,9 @@ export async function* dispatchSandboxStreamRich( ) yield { kind: 'usage', usage } } finally { - await closeSandboxIterator(iterator) + if (heartbeatTimer !== undefined) clearInterval(heartbeatTimer) + signal?.removeEventListener('abort', forwardAbort) + if (iterator) await closeSandboxIterator(iterator) } } @@ -1650,7 +1743,12 @@ export function estimateBillableInputTokens(agent: AgentMeta, userMessage: strin } /** A tokenizer cannot emit more tokens than the UTF-8 bytes it consumes. */ -export function maximumBillableInputTokens(agent: AgentMeta, userMessage: string): number { +export function maximumBillableInputTokens(agent: AgentMeta, userMessage: string): number +export function maximumBillableInputTokens(agent: AgentMeta, messages: readonly ChatMessage[]): number +export function maximumBillableInputTokens(agent: AgentMeta, userMessageOrMessages: string | readonly ChatMessage[]): number { const encoder = new TextEncoder() - return encoder.encode(userMessage).byteLength + encoder.encode(agent.systemPrompt ?? '').byteLength + const prompt = typeof userMessageOrMessages === 'string' + ? userMessageOrMessages + : JSON.stringify(userMessageOrMessages) + return encoder.encode(prompt).byteLength + encoder.encode(agent.systemPrompt ?? '').byteLength } diff --git a/src/index.ts b/src/index.ts index 3573f1d..c412991 100644 --- a/src/index.ts +++ b/src/index.ts @@ -10,6 +10,7 @@ export { export { verifyX402, verifyMpp, + verifyMppCredential, defaultVerifyApiKey, isApiKeyAuthEnabled, isMppAuthEnabled, @@ -46,6 +47,8 @@ export { type MemoryPaymentOperationsOptions, type PaymentAuthorizationContext, type PaymentOperation, + type PaymentOperationNotFound, + type PaymentOperationRecoveryResult, type PaymentOperationState, type PaymentOperations, type PaymentSettlementInput, @@ -76,7 +79,9 @@ export { MemoryNonceStore, KvNonceStore, isAtomicNonceStore, + type AtomicKvNonceClaim, type AtomicNonceStore, + type KvNonceStoreOptions, type NonceStore, } from './nonce-store' export { @@ -114,7 +119,7 @@ export type { // --- A2A protocol surface (Google Agent-to-Agent) --- // Types + task-store adapter. Handlers are wired automatically by -// createAgentGateway when `GatewayConfig.a2a` (or its default) is honored; +// createAgentGateway with an in-memory store by default; // consumers only import these to BYO a durable TaskStore (D1, postgres, DO) // or to declare richer AgentMeta.skills for the Agent Card. export { InMemoryTaskStore, type TaskStore } from './a2a/task-store' diff --git a/src/middleware.ts b/src/middleware.ts index 819d220..c87c1c0 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -7,6 +7,8 @@ import { type GatewayState, authenticateAndGuard, beginPaymentExecution, + markPaymentExecutionStarted, + renewPaymentExecution, claimPayment, dispatchSandboxStreamRich, releasePayment, @@ -126,9 +128,7 @@ export function createAgentGateway(inputConfig: GatewayConfig) { if (config.mpp && mppMethod !== 'blueprintevm' && !mppAuthenticator) { throw new Error('createAgentGateway: generic MPP methods require credential authentication') } - const needsAtomicNonce = config.x402.paymentProtocolVersion === 2 || - (mppMethod !== 'blueprintevm' && config.mpp?.charge !== undefined) - if (needsAtomicNonce && config.nonceStore && !isAtomicNonceStore(config.nonceStore)) { + if (config.nonceStore && !isAtomicNonceStore(config.nonceStore)) { throw new Error('createAgentGateway: durable payment ownership requires an atomic nonce store') } const needsRecovery = config.x402.paymentProtocolVersion === 2 || @@ -296,29 +296,6 @@ export function createAgentGateway(inputConfig: GatewayConfig) { ) } - try { - await beginPaymentExecution(authz, config) - } catch { - try { - await releasePayment(authz, config, 'payment execution fence was lost') - } catch (releaseError) { - console.error( - `[agent-gateway] payment release failed for ${authz.requestId}:`, - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) - } - return c.json( - { - error: { - message: 'Payment authorization failed', - type: 'payment_required', - code: 'payment_execution_fence_lost', - }, - }, - { status: 402, headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': authz.requestId } }, - ) - } - return streamChatCompletions(c, authz, config, obs) }) @@ -329,9 +306,24 @@ export function createAgentGateway(inputConfig: GatewayConfig) { // regardless of which protocol the caller used. const taskStore = config.a2a?.taskStore ?? new InMemoryTaskStore() const pushStore = config.a2a?.pushStore - const a2a = createA2AHandlers({ config, state, taskStore, pushStore }) - gw.get('/:slug/.well-known/agent.json', a2a.handleAgentCard) - gw.post('/:slug', a2a.handleJsonRpc) + try { + const a2a = createA2AHandlers({ config, state, taskStore, pushStore }) + gw.get('/:slug/.well-known/agent.json', a2a.handleAgentCard) + gw.post('/:slug', a2a.handleJsonRpc) + } catch (error) { + // An older custom store must not take down the OpenAI surface. Keep the + // A2A surface unavailable until its owner supplies atomic methods. + console.error( + '[agent-gateway] A2A is unavailable until its task store is upgraded:', + error instanceof Error ? error.message : String(error), + ) + const unavailable = (c: import('hono').Context) => c.json( + { error: 'A2A task persistence is not configured for concurrent workers' }, + 503, + ) + gw.get('/:slug/.well-known/agent.json', unavailable) + gw.post('/:slug', unavailable) + } return gw } @@ -405,11 +397,14 @@ function streamChatCompletions( abortController.signal, undefined, maxOutputTokens, - undefined, + () => beginPaymentExecution(authz, config), authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, - () => { + async () => { if (authz.paymentRecoveryId) workObserved = true + await markPaymentExecutionStarted(authz, config) }, + authz.executionBudget.maxInputTokens, + () => renewPaymentExecution(authz, config), )) { if (event.kind === 'text') { sendChunk(event.delta) diff --git a/src/nonce-store.ts b/src/nonce-store.ts index 352ac5b..2bf7a8f 100644 --- a/src/nonce-store.ts +++ b/src/nonce-store.ts @@ -66,6 +66,11 @@ export class MemoryNonceStore implements NonceStore { return true } + async markSeen(nonce: string, ttlSeconds: number): Promise { + this.evictExpired() + this.seen.set(nonce, { expiresAt: Date.now() + ttlSeconds * 1000 }) + } + private evictExpired() { const now = Date.now() // Evict at most every 60 seconds to avoid O(n) on every request @@ -89,51 +94,88 @@ export class MemoryNonceStore implements NonceStore { export interface KVNamespace { get(key: string, options?: { type?: 'text' | 'json' }): Promise put(key: string, value: string, options?: { expirationTtl?: number }): Promise - /** Linearizable create-if-absent extension. Cloudflare KV does not provide it. */ + /** Optional linearizable create-if-absent extension. Cloudflare KV does not provide it. */ putIfAbsent?(key: string, value: string, options?: { expirationTtl?: number }): Promise delete(key: string): Promise } +/** Atomic claim supplied by D1, a Durable Object, or another linearizable store. */ +export type AtomicKvNonceClaim = ( + key: string, + ttlSeconds: number, + ownerId?: string, +) => Promise + +export interface KvNonceStoreOptions { + /** + * Claim the fully namespaced key atomically. + * The callback must make same-owner retries idempotent. + */ + atomicClaim?: AtomicKvNonceClaim +} + /** * KV-backed NonceStore for distributed Cloudflare Workers deployments. * * Why this exists: MemoryNonceStore works on a single worker instance, but * Cloudflare routes requests across multiple isolates. Without shared state, * an attacker could retry a replayed nonce against a different isolate and - * have it accepted. Version 2 requires an atomic binding for payment claims. - * This store requires an atomic binding for every payment claim. + * have it accepted. Cloudflare KV has no conditional write, so a plain KV + * binding is not an atomic payment store. Supply `atomicClaim` from D1, + * Durable Objects, or another linearizable service before using this store + * for paid requests. * * Usage: * const nonceStore = new KvNonceStore(env.NONCE_KV, 'x402') * createAgentGateway({ ...config, nonceStore }) */ export class KvNonceStore implements NonceStore { + private readonly atomicClaim?: AtomicKvNonceClaim + constructor( private readonly kv: KVNamespace, /** Key prefix to namespace within a shared KV (default: "nonce"). */ private readonly prefix: string = 'nonce', - ) {} + options: KvNonceStoreOptions = {}, + ) { + this.atomicClaim = options.atomicClaim ?? ( + kv.putIfAbsent + ? async (key, ttlSeconds, ownerId) => { + const value = ownerId ?? '1' + if (ownerId !== undefined) { + const existing = await kv.get(key) + if (existing !== null) return existing === ownerId + } + const inserted = await kv.putIfAbsent!(key, value, { expirationTtl: ttlSeconds }) + if (inserted || ownerId === undefined) return inserted + return (await kv.get(key)) === ownerId + } + : undefined + ) + } async hasSeen(nonce: string): Promise { return (await this.kv.get(this.key(nonce))) !== null } + async markSeen(nonce: string, ttlSeconds: number): Promise { + const ttl = Math.max(ttlSeconds, 60) + await this.kv.put(this.key(nonce), '1', { expirationTtl: ttl }) + } + async claim(nonce: string, ttlSeconds: number, ownerId?: string): Promise { - if (!this.kv.putIfAbsent) { - throw new Error('KvNonceStore requires an atomic putIfAbsent binding for payment claims') + if (!this.atomicClaim) { + throw new Error( + 'KvNonceStore requires an atomicClaim backed by D1, Durable Objects, or an atomic KV extension', + ) } const ttl = Math.max(ttlSeconds, 60) - const key = this.key(nonce) - const value = ownerId ?? '1' - if (ownerId !== undefined) { - const existing = await this.kv.get(key) - if (existing !== null) return existing === ownerId - } - const inserted = await this.kv.putIfAbsent(key, value, { expirationTtl: ttl }) - if (inserted || ownerId === undefined) return inserted - // Another isolate may have won between get and putIfAbsent. Re-read so a - // retry by the same durable operation remains idempotent. - return (await this.kv.get(key)) === ownerId + return this.atomicClaim(this.key(nonce), ttl, ownerId) + } + + /** Used by gateway validation to reject plain, non-atomic KV bindings. */ + hasAtomicClaim(): boolean { + return this.atomicClaim !== undefined } private key(nonce: string): string { @@ -141,26 +183,22 @@ export class KvNonceStore implements NonceStore { } } -/** Claim through the atomic contract or the explicit 0.7.1 legacy path. */ +/** Claim through the one atomic contract used by every payment path. */ export async function claimStoredNonce( store: NonceStore, nonce: string, ttlSeconds: number, ownerId?: string, ): Promise { - if (typeof store.claim === 'function') return store.claim(nonce, ttlSeconds, ownerId) - if (ownerId !== undefined) { - throw new Error('NonceStore.claim is required for atomic payment ownership') - } - if (typeof store.markSeen !== 'function') { - throw new Error('NonceStore.markSeen is required for legacy payment replay protection') + if (typeof store.claim !== 'function') { + throw new Error('NonceStore.claim is required for atomic payment replay protection') } - if (await store.hasSeen(nonce)) return false - await store.markSeen(nonce, ttlSeconds) - return true + return store.claim(nonce, ttlSeconds, ownerId) } /** Durable payment paths must use a store with a single atomic claim operation. */ export function isAtomicNonceStore(store: NonceStore): store is AtomicNonceStore { + const kvStore = store as NonceStore & { hasAtomicClaim?: () => boolean } + if (typeof kvStore.hasAtomicClaim === 'function' && !kvStore.hasAtomicClaim()) return false return typeof store.claim === 'function' } diff --git a/src/payment-operations.ts b/src/payment-operations.ts index e718fe8..01d374d 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -16,6 +16,15 @@ export type PaymentOperationState = | 'reclaimable' | 'reclaimed' +/** Fenced result for a recovery lookup that found no provider operation. */ +export interface PaymentOperationNotFound { + protocolVersion: typeof PAYMENT_PROTOCOL_VERSION + operationId: string + state: 'not-found' +} + +export type PaymentOperationRecoveryResult = PaymentOperation | PaymentOperationNotFound + /** Durable ownership of one signed payment authorization. */ export interface PaymentOperation { protocolVersion: typeof PAYMENT_PROTOCOL_VERSION @@ -79,7 +88,7 @@ export interface PaymentOperations { * Repeated calls must recover an ambiguous acknowledgement by operationId. */ releasePayment(operation: PaymentOperation, reason: string): Promise - reclaimPayment(operationId: string): Promise + reclaimPayment(operationId: string): Promise } export interface MemoryPaymentOperationsOptions { @@ -263,7 +272,11 @@ export class MemoryPaymentOperations implements PaymentOperations { if (current.state === 'releasing') { const flight = this.releaseFlights.get(current.operationId) if (flight) return flight - return this.reclaimPayment(current.operationId) + const recovered = await this.reclaimPayment(current.operationId) + if (recovered.state === 'not-found') { + throw new Error('payment operation disappeared during release recovery') + } + return recovered } if (current.state !== 'claimed' && current.state !== 'executing') { throw new Error(`cannot release payment in state ${current.state}`) @@ -273,9 +286,15 @@ export class MemoryPaymentOperations implements PaymentOperations { return this.runRelease(releasing, reason) } - async reclaimPayment(operationId: string): Promise { + async reclaimPayment(operationId: string): Promise { const current = this.operations.get(operationId) - if (!current) throw new Error('payment operation was not found') + if (!current) { + return { + protocolVersion: PAYMENT_PROTOCOL_VERSION, + operationId, + state: 'not-found', + } + } if (current.state === 'reclaimed') return current if (current.state === 'executing' || current.state === 'retained') { throw new Error(`cannot reclaim payment in state ${current.state} without a usage receipt`) diff --git a/src/payment-recovery-worker.ts b/src/payment-recovery-worker.ts index 32e6ebd..f88d852 100644 --- a/src/payment-recovery-worker.ts +++ b/src/payment-recovery-worker.ts @@ -136,6 +136,13 @@ async function reconcileLeased( throw new Error('x402 payment recovery is not configured') } const recovered = await config.x402.paymentOperations.reclaimPayment(record.payment.operationId) + if (recovered.state === 'not-found') { + if (recovered.operationId !== record.payment.operationId) { + throw new Error('x402 recovery operation id mismatch') + } + await completeRecord(record, config) + return + } if (recovered.state !== 'released' && recovered.state !== 'reclaimed') { throw new Error(`ambiguous x402 claim recovered in state ${recovered.state}`) } @@ -245,6 +252,13 @@ async function releaseRecoveredPayment( if (!record.payment.operation) { if (!config.x402.paymentOperations) throw new Error('x402 payment recovery is not configured') const recovered = await config.x402.paymentOperations.reclaimPayment(record.payment.operationId) + if (recovered.state === 'not-found') { + if (recovered.operationId !== record.payment.operationId) { + throw new Error('x402 recovery operation id mismatch') + } + await completeRecord(record, config) + return + } if (recovered.state !== 'released' && recovered.state !== 'reclaimed') { throw new Error(`x402 release recovered in state ${recovered.state}`) } diff --git a/src/types.ts b/src/types.ts index 136a04e..b7ff4aa 100644 --- a/src/types.ts +++ b/src/types.ts @@ -331,6 +331,15 @@ export interface GatewayConfig { /** Output token limit used when a request omits `max_tokens`. Defaults to 1024. */ defaultOutputTokens?: number + /** + * Return a safe upper bound for the complete provider input. + * Include system, chat framing, retained history, tools, harness, and workspace context. + */ + inputTokenBound?: (input: { + agent: AgentMeta + messages: ChatMessage[] + }) => number + /** Hidden provider spend limits included in the pre-execution payment quote. */ executionBudget?: { maxReasoningTokens?: number @@ -363,8 +372,8 @@ export interface GatewayConfig { observer?: import('./observer').GatewayObserver /** - * A2A protocol configuration. When set, the gateway exposes the A2A - * surface alongside its OpenAI-compatible endpoints: + * A2A protocol configuration. The gateway exposes A2A with an in-memory + * task store by default. Set this object to provide durable storage or push: * GET /:slug/.well-known/agent.json — AgentCard discovery * POST /:slug — JSON-RPC 2.0 endpoint * methods: message/send, message/stream, tasks/get, tasks/cancel diff --git a/src/verify.ts b/src/verify.ts index bfb3152..20980fc 100644 --- a/src/verify.ts +++ b/src/verify.ts @@ -195,7 +195,7 @@ export async function verifyX402( * Returns authenticated identity if valid, null otherwise. * In demo mode, accepts any well-formed Payment header with an identity. */ -export async function verifyMpp( +export async function verifyMppCredential( authHeader: string, config: MppConfig, x402Config: X402Config, @@ -305,6 +305,29 @@ export async function verifyMpp( } } +/** + * Verify an MPP credential using the 0.7.1 public return shape. + * Rich durable callers use verifyMppCredential instead. + */ +export async function verifyMpp( + authHeader: string, + config: MppConfig, + x402Config: X402Config, + nonceStore?: NonceStore, + minimumAmount = 1n, + markNonce = true, +): Promise { + const authenticated = await verifyMppCredential( + authHeader, + config, + x402Config, + nonceStore, + minimumAmount, + markNonce, + ) + return authenticated?.consumerId ?? null +} + /** * Default API key verifier — accepts any `sk_agent_*` key (demo mode). * Override in GatewayConfig.verifyApiKey for production. diff --git a/tests/a2a-atomicity.test.ts b/tests/a2a-atomicity.test.ts index d700a21..156c0cf 100644 --- a/tests/a2a-atomicity.test.ts +++ b/tests/a2a-atomicity.test.ts @@ -150,7 +150,7 @@ function recoveredOperation(operation: PaymentOperation) { } describe('A2A task atomicity and restart recovery', () => { - it('rejects a non-atomic task store outside explicit demo mode', () => { + it('keeps OpenAI available and returns 503 for a non-atomic A2A store', async () => { const legacyStore: TaskStore = { get: async () => undefined, put: async () => undefined, @@ -158,9 +158,14 @@ describe('A2A task atomicity and restart recovery', () => { } const config = atomicityConfig(legacyStore, { runs: 0, records: 0, settlements: 0 }) - expect(() => createAgentGateway(config)).toThrow( - /A2A production task store must implement createIfAbsent and compareAndSet/, - ) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const response = await app.request(`/v1/agents/${agent.slug}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({}), + }) + expect(response.status).toBe(503) }) it('lets exactly one of two workers create, execute, and settle a task', async () => { diff --git a/tests/kv-stores.test.ts b/tests/kv-stores.test.ts index 08e83d5..7e08abe 100644 --- a/tests/kv-stores.test.ts +++ b/tests/kv-stores.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, afterEach } from 'vitest' -import { KvNonceStore, type KVNamespace as NonceKV } from '../src/nonce-store' +import { isAtomicNonceStore, KvNonceStore, type KVNamespace as NonceKV } from '../src/nonce-store' import { KvRateLimitStore, checkRateLimit } from '../src/rate-limit' import type { KVNamespace as RlKV } from '../src/rate-limit' @@ -110,8 +110,31 @@ describe('KvNonceStore', () => { } const store = new KvNonceStore(cloudflareKv) - await expect(store.claim('legacy', 300)).rejects.toThrow('atomic putIfAbsent') - await expect(store.claim('version-2', 300, 'operation-1')).rejects.toThrow('atomic putIfAbsent') + expect(isAtomicNonceStore(store)).toBe(false) + await expect(store.claim('legacy', 300)).rejects.toThrow('atomicClaim') + await expect(store.claim('version-2', 300, 'operation-1')).rejects.toThrow('atomicClaim') + }) + + it('accepts an explicitly supplied atomic claim backend for standard KV', async () => { + const backing = new StubKV() + const cloudflareKv: NonceKV = { + get: backing.get.bind(backing), + put: backing.put.bind(backing), + delete: backing.delete.bind(backing), + } + const claims = new Map() + const store = new KvNonceStore(cloudflareKv, 'nonce', { + atomicClaim: async (key, ttlSeconds, ownerId) => { + const existing = claims.get(key) + if (existing !== undefined) return ownerId !== undefined && existing === ownerId + claims.set(key, ownerId ?? '1') + await cloudflareKv.put(key, ownerId ?? '1', { expirationTtl: ttlSeconds }) + return true + }, + }) + + expect(isAtomicNonceStore(store)).toBe(true) + expect(await store.claim('backend', 300, 'operation-1')).toBe(true) }) }) diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index 3a04b77..f44ff13 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -361,7 +361,7 @@ describe('POST /:slug/chat/completions — auth paths', () => { const body = await res.json() as { error: { payment_methods: string[]; x402: Record } } expect(body.error.payment_methods).toContain('x402') expect(body.error.x402.operator).toBe(operatorAddress) - expect(body.error.x402.required_amount).toBe('184861') + expect(body.error.x402.required_amount).toBe('185460') expect(body.error.x402.max_output_tokens).toBe(1024) }) diff --git a/tests/nonce-store.test.ts b/tests/nonce-store.test.ts index 0a38aaf..d1d80dd 100644 --- a/tests/nonce-store.test.ts +++ b/tests/nonce-store.test.ts @@ -67,17 +67,15 @@ describe('MemoryNonceStore', () => { }) describe('claimStoredNonce', () => { - it('keeps the 0.7.1 check-and-mark store contract for legacy claims', async () => { + it('fails closed for the 0.7.1 check-and-mark store contract', async () => { const seen = new Set() const legacyStore: NonceStore = { hasSeen: async (nonce) => seen.has(nonce), markSeen: async (nonce) => { seen.add(nonce) }, } - expect(await claimStoredNonce(legacyStore, 'legacy', 60)).toBe(true) - expect(await claimStoredNonce(legacyStore, 'legacy', 60)).toBe(false) - await expect(claimStoredNonce(legacyStore, 'legacy', 60, 'operation-1')) - .rejects.toThrow('atomic payment ownership') + await expect(claimStoredNonce(legacyStore, 'legacy', 60)) + .rejects.toThrow('atomic payment replay protection') }) it('fails closed for a store without either replay contract', async () => { @@ -85,7 +83,7 @@ describe('claimStoredNonce', () => { hasSeen: async () => false, } as unknown as NonceStore - await expect(claimStoredNonce(legacyStore, 'legacy', 60)).rejects.toThrow('markSeen') + await expect(claimStoredNonce(legacyStore, 'legacy', 60)).rejects.toThrow('atomic payment replay protection') }) }) diff --git a/tests/payment-recovery.test.ts b/tests/payment-recovery.test.ts index 745a0cb..d1c53a3 100644 --- a/tests/payment-recovery.test.ts +++ b/tests/payment-recovery.test.ts @@ -493,7 +493,8 @@ describe('generic MPP charge lifecycle', () => { releaseObserver() const response = await pending - expect(response.status).toBe(402) + // The streaming response is committed before a second worker fences the row. + expect(response.status).toBe(200) expect(runs).toBe(0) expect(lifecycle.confirmations).toBe(1) expect(lifecycle.refunds).toBe(1) @@ -738,7 +739,7 @@ describe('durable OpenAI recovery', () => { await response.text() const pending = await recoveryStore.get(operationId) - expect(pending?.state).toBe('executing') + expect(pending?.state).toBe('claimed') expect(pending?.lease).toBeUndefined() expect(operations.get(operationId)?.state).toBe('claimed') }) diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts new file mode 100644 index 0000000..d40cacf --- /dev/null +++ b/tests/pr11-regressions.test.ts @@ -0,0 +1,565 @@ +import { Hono } from 'hono' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { InMemoryTaskStore } from '../src/a2a/task-store' +import { createAgentGateway } from '../src/middleware' +import { dispatchSandboxStreamRich, requiredX402Amount } from '../src/dispatch' +import { MemoryNonceStore, claimStoredNonce, type NonceStore } from '../src/nonce-store' +import { MemoryPaymentOperations } from '../src/payment-operations' +import { MemoryPaymentRecoveryStore, type PaymentRecoveryRecord } from '../src/payment-recovery' +import { recoverPayment } from '../src/payment-recovery-worker' +import { verifyMpp } from '../src/verify' +import type { AgentMeta, GatewayConfig, SandboxStreamEvent } from '../src/types' +import type { MppConfig } from '../src/types' + +const operatorAddress = '0x1111111111111111111111111111111111111111' +const commitment = `0x${'ab'.repeat(32)}` + +const agent: AgentMeta = { + id: 'agent-pr11', + ownerId: 'owner-pr11', + slug: 'pr11', + systemPrompt: 'You are a test agent.', + pricePerTokenUsd: 0.000001, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, +} + +afterEach(() => { + vi.useRealTimers() +}) + +function paymentHeader(nonce: string): string { + return JSON.stringify({ + commitment, + signature: '0xsig', + operator: operatorAddress, + amount: '1000000000', + nonce, + expiry: String(Math.floor(Date.now() / 1000) + 600), + }) +} + +function usage(inputTokens = 1) { + return { + inputTokens, + outputTokens: 1, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.000002, + budgetEnforced: true, + } +} + +function sandbox(events: SandboxStreamEvent[] = [ + { type: 'sandbox.usage', data: { usage: usage() } }, +]): GatewayConfig['getSandbox'] extends (...args: never[]) => infer R ? R : never { + return { + async *streamPrompt() { + yield* events + }, + } as Awaited> +} + +function durableConfig( + overrides: Partial = {}, +): GatewayConfig { + return { + resolveAgent: async () => agent, + getSandbox: async () => sandbox(), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations({ onReclaim: async () => undefined }), + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, + ...overrides, + } +} + +describe('PR #11 production regressions', () => { + it('does not mark payment execution before sandbox acquisition succeeds', async () => { + const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) + let sandboxReady = false + let beganBeforeSandbox = false + const originalBegin = operations.beginPaymentExecution.bind(operations) + operations.beginPaymentExecution = async (operation) => { + if (!sandboxReady) beganBeforeSandbox = true + return originalBegin(operation) + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + getSandbox: async () => { + sandboxReady = true + return sandbox() + }, + }))) + + const response = await app.request('/v1/agents/pr11/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('execution-order'), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'run' }] }), + }) + await response.text() + + expect(beganBeforeSandbox).toBe(false) + }) + + it('renews the live execution callback while a provider stream is quiet', async () => { + vi.useFakeTimers() + const controller = new AbortController() + let heartbeatCount = 0 + const config: GatewayConfig = { + ...durableConfig(), + paymentRecovery: { store: new MemoryPaymentRecoveryStore(), receiptTimeoutMs: 300 }, + getSandbox: async () => ({ + async *streamPrompt(_message: string, options?: { signal?: AbortSignal }) { + await new Promise((resolve) => { + options?.signal?.addEventListener('abort', () => resolve(), { once: true }) + }) + }, + }), + } + + const consume = (async () => { + for await (const _event of dispatchSandboxStreamRich( + agent, + 'run', + commitment, + config, + controller.signal, + undefined, + 4, + undefined, + true, + undefined, + undefined, + async () => { heartbeatCount += 1 }, + )) { + // The stream exits only after the request aborts. + } + })() + await Promise.resolve() + await Promise.resolve() + await vi.advanceTimersByTimeAsync(1_000) + + expect(heartbeatCount).toBeGreaterThan(0) + controller.abort() + await consume + }) + + it('does not execute after a cancellation wins on another worker', async () => { + const taskStore = new InMemoryTaskStore() + let sandboxEntered!: () => void + const sandboxReady = new Promise((resolve) => { sandboxEntered = resolve }) + let releaseSandbox!: () => void + const sandboxReleased = new Promise((resolve) => { releaseSandbox = resolve }) + let runs = 0 + + const makeConfig = (worker: 'runner' | 'canceler'): GatewayConfig => ({ + resolveAgent: async () => agent, + getSandbox: async () => { + if (worker === 'runner') { + sandboxEntered() + await sandboxReleased + } + return { + async *streamPrompt() { + runs += 1 + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + } + }, + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => true, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + }) + + const runner = new Hono() + runner.route('/v1/agents', createAgentGateway(makeConfig('runner'))) + const canceler = new Hono() + canceler.route('/v1/agents', createAgentGateway(makeConfig('canceler'))) + + const request = runner.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9001'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'pr11-cancel-race', + contextId: 'pr11-context', + messageId: 'pr11-message', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + await sandboxReady + + const cancel = await canceler.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'pr11-cancel-race' }, + }), + }) + expect(cancel.status).toBe(200) + releaseSandbox() + const runnerResponse = await request + await runnerResponse.text() + + expect(runs).toBe(0) + expect((await taskStore.get('pr11-cancel-race'))?.status.state).toBe('canceled') + }) + + it('quotes retained A2A history before charging a continuation', async () => { + const taskStore = new InMemoryTaskStore() + let invocations = 0 + const longHistory = 'history '.repeat(500) + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + invocations += 1 + if (invocations === 1) { + yield { + type: 'input-required', + data: { inputRequired: { prompt: 'What next?' } }, + } + yield { type: 'sandbox.usage', data: { usage: usage(1) } } + return + } + yield { type: 'sandbox.usage', data: { usage: usage(200) } } + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => true, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const send = (nonce: string, text: string) => app.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader(nonce), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: nonce, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'pr11-continuation', + contextId: 'pr11-context', + messageId: `message-${nonce}`, + parts: [{ kind: 'text', text }], + }, + }, + }), + }) + + const first = await send('9004', longHistory) + expect((await first.json() as { result?: { status?: { state?: string } } }).result?.status?.state) + .toBe('input-required') + const second = await send('9005', 'continue') + + expect(second.status).toBe(200) + expect((await second.json() as { result?: { status?: { state?: string } } }).result?.status?.state) + .toBe('completed') + }) + + it('uses the configured complete provider input bound before quoting', async () => { + let quotedMessages: Array<{ role: string; content: string }> | undefined + const inputTokenBound = ({ messages }: { messages: Array<{ role: string; content: string }> }) => { + quotedMessages = messages + return 4_096 + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + maxOutputTokens: 1_024, + defaultOutputTokens: 1_024, + inputTokenBound, + }))) + + const response = await app.request('/v1/agents/pr11/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + messages: [ + { role: 'user', content: 'first turn' }, + { role: 'assistant', content: 'prior answer' }, + { role: 'user', content: 'current turn' }, + ], + }), + }) + const body = await response.json() as { + error?: { x402?: { required_amount?: string } } + } + + expect(response.status).toBe(402) + expect(quotedMessages).toEqual([ + { role: 'user', content: 'first turn' }, + { role: 'assistant', content: 'prior answer' }, + { role: 'user', content: 'current turn' }, + ]) + expect(body.error?.x402?.required_amount).toBe( + requiredX402Amount(agent.pricePerTokenUsd, 4_096, 1_024, 6, 1_024, 1_024) + .toString(), + ) + }) + + it('reconciles an x402 claiming row when the provider has no operation', async () => { + const now = Date.now() + const id = 'x402:missing-provider-operation' + const store = new MemoryPaymentRecoveryStore() + const record: PaymentRecoveryRecord = { + version: 1, + id, + revision: 0, + state: 'claiming', + payment: { kind: 'x402', operationId: id }, + attribution: { + requestId: 'request-pr11', + agentId: agent.id, + agentSlug: agent.slug, + consumerId: commitment, + paymentMethod: 'x402', + startMs: now, + pricePerTokenUsd: agent.pricePerTokenUsd, + platformFeePercent: agent.platformFeePercent, + requiredAmount: '1', + currencyDecimals: 6, + maxOutputTokens: 1, + executionBudget: { + maxInputTokens: 1, + maxOutputTokens: 1, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 0.000001, + }, + }, + workStarted: false, + usageRecorded: false, + attempts: 0, + nextAttemptAt: now, + lease: { id: 'live-fence', expiresAt: now - 1 }, + createdAt: now, + updatedAt: now, + } + await store.createIfAbsent(record) + const config = durableConfig({ + paymentRecovery: { store }, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations({ onReclaim: async () => undefined }), + }, + }) + + const recovered = await recoverPayment(id, config, { force: true }) + + expect(recovered?.state).toBe('reconciled') + }) + + it('does not leave an abandoned working A2A task after payment recovery', async () => { + const taskStore = new InMemoryTaskStore() + const recoveryStore = new MemoryPaymentRecoveryStore() + const recoveryId = 'x402:abandoned-working-task' + const now = Date.now() + await recoveryStore.createIfAbsent({ + version: 1, + id: recoveryId, + revision: 0, + state: 'reconciled', + payment: { kind: 'x402', operationId: recoveryId }, + attribution: { + requestId: 'request-abandoned', + agentId: agent.id, + agentSlug: agent.slug, + consumerId: commitment, + paymentMethod: 'x402', + startMs: now, + pricePerTokenUsd: agent.pricePerTokenUsd, + platformFeePercent: agent.platformFeePercent, + requiredAmount: '1', + currencyDecimals: 6, + maxOutputTokens: 1, + executionBudget: { + maxInputTokens: 1, + maxOutputTokens: 1, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 0.000001, + }, + }, + workStarted: true, + usageRecorded: false, + attempts: 0, + nextAttemptAt: Number.MAX_SAFE_INTEGER, + lease: undefined, + createdAt: now, + updatedAt: now, + }) + await taskStore.put({ + kind: 'task', + id: 'abandoned-task', + contextId: 'abandoned-context', + status: { state: 'working', timestamp: new Date(now).toISOString() }, + metadata: { gatewayPaymentRecovery: { version: 1, id: recoveryId } }, + }) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + paymentRecovery: { store: recoveryStore }, + a2a: { taskStore }, + }))) + + const response = await app.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/get', + params: { id: 'abandoned-task' }, + }), + }) + const body = await response.json() as { result?: { status?: { state?: string } } } + + expect(response.status).toBe(200) + expect(body.result?.status?.state).toBe('failed') + }) + + it('rejects a legacy check-then-mark nonce store instead of racing two payments', async () => { + const seen = new Set() + const legacyStore: NonceStore = { + hasSeen: async () => false, + markSeen: async (nonce) => { + await Promise.resolve() + seen.add(nonce) + }, + } + + const results = await Promise.allSettled([ + claimStoredNonce(legacyStore, 'replayed', 60), + claimStoredNonce(legacyStore, 'replayed', 60), + ]) + + expect(results.every((result) => result.status === 'rejected')).toBe(true) + expect(seen).toHaveLength(0) + }) + + it('preserves the 0.7.1 verifyMpp return contract', async () => { + const config: MppConfig = { + realm: 'gateway.test', + method: 'stripe', + verifySigner: async () => 'consumer-pr11', + } + const credential = Buffer.from(JSON.stringify({ sharedPaymentToken: 'token' })).toString('base64url') + + const result = await verifyMpp( + `Payment stripe ${credential}`, + config, + { operatorAddress, chainId: 1, demoMode: true }, + ) + + expect(result).toBe('consumer-pr11') + }) + + it('keeps the public markSeen method on the in-memory nonce store', async () => { + const store = new MemoryNonceStore() + expect(typeof store.markSeen).toBe('function') + await store.markSeen!('compatibility', 60) + expect(await store.hasSeen('compatibility')).toBe(true) + }) + + it('does not initialize A2A when an OpenAI-only gateway omits it', () => { + expect(() => createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => sandbox(), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: false, + verifySigner: async () => true, + }, + })).not.toThrow() + }) + + it('keeps the OpenAI surface available when an old A2A store is configured', async () => { + const legacyTaskStore = { + get: async () => undefined, + put: async () => undefined, + delete: async () => undefined, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => sandbox(), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: false, + verifySigner: async () => true, + }, + a2a: { taskStore: legacyTaskStore }, + })) + + const discovery = await app.request('/v1/agents/pr11/chat/completions') + const a2a = await app.request('/v1/agents/pr11', { method: 'POST', body: '{}' }) + + expect(discovery.status).toBe(200) + expect(a2a.status).toBe(503) + }) +}) diff --git a/tests/verify.test.ts b/tests/verify.test.ts index 682c54f..6c9f8f3 100644 --- a/tests/verify.test.ts +++ b/tests/verify.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, beforeEach } from 'vitest' -import { verifyX402, verifyMpp, defaultVerifyApiKey } from '../src/verify' +import { verifyX402, verifyMppCredential, defaultVerifyApiKey } from '../src/verify' import { MemoryNonceStore, type NonceStore } from '../src/nonce-store' import type { X402Config, MppConfig } from '../src/types' @@ -116,7 +116,7 @@ describe('verifyX402', () => { expect(ttls[0]).toBeLessThanOrEqual(7200) }) - it('keeps 0.7.1 custom nonce stores working on the legacy path', async () => { + it('fails closed for a 0.7.1 custom nonce store without atomic claims', async () => { const seen = new Set() const nonceStore = { hasSeen: async (nonce: string) => seen.has(nonce), @@ -124,7 +124,7 @@ describe('verifyX402', () => { } as unknown as NonceStore const payload = buildSpendAuth({ nonce: '101' }) - expect(await verifyX402(payload, baseConfig, nonceStore)).toBe('0xCommitmentAlice') + expect(await verifyX402(payload, baseConfig, nonceStore)).toBeNull() expect(await verifyX402(payload, baseConfig, nonceStore)).toBeNull() }) @@ -190,7 +190,7 @@ describe('verifyMpp', () => { it('parses a valid Payment header and returns the signer', async () => { const header = buildCredential({ commitment: '0xAlice', operator: operatorAddress, amount: '1000', nonce: '5' }) - expect(await verifyMpp(header, mppConfig, baseConfig)).toMatchObject({ + expect(await verifyMppCredential(header, mppConfig, baseConfig)).toMatchObject({ consumerId: '0xAlice', replayKey: '0xalice:5', }) @@ -215,11 +215,11 @@ describe('verifyMpp', () => { const productionX402: X402Config = { ...baseConfig, demoMode: false } const nonceStore = new MemoryNonceStore() - expect(await verifyMpp(header, config, productionX402, nonceStore)).toMatchObject({ + expect(await verifyMppCredential(header, config, productionX402, nonceStore)).toMatchObject({ consumerId: 'mpp:alice', replayKey: '0xalice:6', }) - expect(await verifyMpp(header, config, productionX402, nonceStore)).toBeNull() + expect(await verifyMppCredential(header, config, productionX402, nonceStore)).toBeNull() expect(seen).toHaveLength(1) expect(seen[0].method).toBe('blueprintevm') expect(seen[0].credential).toContain('commitment') @@ -243,7 +243,7 @@ describe('verifyMpp', () => { }, } - await expect(verifyMpp( + await expect(verifyMppCredential( header, legacyConfig, { ...baseConfig, demoMode: false }, @@ -264,7 +264,7 @@ describe('verifyMpp', () => { expiry: String(Math.floor(Date.now() / 1000) + 600), } expect(await verifyX402(JSON.stringify({ ...payload, nonce: '1' }), baseConfig, nonceStore)).toBe('0xAlice') - expect(await verifyMpp( + expect(await verifyMppCredential( buildCredential({ ...payload, commitment: '0xALICE' }), mppConfig, baseConfig, @@ -289,7 +289,7 @@ describe('verifyMpp', () => { }, } - expect(await verifyMpp(header, config, { ...baseConfig, demoMode: false }, undefined, 20000n)).toBeNull() + expect(await verifyMppCredential(header, config, { ...baseConfig, demoMode: false }, undefined, 20000n)).toBeNull() expect(calls).toBe(0) }) @@ -302,14 +302,14 @@ describe('verifyMpp', () => { expiry: String(Math.floor(Date.now() / 1000) + 600), }).replace('Payment blueprintevm ', 'Payment BLUEPRINTEVM ') - expect(await verifyMpp( + expect(await verifyMppCredential( underfunded, { realm: 'agents.tangle.tools' }, baseConfig, undefined, 20000n, )).toBeNull() - expect(await verifyMpp( + expect(await verifyMppCredential( underfunded, mppConfig, baseConfig, @@ -321,7 +321,7 @@ describe('verifyMpp', () => { it('rejects MPP in production when no method verifier is configured', async () => { const header = buildCredential({ commitment: '0xAlice', operator: operatorAddress }) const productionX402: X402Config = { ...baseConfig, demoMode: false } - expect(await verifyMpp(header, mppConfig, productionX402)).toBeNull() + expect(await verifyMppCredential(header, mppConfig, productionX402)).toBeNull() }) it('does not reuse the x402 verifier for a different MPP method', async () => { @@ -332,34 +332,34 @@ describe('verifyMpp', () => { demoMode: false, verifySigner: async () => true, } - expect(await verifyMpp(header, { ...mppConfig, method: 'stripe' }, productionX402)).toBeNull() + expect(await verifyMppCredential(header, { ...mppConfig, method: 'stripe' }, productionX402)).toBeNull() }) it('falls back to the `from` field when no `commitment` present — regression: EIP-3009 wallets expose `from` only', async () => { const header = buildCredential({ from: '0xWallet', to: operatorAddress, value: '1000' }) - expect(await verifyMpp(header, mppConfig, baseConfig)).toMatchObject({ + expect(await verifyMppCredential(header, mppConfig, baseConfig)).toMatchObject({ consumerId: '0xWallet', }) }) it('rejects malformed Payment header shape', async () => { - expect(await verifyMpp('Bearer sk_agent_123', mppConfig, baseConfig)).toBeNull() - expect(await verifyMpp('Payment', mppConfig, baseConfig)).toBeNull() - expect(await verifyMpp('Payment blueprintevm', mppConfig, baseConfig)).toBeNull() + expect(await verifyMppCredential('Bearer sk_agent_123', mppConfig, baseConfig)).toBeNull() + expect(await verifyMppCredential('Payment', mppConfig, baseConfig)).toBeNull() + expect(await verifyMppCredential('Payment blueprintevm', mppConfig, baseConfig)).toBeNull() }) it('rejects bad base64url — regression: decode crash must return null', async () => { - expect(await verifyMpp('Payment blueprintevm !@#$not-b64$#@!', mppConfig, baseConfig)).toBeNull() + expect(await verifyMppCredential('Payment blueprintevm !@#$not-b64$#@!', mppConfig, baseConfig)).toBeNull() }) it('rejects operator mismatch', async () => { const header = buildCredential({ commitment: '0xAlice', operator: '0xWrongOp', amount: '1000' }) - expect(await verifyMpp(header, mppConfig, baseConfig)).toBeNull() + expect(await verifyMppCredential(header, mppConfig, baseConfig)).toBeNull() }) it('rejects non-numeric amount/nonce — regression: BigInt throw should become null, not crash', async () => { const header = buildCredential({ commitment: '0xAlice', operator: operatorAddress, amount: 'not-a-number' }) - expect(await verifyMpp(header, mppConfig, baseConfig)).toBeNull() + expect(await verifyMppCredential(header, mppConfig, baseConfig)).toBeNull() }) }) From 58fe4790fa4948aaa0f1b92b0299e7c2bbbf39b9 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 03:02:42 -0600 Subject: [PATCH 24/32] fix(gateway): harden A2A recovery and request boundaries --- README.md | 5 +- docs/a2a-long-horizon.md | 7 +- src/a2a/handler.ts | 56 +++++++++- src/a2a/push-notifications.ts | 12 +- src/a2a/task-store-sql.ts | 79 ++++++++++---- src/a2a/task-store.ts | 3 +- tests/a2a-lifecycle.test.ts | 2 - tests/pr11-regressions.test.ts | 193 +++++++++++++++++++++++++++++++++ 8 files changed, 325 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index e2a3f58..d4cccb2 100644 --- a/README.md +++ b/README.md @@ -109,9 +109,12 @@ Wire protocol handlers only translate their request and response shapes. The gateway speaks Google's A2A protocol alongside its OpenAI-compatible surface: discovery via `.well-known/agent.json`, JSON-RPC 2.0 dispatch for `message/send`, `message/stream`, `tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and the four `tasks/pushNotificationConfig/*` methods. Long-horizon agents — durable tasks across worker restarts, webhook delivery on terminal state, `input-required` pauses with multi-turn continuation — are documented in [`docs/a2a-long-horizon.md`](./docs/a2a-long-horizon.md). Production A2A task control requires `a2a.authorizeTaskAccess`; explicit demo mode is the local-test exception. Custom production task stores must implement atomic `createIfAbsent` and `compareAndSet` methods. -Task stores must retain gateway recovery metadata until reconciliation clears it. +Task stores must retain payment recovery metadata until reconciliation clears it. +The short-lived `gatewaySubmission` marker is not a payment recovery record and may expire with its task. The bundled memory and SQL stores enforce this rule even after the normal task TTL. Push destinations must use HTTPS without URL credentials. +Push delivery does not follow redirects. +Tasks created before this release have no recorded origin and fail closed; migrate them with a verified owner binding or let them expire. ## Tier diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index babb2f8..f1df802 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -39,11 +39,16 @@ Usage attribution must atomically upsert by `requestId`. The finalization record stores whether attribution was acknowledged, so recovery does not repeat an acknowledged usage event. If a legacy record is malformed or has no recoverable operation, the gateway expires the task as failed. If work exists without a receipt, the outbox settles the original quoted ceiling after its configured timeout. -The task-store TTL cannot delete a task while any gateway recovery marker remains. +The task-store TTL cannot delete a task while any payment recovery marker remains. +The short-lived `gatewaySubmission` marker is not a payment recovery marker. +An abandoned submission without a payment marker may expire after the normal task TTL. Task control methods (`tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and push configuration methods) require `a2a.authorizeTaskAccess` in production. The hook receives the task and request headers so the application can enforce task ownership. Explicit `x402.demoMode` permits these methods without the hook for local tests only. +This is a fail-closed upgrade boundary. +Tasks stored before this release do not have a `gatewayOrigin` binding and cannot pass the default ownership check. +Migrate those records with a verified owner binding, or allow them to expire before enabling production control methods. ### D1 (Cloudflare Workers) diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 17f9025..4f4e8f1 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -144,6 +144,9 @@ const TASK_ORIGIN_METADATA_KEY = 'gatewayOrigin' const TASK_SUBMISSION_METADATA_KEY = 'gatewaySubmission' const TASK_SUBMISSION_RECOVERY_METADATA_KEY = 'gatewaySubmissionRecovery' const TASK_SUBMISSION_LEASE_MS = 5 * 60 * 1000 +const MAX_A2A_BODY_BYTES = 64 * 1024 + +class RequestBodyTooLargeError extends Error {} /** * Per-gateway in-process registry of cancellable runs. Keyed by task id; @@ -221,14 +224,20 @@ export function createA2AHandlers(deps: A2AHandlerDeps) { // Body size limit (DoS prevention) — mirrors the OpenAI-compat handler. const contentLength = Number.parseInt(c.req.header('Content-Length') ?? '0', 10) - if (contentLength > 65536) { + if (contentLength > MAX_A2A_BODY_BYTES) { return c.json(fail(null, A2A_ERROR_CODES.INVALID_REQUEST, 'request body too large (max 64KB)'), 413) } let raw: unknown try { - raw = await c.req.json() - } catch { + raw = await readJsonBody(c.req.raw) + } catch (error) { + if (error instanceof RequestBodyTooLargeError) { + return c.json( + fail(null, A2A_ERROR_CODES.INVALID_REQUEST, 'request body too large (max 64KB)'), + 413, + ) + } return c.json(fail(null, A2A_ERROR_CODES.PARSE_ERROR, 'invalid JSON'), 400) } const parsed = parseEnvelope(raw) @@ -2455,12 +2464,51 @@ function agentMessage(task: Task, text: string): Message { kind: 'message', role: 'agent', parts: [{ kind: 'text', text }], - messageId: `${task.id}-status-${task.status.state}-${nowIso()}`, + messageId: `${task.id}-status-${task.status.state}-${stableMessageDigest(text)}`, taskId: task.id, contextId: task.contextId, } } +function stableMessageDigest(value: string): string { + let hash = 2166136261 + for (let index = 0; index < value.length; index += 1) { + hash ^= value.charCodeAt(index) + hash = Math.imul(hash, 16777619) + } + return (hash >>> 0).toString(16).padStart(8, '0') +} + +async function readJsonBody(request: Request): Promise { + if (!request.body) return await request.json() + const reader = request.body.getReader() + const chunks: Uint8Array[] = [] + let total = 0 + try { + while (true) { + const { done, value } = await reader.read() + if (done) break + if (value) { + total += value.byteLength + if (total > MAX_A2A_BODY_BYTES) { + await reader.cancel().catch(() => undefined) + throw new RequestBodyTooLargeError('request body too large') + } + chunks.push(value) + } + } + } finally { + reader.releaseLock() + } + const body = new Uint8Array(total) + let offset = 0 + for (const chunk of chunks) { + body.set(chunk, offset) + offset += chunk.byteLength + } + return JSON.parse(new TextDecoder().decode(body)) as unknown +} + /** * Fire-and-forget push delivery. Idempotent w.r.t. push: if the task hasn't * reached a terminal state, this is a no-op. The dispatch logs failures via diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index 36ba3d5..ff3e97c 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -252,13 +252,23 @@ export async function deliverPushNotifications(args: { if (url.protocol !== 'https:' || url.username !== '' || url.password !== '') { throw new Error('push notification URL must use https without credentials') } - const res = await fetcher(config.url, { method: 'POST', headers, body }) + const res = await fetcher(config.url, { + method: 'POST', + headers, + body, + // Never follow a user-controlled redirect. The redirected destination + // could be an internal HTTP service or instance metadata endpoint. + redirect: 'manual', + }) result = { taskId: args.task.id, configId: config.id, url: config.url, ok: res.ok, status: res.status, + ...(res.status >= 300 && res.status < 400 + ? { error: 'push notification redirect rejected' } + : {}), } } catch (err) { result = { diff --git a/src/a2a/task-store-sql.ts b/src/a2a/task-store-sql.ts index b4ba33f..dcf1108 100644 --- a/src/a2a/task-store-sql.ts +++ b/src/a2a/task-store-sql.ts @@ -123,6 +123,34 @@ export class SqlTaskStore implements TaskStore { return this.opts.table ?? 'a2a_tasks' } + private async readRow(id: string): Promise<{ + payload: string + updatedAt: number + } | undefined> { + const rows = await this.db.query<{ payload: string; updated_at: number }>( + `SELECT payload, updated_at FROM ${this.table} WHERE id = ?`, + [id], + ) + const row = rows[0] + return row ? { payload: row.payload, updatedAt: row.updated_at } : undefined + } + + private isExpired(updatedAt: number, task: Task): boolean { + return Date.now() - updatedAt > this.ttlMs && !hasPendingPaymentRecovery(task) + } + + private async deleteObservedRow( + id: string, + payload: string, + updatedAt: number, + ): Promise { + const result = await this.db.exec( + `DELETE FROM ${this.table} WHERE id = ? AND payload = ? AND updated_at = ?`, + [id, payload, updatedAt], + ) + return result.rowsAffected + } + /** Idempotent. Call once at deploy. */ async migrate(): Promise { await this.db.exec(TASKS_TABLE_DDL(this.table)) @@ -130,25 +158,27 @@ export class SqlTaskStore implements TaskStore { } async get(id: string): Promise { - const rows = await this.db.query<{ payload: string; updated_at: number }>( - `SELECT payload, updated_at FROM ${this.table} WHERE id = ?`, - [id], - ) - const row = rows[0] + const row = await this.readRow(id) if (!row) return undefined const task = JSON.parse(row.payload) as Task - if (Date.now() - row.updated_at > this.ttlMs && !hasPendingPaymentRecovery(task)) { + if (this.isExpired(row.updatedAt, task)) { // Delete only the version that was observed as stale. A refresh can reuse - // the same payload, so payload equality alone does not protect the row. - void this.db.exec( - `DELETE FROM ${this.table} WHERE id = ? AND payload = ? AND updated_at = ?`, - [id, row.payload, row.updated_at], - ) + // the same payload, so payload equality and updated_at both fence the row. + await this.deleteObservedRow(id, row.payload, row.updatedAt) return undefined } return task } + private async insert(task: Task): Promise { + const payload = JSON.stringify(task) + const result = await this.db.exec( + `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at) VALUES (?, ?, ?, ?, ?)`, + [task.id, task.contextId, task.status.state, payload, Date.now()], + ) + return result.rowsAffected + } + async put(task: Task): Promise { const payload = JSON.stringify(task) const updatedAt = Date.now() @@ -168,18 +198,25 @@ export class SqlTaskStore implements TaskStore { } async createIfAbsent(task: Task): Promise { - const payload = JSON.stringify(task) try { - const result = await this.db.exec( - `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at) VALUES (?, ?, ?, ?, ?)`, - [task.id, task.contextId, task.status.state, payload, Date.now()], - ) - return result.rowsAffected === 1 + return (await this.insert(task)) === 1 } catch (error) { - // SQL dialects report duplicate primary keys as errors. Convert only a - // confirmed existing row into the protocol-level "already exists" result. - if (await this.get(task.id)) return false - throw error + // SQL dialects report duplicate primary keys as errors. Inspect the raw + // row so an expired row can be removed and retried in the same call. + const row = await this.readRow(task.id) + if (!row) throw error + const existing = JSON.parse(row.payload) as Task + if (!this.isExpired(row.updatedAt, existing)) return false + + await this.deleteObservedRow(task.id, row.payload, row.updatedAt) + try { + return (await this.insert(task)) === 1 + } catch (retryError) { + // Another writer may have won the retry after the stale row was + // removed. Return the normal idempotency result in that case. + if (await this.get(task.id)) return false + throw retryError + } } } diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index 7746875..cae2f88 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -23,10 +23,9 @@ const PAYMENT_RECOVERY_KEYS = [ 'gatewayFinalizing', 'gatewayPaymentRelease', 'gatewayPaymentRecovery', - 'gatewaySubmission', ] as const -/** Recovery-bearing tasks must remain readable until reconciliation clears the marker. */ +/** Payment recovery tasks must remain readable until reconciliation clears the marker. */ export function hasPendingPaymentRecovery(task: Task): boolean { return PAYMENT_RECOVERY_KEYS.some((key) => task.metadata?.[key] !== undefined) } diff --git a/tests/a2a-lifecycle.test.ts b/tests/a2a-lifecycle.test.ts index 5b2ff5b..30fcbe8 100644 --- a/tests/a2a-lifecycle.test.ts +++ b/tests/a2a-lifecycle.test.ts @@ -305,8 +305,6 @@ describe('A2A lifecycle recovery and ownership', () => { const created = await taskStore.get('task-created-before-claim') expect(created?.status.state).toBe('submitted') expect(created?.metadata?.gatewaySubmission).toBeDefined() - await taskStore.delete('task-created-before-claim') - expect(await taskStore.get('task-created-before-claim')).toBeDefined() const submission = created?.metadata?.gatewaySubmission as { lease: { id: string; expiresAt: number } } diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index d40cacf..bf85c71 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -2,6 +2,7 @@ import { Hono } from 'hono' import { afterEach, describe, expect, it, vi } from 'vitest' import { InMemoryTaskStore } from '../src/a2a/task-store' +import { SqlTaskStore, type SqlAdapter } from '../src/a2a/task-store-sql' import { createAgentGateway } from '../src/middleware' import { dispatchSandboxStreamRich, requiredX402Amount } from '../src/dispatch' import { MemoryNonceStore, claimStoredNonce, type NonceStore } from '../src/nonce-store' @@ -10,6 +11,7 @@ import { MemoryPaymentRecoveryStore, type PaymentRecoveryRecord } from '../src/p import { recoverPayment } from '../src/payment-recovery-worker' import { verifyMpp } from '../src/verify' import type { AgentMeta, GatewayConfig, SandboxStreamEvent } from '../src/types' +import type { Task } from '../src/a2a/types' import type { MppConfig } from '../src/types' const operatorAddress = '0x1111111111111111111111111111111111111111' @@ -562,4 +564,195 @@ describe('PR #11 production regressions', () => { expect(discovery.status).toBe(200) expect(a2a.status).toBe(503) }) + + it('does not retain an abandoned submission marker as payment recovery', async () => { + vi.useFakeTimers() + const store = new InMemoryTaskStore(10) + const task: Task = { + kind: 'task', + id: 'expired-submission', + contextId: 'expired-context', + status: { state: 'submitted', timestamp: new Date().toISOString() }, + metadata: { + gatewaySubmission: { + version: 1, + lease: { id: 'submission-lease', expiresAt: Date.now() + 5 * 60 * 1000 }, + agentId: agent.id, + agentSlug: agent.slug, + requestId: 'submission-request', + consumerId: commitment, + }, + }, + } + await store.put(task) + await vi.advanceTimersByTimeAsync(11) + + expect(await store.get(task.id)).toBeUndefined() + }) + + it('keeps generated input-required message ids stable across retries', async () => { + const makeApp = (nonce: string) => { + const app = new Hono() + const taskStore = new InMemoryTaskStore() + app.route('/v1/agents', createAgentGateway({ + resolveAgent: async () => agent, + getSandbox: async () => sandbox([ + { type: 'input-required', data: { inputRequired: { prompt: 'Need one more detail' } } }, + { type: 'sandbox.usage', data: { usage: usage() } }, + ]), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => true, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + })) + return app.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader(nonce), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: nonce, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'stable-input-required', + contextId: 'stable-context', + messageId: 'stable-request', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + } + + const first = await makeApp('stable-1') + const second = await makeApp('stable-2') + const firstBody = await first.json() as { result?: { status?: { message?: { messageId?: string } } } } + const secondBody = await second.json() as { result?: { status?: { message?: { messageId?: string } } } } + + expect(firstBody.result?.status?.message?.messageId).toBe( + secondBody.result?.status?.message?.messageId, + ) + }) + + it('rejects oversized chunked A2A bodies without trusting Content-Length', async () => { + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig())) + const response = await app.fetch(new Request('http://gateway.test/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/get', + params: { id: 'oversized' }, + padding: 'x'.repeat(70_000), + }), + })) + + expect(response.status).toBe(413) + }) + + it('retries SQL task creation after removing an expired colliding row', async () => { + interface Row { + id: string + context_id: string + state: string + payload: string + updated_at: number + } + const rows = new Map() + const db: SqlAdapter = { + async exec(sql, params = []) { + const statement = sql.trim() + if (statement.startsWith('CREATE TABLE') || statement.startsWith('CREATE INDEX')) { + return { rowsAffected: 0 } + } + if (statement.startsWith('INSERT INTO')) { + const [id, contextId, state, payload, updatedAt] = params as [ + string, + string, + string, + string, + number, + ] + if (rows.has(id)) throw new Error('duplicate primary key') + rows.set(id, { id, context_id: contextId, state, payload, updated_at: updatedAt }) + return { rowsAffected: 1 } + } + if (statement.startsWith('DELETE FROM')) { + const [id, payload, updatedAt] = params as [string, string, number] + const row = rows.get(id) + if (!row || row.payload !== payload || row.updated_at !== updatedAt) { + return { rowsAffected: 0 } + } + rows.delete(id) + return { rowsAffected: 1 } + } + throw new Error(`unrecognized SQL: ${statement}`) + }, + async query(_sql: string, params: readonly unknown[] = []): Promise { + const row = rows.get(params[0] as string) + return (row ? [{ payload: row.payload, updated_at: row.updated_at }] : []) as TRow[] + }, + } + const store = new SqlTaskStore(db, { ttlMs: 10 }) + const oldTask: Task = { + kind: 'task', + id: 'sql-expired-task', + contextId: 'sql-context', + status: { state: 'submitted', timestamp: new Date().toISOString() }, + } + expect(await store.createIfAbsent(oldTask)).toBe(true) + rows.get(oldTask.id)!.updated_at = Date.now() - 100 + const replacement: Task = { + ...oldTask, + status: { state: 'failed', timestamp: new Date().toISOString() }, + } + + expect(await store.createIfAbsent(replacement)).toBe(true) + expect((await store.get(replacement.id))?.status.state).toBe('failed') + }) + + it('does not follow push notification redirects', async () => { + const { deliverPushNotifications } = await import('../src/a2a/push-notifications') + const pushStore = { + async list() { + return [{ id: 'redirect', url: 'https://hook.example/redirect' }] + }, + } + const fetcher = vi.fn(async (_url: string, _init: RequestInit) => ( + new Response(null, { status: 302, headers: { Location: 'http://169.254.169.254/' } }) + )) + const results = await deliverPushNotifications({ + task: { + kind: 'task', + id: 'redirect-task', + contextId: 'redirect-context', + status: { state: 'completed', timestamp: new Date().toISOString() }, + }, + store: { + list: pushStore.list, + set: async () => undefined, + get: async () => undefined, + delete: async () => undefined, + }, + webhookSecret: undefined, + fetcher: fetcher as unknown as typeof fetch, + }) + + expect(results[0]).toMatchObject({ ok: false, status: 302 }) + expect(results[0]?.error).toContain('redirect') + expect(fetcher.mock.calls[0]?.[1]).toMatchObject({ redirect: 'manual' }) + }) }) From 6381dd9ec95bb9dba1d7c144844ea7acd4fe4f84 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 03:36:48 -0600 Subject: [PATCH 25/32] fix(gateway): close live cancellation and recovery races --- README.md | 1 + docs/a2a-long-horizon.md | 2 + src/a2a/handler.ts | 52 ++++++++------- src/a2a/push-notifications.ts | 91 +++++++++++++++++++++++++- src/dispatch.ts | 15 ++++- src/index.ts | 1 + src/payment-recovery-worker.ts | 41 +++++++----- src/types.ts | 5 ++ tests/payment-recovery.test.ts | 101 +++++++++++++++++++++++++++++ tests/pr11-regressions.test.ts | 114 ++++++++++++++++++++++++++++++++- 10 files changed, 377 insertions(+), 46 deletions(-) diff --git a/README.md b/README.md index d4cccb2..27b0de6 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,7 @@ The short-lived `gatewaySubmission` marker is not a payment recovery record and The bundled memory and SQL stores enforce this rule even after the normal task TTL. Push destinations must use HTTPS without URL credentials. Push delivery does not follow redirects. +Production push delivery also requires `a2a.pushUrlValidator` to reject private DNS destinations. Tasks created before this release have no recorded origin and fail closed; migrate them with a verified owner binding or let them expire. ## Tier diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index f1df802..79e9cfb 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -49,6 +49,8 @@ Explicit `x402.demoMode` permits these methods without the hook for local tests This is a fail-closed upgrade boundary. Tasks stored before this release do not have a `gatewayOrigin` binding and cannot pass the default ownership check. Migrate those records with a verified owner binding, or allow them to expire before enabling production control methods. +Push registration rejects reserved IP literals and private hostname suffixes. +Production push delivery also requires `a2a.pushUrlValidator`, which should apply the deployment's DNS-aware private-network policy. ### D1 (Cloudflare Workers) diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 4f4e8f1..6fb559b 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -44,12 +44,12 @@ import { buildAgentCard } from './agent-card' import { claimTaskExecution, clearTaskExecution, - hasActiveTaskExecution, renewTaskExecution, } from './execution-fence' import { fail, ok, parseEnvelope } from './jsonrpc' import { deliverPushNotifications, + validatePushNotificationUrl, type PushNotificationStore, type TaskPushNotificationConfig, } from './push-notifications' @@ -180,11 +180,6 @@ class CancelRegistry { return this.finalizing.has(taskId) } - hasController(taskId: string): boolean { - const controller = this.controllers.get(taskId) - return controller !== undefined && !controller.signal.aborted - } - cancel(taskId: string): boolean { if (this.finalizing.has(taskId)) return false const c = this.controllers.get(taskId) @@ -922,8 +917,7 @@ async function handleTasksCancel( if ( isTaskFinalizing(task) || - cancels.isFinalizing(task.id) || - (hasActiveTaskExecution(task) && !cancels.hasController(task.id)) + cancels.isFinalizing(task.id) ) { return c.json( fail( @@ -933,10 +927,7 @@ async function handleTasksCancel( ), ) } - const canceled: Task = { - ...task, - status: { state: 'canceled', timestamp: nowIso() }, - } + const canceled = withStatus(task, 'canceled') const transitioned = await compareAndSetTask(deps.taskStore, task, canceled) if (!transitioned) { const current = await deps.taskStore.get(task.id) @@ -1047,8 +1038,30 @@ async function handlePushSet( } const accessError = await authorizeTaskAccess(c, req, task, deps) if (accessError) return accessError - if (!isHttpsUrl(params.pushNotificationConfig.url)) { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url must use https')) + const pushUrl = validatePushNotificationUrl(params.pushNotificationConfig.url) + if (!pushUrl) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url is not a safe HTTPS destination'), + ) + } + const urlValidator = deps.config.a2a?.pushUrlValidator + if (!deps.config.x402.demoMode && !urlValidator) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'production push URL validation is not configured'), + ) + } + let allowedByHostPolicy = true + try { + if (urlValidator) allowedByHostPolicy = await urlValidator(pushUrl) + } catch (error) { + allowedByHostPolicy = false + console.error( + `[a2a] push URL policy failed for task ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + } + if (!allowedByHostPolicy) { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url was rejected')) } await deps.pushStore.set(params.taskId, params.pushNotificationConfig) const stored = await deps.pushStore.get(params.taskId, params.pushNotificationConfig.id) @@ -1597,15 +1610,6 @@ async function authorizeTaskAccess( return c.json(fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task access denied'), 403) } -function isHttpsUrl(value: string): boolean { - try { - const url = new URL(value) - return url.protocol === 'https:' && url.username === '' && url.password === '' - } catch { - return false - } -} - function bindRequestAbort(requestSignal: AbortSignal, controller: AbortController): () => void { const abort = () => controller.abort() if (requestSignal.aborted) abort() @@ -2523,6 +2527,8 @@ async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise store: deps.pushStore, webhookSecret: deps.config.a2a?.webhookSecret, fetcher: deps.config.a2a?.pushFetcher, + urlValidator: deps.config.a2a?.pushUrlValidator, + requireUrlValidator: !deps.config.x402.demoMode, onDelivery: (result) => { if (!result.ok) { deps.state.obs?.onStreamError?.( diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index ff3e97c..d9e2ca8 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -90,6 +90,83 @@ export interface PushNotificationStore { delete(taskId: string, configId: string): Promise } +/** + * Validate a push destination before the gateway sends task data to it. + * + * The default policy rejects URL credentials, non-HTTPS schemes, IP literals + * in reserved ranges, and common private hostnames. Production deployments + * should also provide `GatewayConfig.a2a.pushUrlValidator` for DNS policy. + */ +export function validatePushNotificationUrl(value: string): URL | undefined { + let url: URL + try { + url = new URL(value) + } catch { + return undefined + } + if ( + url.protocol !== 'https:' || + url.username !== '' || + url.password !== '' || + isPrivatePushHostname(url.hostname) + ) { + return undefined + } + return url +} + +function isPrivatePushHostname(value: string): boolean { + const hostname = value.toLowerCase().replace(/\.$/, '') + const ipv4 = parseIpv4(hostname) + if (ipv4) { + const [first, second] = ipv4 + return first === 0 || + first === 10 || + first === 127 || + (first === 100 && second >= 64 && second <= 127) || + (first === 169 && second === 254) || + (first === 172 && second >= 16 && second <= 31) || + (first === 192 && second === 0) || + (first === 192 && second === 168) || + (first === 198 && (second === 18 || second === 19)) || + (first === 203 && second === 0) || + first >= 224 + } + const ipv6 = hostname.replace(/^\[|\]$/g, '') + if ( + ipv6 === '::' || + ipv6 === '::1' || + ipv6.startsWith('fc') || + ipv6.startsWith('fd') || + ipv6.startsWith('fe8') || + ipv6.startsWith('fe9') || + ipv6.startsWith('fea') || + ipv6.startsWith('feb') || + ipv6.startsWith('::ffff:127.') || + ipv6.startsWith('::ffff:10.') || + ipv6.startsWith('::ffff:192.168.') || + ipv6.startsWith('::ffff:169.254.') + ) return true + return hostname === 'localhost' || + hostname === 'localhost.localdomain' || + hostname === 'metadata' || + hostname === 'metadata.google.internal' || + hostname.endsWith('.localhost') || + hostname.endsWith('.local') || + hostname.endsWith('.internal') || + hostname.endsWith('.intranet') || + hostname.endsWith('.lan') || + hostname.endsWith('.home') +} + +function parseIpv4(value: string): [number, number, number, number] | undefined { + const parts = value.split('.') + if (parts.length !== 4 || parts.some((part) => !/^\d{1,3}$/.test(part))) return undefined + const numbers = parts.map(Number) + if (numbers.some((part) => part > 255)) return undefined + return numbers as [number, number, number, number] +} + export class InMemoryPushNotificationStore implements PushNotificationStore { private readonly byTask = new Map>() @@ -222,6 +299,10 @@ export async function deliverPushNotifications(args: { webhookSecret: string | undefined /** Inject for tests. Defaults to global `fetch`. */ fetcher?: typeof fetch + /** Optional DNS-aware host policy for production deployments. */ + urlValidator?: (url: URL) => boolean | Promise + /** Require `urlValidator` before sending from a production gateway. */ + requireUrlValidator?: boolean /** Optional callback so the gateway's observer can log delivery outcomes. */ onDelivery?: (result: PushDeliveryResult) => void }): Promise { @@ -249,8 +330,14 @@ export async function deliverPushNotifications(args: { let result: PushDeliveryResult try { const url = new URL(config.url) - if (url.protocol !== 'https:' || url.username !== '' || url.password !== '') { - throw new Error('push notification URL must use https without credentials') + if (!validatePushNotificationUrl(config.url)) { + throw new Error('push notification URL is not a safe HTTPS destination') + } + if (args.requireUrlValidator && !args.urlValidator) { + throw new Error('push notification URL validation is not configured') + } + if (args.urlValidator && !await args.urlValidator(url)) { + throw new Error('push notification URL was rejected by host policy') } const res = await fetcher(config.url, { method: 'POST', diff --git a/src/dispatch.ts b/src/dispatch.ts index 767445c..c61d1c6 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -1341,8 +1341,11 @@ export async function* dispatchSandboxStreamRich( yield { kind: 'usage', usage } } finally { if (heartbeatTimer !== undefined) clearInterval(heartbeatTimer) + const pendingHeartbeat = heartbeatInFlight + if (pendingHeartbeat) await pendingHeartbeat signal?.removeEventListener('abort', forwardAbort) if (iterator) await closeSandboxIterator(iterator) + if (heartbeatError !== undefined && !signal?.aborted) throw heartbeatError } } @@ -1532,15 +1535,23 @@ async function markRecoverySettling( if (!recovery || !authz.paymentRecoveryId) return const fenceId = requirePaymentRecoveryFence(authz) await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => { - return { + const next: PaymentRecoveryRecord = { ...record, state: 'settling', payment: recoveryTarget(authz, record.payment), workStarted: true, - usage, settlementBasis, nextAttemptAt: Date.now(), } + // A quoted-ceiling settlement has no provider receipt. Keep the durable + // basis and original amount, then rebuild the synthetic accounting input + // on each retry instead of persisting a lossy floating-point surrogate. + if (settlementBasis !== 'quoted-ceiling' || record.usage !== undefined) { + next.usage = usage + } else { + delete next.usage + } + return next }) } diff --git a/src/index.ts b/src/index.ts index c412991..514cb64 100644 --- a/src/index.ts +++ b/src/index.ts @@ -133,6 +133,7 @@ export { export { deliverPushNotifications, InMemoryPushNotificationStore, + validatePushNotificationUrl, type PushDeliveryResult, type PushNotificationAuthentication, type PushNotificationConfig, diff --git a/src/payment-recovery-worker.ts b/src/payment-recovery-worker.ts index f88d852..670b5e7 100644 --- a/src/payment-recovery-worker.ts +++ b/src/payment-recovery-worker.ts @@ -81,7 +81,7 @@ export async function recoverPayment( const current = await recovery.store.get(recoveryId) if (!current || current.state === 'reconciled') return current if (!options.force && current.nextAttemptAt > scanNow) return current - const now = recoveryNow(options) + const now = options.now ?? recoveryNow(options) const leased = await acquireLease( current.id, @@ -99,7 +99,7 @@ export async function recoverPayment( await reconcileLeased(ready, config, now) } catch (error) { const message = error instanceof Error ? error.message : String(error) - const failedAt = recoveryNow(options) + const failedAt = options.now ?? recoveryNow(options) try { await updateOwnedPaymentRecovery(recovery.store, recoveryId, fenceId, (record) => ({ ...record, @@ -126,7 +126,7 @@ async function reconcileLeased( now: number, ): Promise { if (record.payment.kind === 'mpp-charge' && !record.payment.operation) { - record = await recoverUnknownMppCharge(record, config) + record = await recoverUnknownMppCharge(record, config, now) } if (record.state === 'reconciled') return @@ -140,25 +140,25 @@ async function reconcileLeased( if (recovered.operationId !== record.payment.operationId) { throw new Error('x402 recovery operation id mismatch') } - await completeRecord(record, config) + await completeRecord(record, config, now) return } if (recovered.state !== 'released' && recovered.state !== 'reclaimed') { throw new Error(`ambiguous x402 claim recovered in state ${recovered.state}`) } - await completeRecord(record, config) + await completeRecord(record, config, now) return } throw new Error('MPP charge confirmation remains unresolved') } if (record.state === 'claimed' && !record.workStarted) { - await releaseRecoveredPayment(record, config, 'request ended before sandbox execution') + await releaseRecoveredPayment(record, config, 'request ended before sandbox execution', now) return } if (record.state === 'releasing') { - await releaseRecoveredPayment(record, config, record.reason ?? 'payment recovery release') + await releaseRecoveredPayment(record, config, record.reason ?? 'payment recovery release', now) return } @@ -192,6 +192,7 @@ async function reconcileLeased( async function recoverUnknownMppCharge( record: PaymentRecoveryRecord, config: GatewayConfig, + now: number, ): Promise { if (record.payment.kind !== 'mpp-charge') return record const method = record.payment.method @@ -206,7 +207,7 @@ async function recoverUnknownMppCharge( throw new Error('MPP recovery operation id mismatch') } if (result.state === 'not-found') { - await completeRecord(record, config) + await completeRecord(record, config, now) return requireRecord(record.id, config) } throw new Error('MPP charge confirmation is still pending') @@ -236,9 +237,10 @@ async function recoverUnknownMppCharge( : 'claimed', payment: { kind: 'mpp-charge', method, operationId, operation: result }, ...(result.state === 'released' - ? { reconciledAt: Date.now(), nextAttemptAt: Number.MAX_SAFE_INTEGER, lease: undefined } + ? { reconciledAt: now, nextAttemptAt: Number.MAX_SAFE_INTEGER, lease: undefined } : {}), }), + now, ) } @@ -246,6 +248,7 @@ async function releaseRecoveredPayment( record: PaymentRecoveryRecord, config: GatewayConfig, reason: string, + now: number, ): Promise { const authz = authorizedRequest(record) if (record.payment.kind === 'x402') { @@ -256,13 +259,13 @@ async function releaseRecoveredPayment( if (recovered.operationId !== record.payment.operationId) { throw new Error('x402 recovery operation id mismatch') } - await completeRecord(record, config) + await completeRecord(record, config, now) return } if (recovered.state !== 'released' && recovered.state !== 'reclaimed') { throw new Error(`x402 release recovered in state ${recovered.state}`) } - await completeRecord(record, config) + await completeRecord(record, config, now) return } authz.paymentOperation = deserializePaymentOperation(record.payment.operation) @@ -288,8 +291,11 @@ async function settleRecoveredPayment( authz.mppChargeOperation = record.payment.operation } - const fallback = !record.usage + const fallback = record.settlementBasis === 'quoted-ceiling' || !record.usage const usage = record.usage ?? quotedCeilingUsage(record) + const settlementBasis = fallback + ? 'quoted-ceiling' + : record.settlementBasis ?? 'usage-receipt' await settleAndRecord( recoveryAgent(record), authz, @@ -298,8 +304,8 @@ async function settleRecoveredPayment( config.observer, { usageAlreadyRecorded: record.usageRecorded, - settlementBasis: fallback ? 'quoted-ceiling' : record.settlementBasis ?? 'usage-receipt', - ...(fallback && record.payment.kind === 'x402' + settlementBasis, + ...(settlementBasis === 'quoted-ceiling' && record.payment.kind === 'x402' ? { paymentAmount: BigInt(record.attribution.requiredAmount) } : {}), }, @@ -387,8 +393,11 @@ async function acquireLease( throw new Error(`payment recovery record ${id} changed too many times`) } -async function completeRecord(record: PaymentRecoveryRecord, config: GatewayConfig): Promise { - const now = Date.now() +async function completeRecord( + record: PaymentRecoveryRecord, + config: GatewayConfig, + now: number, +): Promise { await updateOwnedPaymentRecovery( config.paymentRecovery!.store, record.id, diff --git a/src/types.ts b/src/types.ts index b7ff4aa..1442c26 100644 --- a/src/types.ts +++ b/src/types.ts @@ -423,6 +423,11 @@ export interface GatewayConfig { * `fetch`. Override for tests or to wire a queue-backed sender. */ pushFetcher?: typeof fetch + /** + * DNS-aware policy for push destinations. Required when production + * push delivery is enabled so private DNS names cannot receive task data. + */ + pushUrlValidator?: (url: URL) => boolean | Promise } } diff --git a/tests/payment-recovery.test.ts b/tests/payment-recovery.test.ts index d1c53a3..573982b 100644 --- a/tests/payment-recovery.test.ts +++ b/tests/payment-recovery.test.ts @@ -15,6 +15,7 @@ import { } from '../src/payment-operations' import { MemoryPaymentRecoveryStore, + serializePaymentOperation, type PaymentRecoveryRecord, } from '../src/payment-recovery' import { recoverPayment, recoverPayments } from '../src/payment-recovery-worker' @@ -968,6 +969,106 @@ describe('durable OpenAI recovery', () => { expect(usage.every((event) => event.settlementBasis === 'quoted-ceiling')).toBe(true) expect(usage.every((event) => event.inputTokens === 0 && event.outputTokens === 0)).toBe(true) }) + + it('reuses the exact quoted ceiling after attribution fails post-settlement', async () => { + const requiredAmount = 1234567890123456789n + const recoveryStore = new MemoryPaymentRecoveryStore() + const settlements: PaymentSettlementInput[] = [] + const operations = new MemoryPaymentOperations({ + onSettle: async (_operation, input) => { settlements.push(input) }, + onReclaim: async () => undefined, + }) + let usageAttempts = 0 + const config: GatewayConfig = { + resolveAgent: async () => agent, + getSandbox: async () => ({ async *streamPrompt() {} }), + recordUsage: async () => { + usageAttempts += 1 + if (usageAttempts === 1) throw new Error('usage store unavailable') + }, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore: new MemoryNonceStore(), + paymentRecovery: { store: recoveryStore, retryDelayMs: 1 }, + } + const executionBudget = { + maxInputTokens: 1, + maxOutputTokens: 1, + maxReasoningTokens: 0, + maxToolTokens: 0, + maxToolCalls: 0, + maxProviderCostUsd: 1, + } + const claimed = await operations.claimPayment( + { + commitment, + nonce: '600', + amount: requiredAmount.toString(), + expiry: String(Math.floor(Date.now() / 1000) + 600), + }, + { + requestId: 'quoted-retry-request', + agentId: agent.id, + requiredAmount: 1n, + maxOutputTokens: 1, + executionBudget, + }, + ) + const executing = await operations.beginPaymentExecution(claimed) + const retained = await operations.retainPayment(executing, 'test recovery') + const recoveryId = retained.operationId + const now = Date.now() + await recoveryStore.createIfAbsent({ + version: 1, + id: recoveryId, + revision: 0, + state: 'retained', + payment: { + kind: 'x402', + operationId: recoveryId, + operation: serializePaymentOperation(retained), + }, + attribution: { + requestId: 'quoted-retry-request', + agentId: agent.id, + agentSlug: agent.slug, + consumerId: commitment, + paymentMethod: 'x402', + startMs: now, + pricePerTokenUsd: agent.pricePerTokenUsd, + platformFeePercent: agent.platformFeePercent, + requiredAmount: requiredAmount.toString(), + currencyDecimals: 18, + maxOutputTokens: 1, + executionBudget, + }, + workStarted: true, + usageRecorded: false, + attempts: 0, + nextAttemptAt: now, + lease: undefined, + createdAt: now, + updatedAt: now, + }) + + await expect(recoverPayment(recoveryId, config, { force: true, now })).rejects.toThrow( + 'usage store unavailable', + ) + const settling = await recoveryStore.get(recoveryId) + expect(settling?.state).toBe('settling') + expect(settling?.settlementBasis).toBe('quoted-ceiling') + expect(settling?.usage).toBeUndefined() + + const recovered = await recoverPayment(recoveryId, config, { force: true, now: now + 1 }) + expect(recovered?.state).toBe('reconciled') + expect(settlements).toHaveLength(1) + expect(settlements[0]?.amount).toBe(requiredAmount) + }) }) describe('A2A recovery retention', () => { diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index bf85c71..ee2394a 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -250,6 +250,87 @@ describe('PR #11 production regressions', () => { expect((await taskStore.get('pr11-cancel-race'))?.status.state).toBe('canceled') }) + it('cancels an execution already fenced by another worker', async () => { + const taskStore = new InMemoryTaskStore() + let executionClaimed!: () => void + const executionReady = new Promise((resolve) => { executionClaimed = resolve }) + let releaseExecution!: () => void + const executionReleased = new Promise((resolve) => { releaseExecution = resolve }) + + const makeConfig = (worker: 'runner' | 'canceler'): GatewayConfig => ({ + resolveAgent: async () => agent, + getSandbox: async () => ({ + async *streamPrompt() { + if (worker === 'runner') { + executionClaimed() + await executionReleased + } + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 1, + authorizePayment: async () => true, + }, + nonceStore: new MemoryNonceStore(), + a2a: { taskStore, authorizeTaskAccess: async () => true }, + }) + + const runner = new Hono() + runner.route('/v1/agents', createAgentGateway(makeConfig('runner'))) + const canceler = new Hono() + canceler.route('/v1/agents', createAgentGateway(makeConfig('canceler'))) + + const running = runner.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9002'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'pr11-active-cancel-race', + contextId: 'pr11-active-cancel-context', + messageId: 'pr11-active-cancel-message', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + await executionReady + + const cancel = await canceler.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'pr11-active-cancel-race' }, + }), + }) + const cancelBody = await cancel.json() as { result?: { status?: { state?: string } } } + expect(cancel.status).toBe(200) + expect(cancelBody.result?.status?.state).toBe('canceled') + + releaseExecution() + const runningResponse = await running + const runningBody = await runningResponse.json() as { result?: { status?: { state?: string } } } + expect(runningBody.result?.status?.state).toBe('canceled') + expect((await taskStore.get('pr11-active-cancel-race'))?.metadata?.gatewayExecution) + .toBeUndefined() + }) + it('quotes retained A2A history before charging a continuation', async () => { const taskStore = new InMemoryTaskStore() let invocations = 0 @@ -407,9 +488,11 @@ describe('PR #11 production regressions', () => { }, }) - const recovered = await recoverPayment(id, config, { force: true }) + const recoveryNow = now + 10_000 + const recovered = await recoverPayment(id, config, { force: true, now: recoveryNow }) expect(recovered?.state).toBe('reconciled') + expect((await store.get(id))?.reconciledAt).toBe(recoveryNow) }) it('does not leave an abandoned working A2A task after payment recovery', async () => { @@ -706,7 +789,7 @@ describe('PR #11 production regressions', () => { return (row ? [{ payload: row.payload, updated_at: row.updated_at }] : []) as TRow[] }, } - const store = new SqlTaskStore(db, { ttlMs: 10 }) + const store = new SqlTaskStore(db, { ttlMs: 1_000 }) const oldTask: Task = { kind: 'task', id: 'sql-expired-task', @@ -714,7 +797,7 @@ describe('PR #11 production regressions', () => { status: { state: 'submitted', timestamp: new Date().toISOString() }, } expect(await store.createIfAbsent(oldTask)).toBe(true) - rows.get(oldTask.id)!.updated_at = Date.now() - 100 + rows.get(oldTask.id)!.updated_at = Date.now() - 10_000 const replacement: Task = { ...oldTask, status: { state: 'failed', timestamp: new Date().toISOString() }, @@ -755,4 +838,29 @@ describe('PR #11 production regressions', () => { expect(results[0]?.error).toContain('redirect') expect(fetcher.mock.calls[0]?.[1]).toMatchObject({ redirect: 'manual' }) }) + + it('rejects direct private push destinations before fetch', async () => { + const { deliverPushNotifications } = await import('../src/a2a/push-notifications') + const fetcher = vi.fn(async () => new Response('unexpected', { status: 200 })) + const results = await deliverPushNotifications({ + task: { + kind: 'task', + id: 'private-task', + contextId: 'private-context', + status: { state: 'completed', timestamp: new Date().toISOString() }, + }, + store: { + list: async () => [{ id: 'private', url: 'https://127.0.0.1/internal' }], + set: async () => undefined, + get: async () => undefined, + delete: async () => undefined, + }, + webhookSecret: undefined, + fetcher: fetcher as unknown as typeof fetch, + }) + + expect(results[0]?.ok).toBe(false) + expect(results[0]?.error).toContain('safe HTTPS destination') + expect(fetcher).not.toHaveBeenCalled() + }) }) From a9834ca213d3ba55aad9e15cd33eda62a3b5460c Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 04:07:04 -0600 Subject: [PATCH 26/32] fix(gateway): close finalization recovery races --- README.md | 2 ++ src/a2a/handler.ts | 18 ++++++++++++++++- src/a2a/push-notifications.ts | 7 +++---- src/dispatch.ts | 19 ++++++++++++++---- src/payment-operations.ts | 15 +++++++++++++++ tests/a2a-atomicity.test.ts | 35 ++++++++++++++++++++++++++++++++++ tests/pr11-regressions.test.ts | 33 ++++++++++++++++++++++++++++++++ 7 files changed, 120 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 27b0de6..988c166 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,8 @@ Run `recoverPayments(config)` from a private scheduled worker. Every live request and worker uses a unique durable fence token. A stale request or worker cannot update a row after another worker takes its lease. Provider settlement, recovery, and release methods must still use the operation ID idempotently. +`paymentOperations.getPaymentOperation` must read the authoritative provider state by operation ID without changing it. +This read is required for recovery of older A2A finalization records that predate the shared payment outbox. The operation store owns claim, execution start, receipt retention, partial settle, release, and expiry reclaim. An executing or retained operation cannot expire into a refund. A retained operation settles from its receipt when one exists. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 6fb559b..c2da736 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -2186,6 +2186,21 @@ async function recoverFinalizationIfNeeded( paymentOperation = deserializePaymentOperation(renewed.paymentOperation) } + let paymentAlreadySettled = false + if (paymentOperation && deps.config.x402.paymentOperations) { + const currentOperation = await deps.config.x402.paymentOperations.getPaymentOperation( + paymentOperation.operationId, + ) + if (currentOperation.state === 'not-found') { + throw new Error('A2A payment operation disappeared during finalization recovery') + } + if (currentOperation.operationId !== paymentOperation.operationId) { + throw new Error('A2A payment operation recovery returned a different operation') + } + paymentOperation = currentOperation + paymentAlreadySettled = currentOperation.state === 'settled' + } + const authz: AuthorizedRequest = { agent, consumerId: renewed.consumerId, @@ -2215,6 +2230,7 @@ async function recoverFinalizationIfNeeded( deps.state.obs, { usageAlreadyRecorded: renewed.usageRecorded === true, + paymentAlreadySettled, onUsageRecorded: async () => { usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) }, @@ -2468,7 +2484,7 @@ function agentMessage(task: Task, text: string): Message { kind: 'message', role: 'agent', parts: [{ kind: 'text', text }], - messageId: `${task.id}-status-${task.status.state}-${stableMessageDigest(text)}`, + messageId: `${task.id}-input-required-${stableMessageDigest(text)}`, taskId: task.id, contextId: task.contextId, } diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index d9e2ca8..7b02055 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -142,10 +142,9 @@ function isPrivatePushHostname(value: string): boolean { ipv6.startsWith('fe9') || ipv6.startsWith('fea') || ipv6.startsWith('feb') || - ipv6.startsWith('::ffff:127.') || - ipv6.startsWith('::ffff:10.') || - ipv6.startsWith('::ffff:192.168.') || - ipv6.startsWith('::ffff:169.254.') + // WHATWG URL normalizes dotted IPv4-mapped IPv6 literals to hexadecimal. + // Reject the whole mapped range instead of matching only dotted forms. + ipv6.startsWith('::ffff:') ) return true return hostname === 'localhost' || hostname === 'localhost.localdomain' || diff --git a/src/dispatch.ts b/src/dispatch.ts index c61d1c6..ae88da1 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -1408,6 +1408,8 @@ async function closeSandboxIterator(iterator: AsyncIterator) export interface SettleAndRecordOptions { /** Skip attribution after a durable finalization marker confirms it ran. */ usageAlreadyRecorded?: boolean + /** Skip provider settlement only after an authoritative read found it settled. */ + paymentAlreadySettled?: boolean /** Persist the caller's recovery marker after attribution succeeds. */ onUsageRecorded?: () => Promise /** Recovery uses the original quoted ceiling when no receipt arrives. */ @@ -1467,10 +1469,19 @@ export async function settleAndRecord( config.x402.currencyDecimals, usage.providerCostUsd, ) - authz.paymentOperation = await config.x402.paymentOperations.settlePayment( - authz.paymentOperation, - { amount, totalCostUsd: totalCost, usage, basis: settlementBasis }, - ) + if (options.paymentAlreadySettled) { + if ( + authz.paymentOperation.state !== 'settled' || + authz.paymentOperation.settledAmount !== amount + ) { + throw new Error('authoritative payment state does not match finalization') + } + } else { + authz.paymentOperation = await config.x402.paymentOperations.settlePayment( + authz.paymentOperation, + { amount, totalCostUsd: totalCost, usage, basis: settlementBasis }, + ) + } // Durable settlement happens first. If attribution storage is // unavailable, recovery must never refund delivered work. if (!options.usageAlreadyRecorded) { diff --git a/src/payment-operations.ts b/src/payment-operations.ts index 01d374d..bc94b88 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -83,6 +83,12 @@ export interface PaymentOperations { operation: PaymentOperation, input: PaymentSettlementInput, ): Promise + /** + * Read the authoritative durable operation without changing its state. + * Recovery uses this to avoid repeating a provider settlement after the + * provider committed but the task finalization write lost its acknowledgement. + */ + getPaymentOperation(operationId: string): Promise /** * Release an unused authorization. * Repeated calls must recover an ambiguous acknowledgement by operationId. @@ -427,6 +433,15 @@ export class MemoryPaymentOperations implements PaymentOperations { return operation ? { ...operation } : undefined } + async getPaymentOperation(operationId: string): Promise { + const operation = this.get(operationId) + return operation ?? { + protocolVersion: PAYMENT_PROTOCOL_VERSION, + operationId, + state: 'not-found', + } + } + private requireCurrent(operation: PaymentOperation): PaymentOperation { const current = this.operations.get(operation.operationId) if (!current) throw new Error('payment operation was not found') diff --git a/tests/a2a-atomicity.test.ts b/tests/a2a-atomicity.test.ts index 156c0cf..3abc7b7 100644 --- a/tests/a2a-atomicity.test.ts +++ b/tests/a2a-atomicity.test.ts @@ -312,6 +312,41 @@ describe('A2A task atomicity and restart recovery', () => { expect(operations.get(executing.operationId)?.state).toBe('settled') expect(counters.records).toBe(1) expect(counters.settlements).toBe(1) + + // Simulate the provider commit succeeding immediately before the task + // finalization write was lost. The durable task still contains the + // pre-settlement operation snapshot, so recovery must inspect the current + // provider state instead of charging the operation again. + await taskStore.put({ + ...task, + id: 'task-settled-finalization', + contextId: 'ctx-settled-finalization', + metadata: { + ...task.metadata, + gatewayFinalizing: { + ...(task.metadata?.gatewayFinalizing as Record), + lease: { id: 'settled-replay-lease', expiresAt: Date.now() - 1 }, + }, + }, + }) + const replayed = await app.request(`/v1/agents/${agent.slug}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/get', + params: { id: 'task-settled-finalization' }, + }), + }) + const replayedBody = await replayed.json() as { result?: Task; error?: unknown } + + expect(replayedBody.error).toBeUndefined() + expect(replayedBody.result?.status.state).toBe('completed') + expect((await taskStore.get('task-settled-finalization'))?.metadata?.gatewayFinalizing) + .toBeUndefined() + expect(counters.records).toBe(2) + expect(counters.settlements).toBe(1) }) it('keeps a settling payment recoverable after a lost settlement acknowledgement', async () => { diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index ee2394a..7a0cdb8 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -863,4 +863,37 @@ describe('PR #11 production regressions', () => { expect(results[0]?.error).toContain('safe HTTPS destination') expect(fetcher).not.toHaveBeenCalled() }) + + it('rejects IPv4-mapped IPv6 private push destinations before fetch', async () => { + const { deliverPushNotifications } = await import('../src/a2a/push-notifications') + const fetcher = vi.fn(async () => new Response('unexpected', { status: 200 })) + const urls = [ + 'https://[::ffff:169.254.169.254]/metadata', + 'https://[::ffff:a9fe:a9fe]/metadata', + ] + + for (const [index, url] of urls.entries()) { + const results = await deliverPushNotifications({ + task: { + kind: 'task', + id: `mapped-private-task-${index}`, + contextId: 'private-context', + status: { state: 'completed', timestamp: new Date().toISOString() }, + }, + store: { + list: async () => [{ id: `mapped-private-${index}`, url }], + set: async () => undefined, + get: async () => undefined, + delete: async () => undefined, + }, + webhookSecret: undefined, + fetcher: fetcher as unknown as typeof fetch, + }) + + expect(results[0]?.ok).toBe(false) + expect(results[0]?.error).toContain('safe HTTPS destination') + } + + expect(fetcher).not.toHaveBeenCalled() + }) }) From 43897380c5c7034b59f5ac5d06885679d5a9bb94 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 04:34:58 -0600 Subject: [PATCH 27/32] fix(a2a): recover stale task leases safely --- README.md | 2 + docs/a2a-long-horizon.md | 2 + src/a2a/execution-fence.ts | 6 ++ src/a2a/handler.ts | 78 ++++++++++----- src/a2a/push-notifications.ts | 4 +- tests/pr11-regressions.test.ts | 176 ++++++++++++++++++++++++++++++++- 6 files changed, 244 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index 988c166..87d40dd 100644 --- a/README.md +++ b/README.md @@ -118,6 +118,8 @@ Push destinations must use HTTPS without URL credentials. Push delivery does not follow redirects. Production push delivery also requires `a2a.pushUrlValidator` to reject private DNS destinations. Tasks created before this release have no recorded origin and fail closed; migrate them with a verified owner binding or let them expire. +The payment claim keeps its submission lease until the atomic submitted-to-working transition. +An expired execution lease fails the working task and preserves its payment recovery markers. ## Tier diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 79e9cfb..9b48272 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -37,6 +37,8 @@ If cancellation must release an unused durable operation, the shared outbox ente An ambiguous release acknowledgement remains in that outbox and is retried by operation ID. Usage attribution must atomically upsert by `requestId`. The finalization record stores whether attribution was acknowledged, so recovery does not repeat an acknowledged usage event. +The payment claim keeps its submission lease until the task changes to `working` in one compare-and-set operation. +An expired execution lease fails the task while leaving payment recovery metadata available for reconciliation. If a legacy record is malformed or has no recoverable operation, the gateway expires the task as failed. If work exists without a receipt, the outbox settles the original quoted ceiling after its configured timeout. The task-store TTL cannot delete a task while any payment recovery marker remains. diff --git a/src/a2a/execution-fence.ts b/src/a2a/execution-fence.ts index d12c3fe..3fdbe79 100644 --- a/src/a2a/execution-fence.ts +++ b/src/a2a/execution-fence.ts @@ -68,6 +68,12 @@ export function hasActiveTaskExecution(task: Task, now = Date.now()): boolean { return marker !== undefined && marker.lease.expiresAt > now } +/** A working task with this marker has lost its execution owner. */ +export function hasExpiredTaskExecution(task: Task, now = Date.now()): boolean { + const marker = readTaskExecution(task) + return marker !== undefined && !hasActiveTaskExecution(task, now) +} + /** Remove the marker when the task reaches a terminal or paused state. */ export function clearTaskExecution(task: Task): Task { if (!task.metadata || !(TASK_EXECUTION_METADATA_KEY in task.metadata)) return task diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index c2da736..50b0658 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -44,6 +44,7 @@ import { buildAgentCard } from './agent-card' import { claimTaskExecution, clearTaskExecution, + hasExpiredTaskExecution, renewTaskExecution, } from './execution-fence' import { fail, ok, parseEnvelope } from './jsonrpc' @@ -301,10 +302,14 @@ async function executeMessageSend( signal: AbortSignal, ): Promise { if (isTerminal(task.status.state)) return c.json(ok(req.id, task)) + const taskWithoutSubmission = clearTaskSubmission(task) let workingTask: Task = task.status.state === 'working' - ? task - : { ...task, status: { state: 'working', timestamp: nowIso() } } - if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { + ? taskWithoutSubmission + : { ...taskWithoutSubmission, status: { state: 'working', timestamp: nowIso() } } + if ( + JSON.stringify(task) !== JSON.stringify(workingTask) && + !await compareAndSetTask(deps.taskStore, task, workingTask) + ) { await releaseTaskPayment(authz, task, deps, 'A2A task changed before execution started', false) if (signal.aborted) { const canceled = await deps.taskStore.get(task.id) @@ -927,26 +932,34 @@ async function handleTasksCancel( ), ) } - const canceled = withStatus(task, 'canceled') - const transitioned = await compareAndSetTask(deps.taskStore, task, canceled) - if (!transitioned) { - const current = await deps.taskStore.get(task.id) - if (current && (isTerminal(current.status.state) || isTaskFinalizing(current))) { + let stillActive = false + let candidate = task + for (let attempt = 0; attempt < 8; attempt += 1) { + if (isTerminal(candidate.status.state)) { return c.json( fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' changed before cancellation`), ) } - return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'task changed before cancellation')) - } - const stillActive = cancels.cancel(task.id) - - // If a stream was active, it'll observe the abort and emit its own final - // status-update AND fire its own push delivery; the dispatcher only fires - // push when the cancel races to terminal state with no active streamer. - if (!stillActive) { - await maybeDeliverPush(canceled, deps) + if (isTaskFinalizing(candidate)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' is being finalized`), + ) + } + const canceled = withStatus(candidate, 'canceled') + if (await compareAndSetTask(deps.taskStore, candidate, canceled)) { + stillActive = cancels.cancel(task.id) + // If a stream was active, it observes the abort and emits its own final + // status update and push delivery. Otherwise this handler owns delivery. + if (!stillActive) await maybeDeliverPush(canceled, deps) + return c.json(ok(req.id, canceled)) + } + const current = await deps.taskStore.get(task.id) + if (!current) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${task.id}' not found`)) + } + candidate = current } - return c.json(ok(req.id, canceled)) + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'task changed before cancellation')) } // ── tasks/resubscribe ───────────────────────────────────────────────────── @@ -1332,8 +1345,6 @@ async function claimTaskPayment( return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment authorization failed')) } - const submissionClear = await clearTaskSubmissionMarker(deps.taskStore, paymentTask) - if (submissionClear.applied) paymentTask = submissionClear.task const current = await deps.taskStore.get(task.id) if (current && JSON.stringify(current) === JSON.stringify(paymentTask)) { return paymentTask @@ -2274,7 +2285,30 @@ async function recoverTaskIfNeeded( const released = await recoverPaymentReleaseIfNeeded(task, deps) const finalized = await recoverFinalizationIfNeeded(released, deps, requestedAgentSlug) const paymentRecovered = await recoverPaymentMarkerIfNeeded(finalized, deps) - return recoverSubmissionIfNeeded(paymentRecovered, deps) + const submissionRecovered = await recoverSubmissionIfNeeded(paymentRecovered, deps) + return recoverExpiredExecutionIfNeeded(submissionRecovered, deps) +} + +async function recoverExpiredExecutionIfNeeded(task: Task, deps: A2AHandlerDeps): Promise { + if ( + task.status.state !== 'working' || + isTaskFinalizing(task) || + !hasExpiredTaskExecution(task) + ) return task + const failed: Task = { + ...withStatus(task, 'failed'), + metadata: { + ...(clearTaskExecution(task).metadata ?? {}), + [EXECUTION_RECOVERY_METADATA_KEY]: { + error: 'A2A execution lease expired before a task result was stored', + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await maybeDeliverPush(failed, deps) + return failed + } + return await deps.taskStore.get(task.id) ?? task } async function recoverSubmissionIfNeeded(task: Task, deps: A2AHandlerDeps): Promise { @@ -2333,7 +2367,7 @@ async function clearReconciledPaymentRecoveryMarker( const record = await deps.config.paymentRecovery.store.get(marker.id) if (record?.state !== 'reconciled') return task const cleared = clearPaymentRecoveryMarker(task) - if (cleared.status.state === 'working') { + if (cleared.status.state === 'working' || cleared.status.state === 'submitted') { const failed: Task = { ...withStatus(cleared, 'failed'), metadata: { diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index 7b02055..a16c72d 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -144,7 +144,9 @@ function isPrivatePushHostname(value: string): boolean { ipv6.startsWith('feb') || // WHATWG URL normalizes dotted IPv4-mapped IPv6 literals to hexadecimal. // Reject the whole mapped range instead of matching only dotted forms. - ipv6.startsWith('::ffff:') + ipv6.startsWith('::ffff:') || + // IPv4-compatible IPv6 literals can also embed loopback/private IPv4. + (ipv6.startsWith('::') && ipv6 !== '::1') ) return true return hostname === 'localhost' || hostname === 'localhost.localdomain' || diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index 7a0cdb8..7277af0 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -1,7 +1,7 @@ import { Hono } from 'hono' import { afterEach, describe, expect, it, vi } from 'vitest' -import { InMemoryTaskStore } from '../src/a2a/task-store' +import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { SqlTaskStore, type SqlAdapter } from '../src/a2a/task-store-sql' import { createAgentGateway } from '../src/middleware' import { dispatchSandboxStreamRich, requiredX402Amount } from '../src/dispatch' @@ -331,6 +331,139 @@ describe('PR #11 production regressions', () => { .toBeUndefined() }) + it('keeps the submission lease until the task enters working state', async () => { + const innerStore = new InMemoryTaskStore() + let transitionStarted!: () => void + const transitionReady = new Promise((resolve) => { transitionStarted = resolve }) + let releaseTransition!: () => void + const transitionReleased = new Promise((resolve) => { releaseTransition = resolve }) + let blocked = false + const taskStore: TaskStore = { + get: (id) => innerStore.get(id), + put: (task) => innerStore.put(task), + createIfAbsent: (task) => innerStore.createIfAbsent(task), + delete: (id) => innerStore.delete(id), + async compareAndSet(expected, next) { + if (!blocked && next.status.state === 'working') { + blocked = true + transitionStarted() + await transitionReleased + } + return innerStore.compareAndSet(expected, next) + }, + } + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + a2a: { taskStore }, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + verifySigner: async () => true, + paymentOperations: new MemoryPaymentOperations({ onReclaim: async () => undefined }), + }, + }))) + + const send = app.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9010'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'submission-lease-task', + contextId: 'submission-lease-context', + messageId: 'submission-lease-message', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + await transitionReady + + const pending = await taskStore.get('submission-lease-task') + expect(pending?.status.state).toBe('submitted') + expect(pending?.metadata?.gatewaySubmission).toBeDefined() + expect(pending?.metadata?.gatewayPaymentRecovery).toBeDefined() + + releaseTransition() + const response = await send + expect(response.status).toBe(200) + expect((await response.json() as { result?: { status?: { state?: string } } }) + .result?.status?.state).toBe('completed') + }) + + it('retries cancellation after a heartbeat changes the task row', async () => { + const innerStore = new InMemoryTaskStore() + let cancelAttempts = 0 + const taskStore: TaskStore = { + get: (id) => innerStore.get(id), + put: (task) => innerStore.put(task), + createIfAbsent: (task) => innerStore.createIfAbsent(task), + delete: (id) => innerStore.delete(id), + async compareAndSet(expected, next) { + if (next.status.state === 'canceled' && cancelAttempts++ === 0) { + await innerStore.put({ + ...expected, + metadata: { + ...(expected.metadata ?? {}), + gatewayExecution: { + version: 1, + requestId: 'heartbeat-owner', + lease: { id: 'heartbeat-owner', expiresAt: Date.now() + 60_000 }, + }, + }, + }) + return false + } + return innerStore.compareAndSet(expected, next) + }, + } + await taskStore.put({ + kind: 'task', + id: 'cancel-heartbeat-task', + contextId: 'cancel-heartbeat-context', + status: { state: 'working', timestamp: new Date().toISOString() }, + metadata: { + gatewayOrigin: { version: 1, agentId: agent.id, agentSlug: agent.slug }, + gatewayExecution: { + version: 1, + requestId: 'heartbeat-owner', + lease: { id: 'heartbeat-owner', expiresAt: Date.now() + 60_000 }, + }, + }, + }) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + a2a: { taskStore }, + }))) + + const response = await app.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/cancel', + params: { id: 'cancel-heartbeat-task' }, + }), + }) + const body = await response.json() as { result?: Task; error?: unknown } + + expect(response.status).toBe(200) + expect(body.error).toBeUndefined() + expect(body.result?.status.state).toBe('canceled') + expect(cancelAttempts).toBe(2) + }) + it('quotes retained A2A history before charging a continuation', async () => { const taskStore = new InMemoryTaskStore() let invocations = 0 @@ -564,6 +697,45 @@ describe('PR #11 production regressions', () => { expect(body.result?.status?.state).toBe('failed') }) + it('fails a working task after its execution fence expires', async () => { + const taskStore = new InMemoryTaskStore() + const now = Date.now() + await taskStore.put({ + kind: 'task', + id: 'expired-execution-task', + contextId: 'expired-execution-context', + status: { state: 'working', timestamp: new Date(now).toISOString() }, + metadata: { + gatewayOrigin: { version: 1, agentId: agent.id, agentSlug: agent.slug }, + gatewayExecution: { + version: 1, + requestId: 'abandoned-execution-request', + lease: { id: 'abandoned-execution-request', expiresAt: now - 1 }, + }, + }, + }) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + a2a: { taskStore }, + }))) + + const response = await app.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/get', + params: { id: 'expired-execution-task' }, + }), + }) + const body = await response.json() as { result?: Task } + + expect(body.result?.status.state).toBe('failed') + expect(body.result?.metadata?.gatewayExecution).toBeUndefined() + expect(body.result?.metadata?.gatewayExecutionRecovery).toBeDefined() + }) + it('rejects a legacy check-then-mark nonce store instead of racing two payments', async () => { const seen = new Set() const legacyStore: NonceStore = { @@ -870,6 +1042,8 @@ describe('PR #11 production regressions', () => { const urls = [ 'https://[::ffff:169.254.169.254]/metadata', 'https://[::ffff:a9fe:a9fe]/metadata', + 'https://[::127.0.0.1]/metadata', + 'https://[::a9fe:a9fe]/metadata', ] for (const [index, url] of urls.entries()) { From a2232ac1e14dee4ea8f879db49459ce417bf17e3 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 05:57:00 -0600 Subject: [PATCH 28/32] fix(gateway): close terminal delivery and v1 payment races --- .../2026-08-15T11-54-51Z/findings.jsonl | 1 + .../2026-08-15T11-54-51Z/manifest.json | 29 +++ .../2026-08-15T11-54-51Z/summary.md | 35 ++++ .agent/skill-runs.jsonl | 3 + README.md | 4 +- docs/a2a-long-horizon.md | 6 +- src/a2a/handler.ts | 70 ++++++- src/a2a/push-notifications.ts | 12 ++ src/dispatch.ts | 11 +- src/middleware.ts | 10 + src/types.ts | 6 +- tests/middleware.test.ts | 2 + tests/pr11-regressions.test.ts | 181 +++++++++++++++++- 13 files changed, 361 insertions(+), 9 deletions(-) create mode 100644 .agent/critical-audit/2026-08-15T11-54-51Z/findings.jsonl create mode 100644 .agent/critical-audit/2026-08-15T11-54-51Z/manifest.json create mode 100644 .agent/critical-audit/2026-08-15T11-54-51Z/summary.md diff --git a/.agent/critical-audit/2026-08-15T11-54-51Z/findings.jsonl b/.agent/critical-audit/2026-08-15T11-54-51Z/findings.jsonl new file mode 100644 index 0000000..48bf9cb --- /dev/null +++ b/.agent/critical-audit/2026-08-15T11-54-51Z/findings.jsonl @@ -0,0 +1 @@ +{"finding_count":0,"verdict":"APPROVE"} diff --git a/.agent/critical-audit/2026-08-15T11-54-51Z/manifest.json b/.agent/critical-audit/2026-08-15T11-54-51Z/manifest.json new file mode 100644 index 0000000..603d85c --- /dev/null +++ b/.agent/critical-audit/2026-08-15T11-54-51Z/manifest.json @@ -0,0 +1,29 @@ +{ + "skill": "/critical-audit", + "target": "agent-gateway PR #11 blocking findings", + "base_ref": "HEAD", + "base_sha": "43897380c5c7034b59f5ac5d06885679d5a9bb94", + "head_sha": "43897380c5c7034b59f5ac5d06885679d5a9bb94", + "working_tree_diff": true, + "files": 9, + "project_type": "TypeScript package", + "reviewers": [ + "A correctness and security", + "B architecture and quality", + "C standards and real-system coverage" + ], + "review_mode": "serial", + "independent_review": "Codex gpt-5.6-sol inspected the diff and ran direct race and timeout smoke checks in read-only mode.", + "checks": [ + "pnpm typecheck", + "TSX_DISABLE_CACHE=1 NODE_OPTIONS=--import tsx pnpm exec vitest run ... tests/pr11-regressions.test.ts", + "pnpm test", + "pnpm build", + "git diff --check" + ], + "findings": 0, + "dropped": 0, + "verdict": "APPROVE", + "not_inspected": "Live SQL provider behavior and external payment provider behavior; local stores and the public operation contract were inspected.", + "timestamp": "2026-08-15T11:54:51Z" +} diff --git a/.agent/critical-audit/2026-08-15T11-54-51Z/summary.md b/.agent/critical-audit/2026-08-15T11-54-51Z/summary.md new file mode 100644 index 0000000..11146cb --- /dev/null +++ b/.agent/critical-audit/2026-08-15T11-54-51Z/summary.md @@ -0,0 +1,35 @@ +# Audit: PR #11 blocking findings — 43897380..43897380 — n=9 files, 0 findings + +**Verdict:** APPROVE — no reproducible blocking defect remains · 0 CRITICAL / 0 HIGH / 0 MEDIUM / 0 LOW +**Worst:** none · cost if shipped unmeasured because no finding remains +**Next:** `/verify` with the full suite and build + +## Scope + +| Field | Value | +|---|---| +| Files | n=9 changed files in the working tree | +| Base..head | `43897380c5c7034b59f5ac5d06885679d5a9bb94..working tree` | +| Project type | TypeScript package | +| Reviewers | A, B, C · serial | +| Not inspected | Live SQL and external provider systems; local durable contracts were inspected | + +## Findings — 0 of 0, ranked + +| # | Sev | file:line | Defect | Failure scenario | Status | Evidence | Fix | Verification | Cost if shipped | Saved if fixed | +|---:|---|---|---|---|---|---|---|---|---:|---:| +| — | — | — | — | — | — | — | — | — | 0 | 0 | + +0 dropped of 0 reviewed findings. + +## Assumptions & unverified + +| Assumption | Finding it would flip | Check that settles it | +|---|---|---| +| Production task stores implement atomic compare-and-set | Duplicate terminal delivery | SQL adapter integration test against the deployment database | +| Version 2 payment operations persist provider ownership by operation ID | Stranded external funds | Provider integration test with an ambiguous acknowledgement | + +## Self-gate + +9/9 passed — failed: none. +1 verdict = decision + 1 number · 2 every finding has file:line · 3 concrete failure scenario · 4 status label · 5 evidence pointer · 6 cost both sides · 7 fix and verification · 8 zero unsupported adjectives · 9 words ≤600 outside tables. diff --git a/.agent/skill-runs.jsonl b/.agent/skill-runs.jsonl index 47dc524..2c5ac0a 100644 --- a/.agent/skill-runs.jsonl +++ b/.agent/skill-runs.jsonl @@ -3,3 +3,6 @@ {"skill":"/harden","ts":"2026-08-15T06:00:39Z","project":"agent-gateway-compat-store","target":"Gateway compatibility, MPP auth, SQL task GC","operatorPrompt":"","durationMin":null,"verdict":"PASS","dispatchedTo":"/critical-audit","operatorOverride":null,"transcriptPath":null,"traceDir":null} {"skill":"/critical-audit","ts":"2026-08-15T06:00:39Z","project":"agent-gateway-compat-store","target":"Gateway compatibility/store commit n=11 files","operatorPrompt":"","durationMin":null,"verdict":"APPROVE","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} {"skill":"/verify","ts":"2026-08-15T06:00:39Z","project":"agent-gateway-compat-store","target":"Gateway compatibility/store commit","operatorPrompt":"","durationMin":null,"verdict":"SHIP_IT","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/harden","ts":"2026-08-15T11:56:05Z","project":"agent-gateway-x402-priced-reservation","target":"PR #11 terminal delivery and x402 v1 safety n=9 files","operatorPrompt":"","durationMin":null,"verdict":"PASS","dispatchedTo":"/critical-audit","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/critical-audit","ts":"2026-08-15T11:56:07Z","project":"agent-gateway-x402-priced-reservation","target":"PR #11 blocking findings n=9 files","operatorPrompt":"","durationMin":null,"verdict":"APPROVE","dispatchedTo":"/verify","operatorOverride":null,"transcriptPath":null,"traceDir":null} +{"skill":"/verify","ts":"2026-08-15T11:56:11Z","project":"agent-gateway-x402-priced-reservation","target":"PR #11 terminal delivery and x402 v1 safety","operatorPrompt":"","durationMin":null,"verdict":"SHIP_IT","dispatchedTo":"/stop","operatorOverride":null,"transcriptPath":null,"traceDir":null} diff --git a/README.md b/README.md index 87d40dd..c5fe633 100644 --- a/README.md +++ b/README.md @@ -46,9 +46,11 @@ app.route('/v1/agents', createAgentGateway({ `x402.verifySigner` is required for production. Set `x402.demoMode: true` only for local development and tests; that explicit mode also enables the built-in `sk_agent_*` demo key verifier. Keep `verifySigner` free of side effects. -Use `authorizePayment` to reserve or claim funds after rate limits, content checks, and product authorization succeed. +Use version 2's `authorizePayment` to reserve or claim funds after rate limits, content checks, and product authorization succeed. For production version 2, set `x402.paymentProtocolVersion: 2`, provide `paymentOperations`, and return its operation from `authorizePayment`. Production version 2 also requires a durable `paymentRecovery.store`. +Production version 1 is read-only and must not configure `authorizePayment`. +Use version 2 whenever authorization can reserve, charge, or otherwise mutate external funds. Run `recoverPayments(config)` from a private scheduled worker. Every live request and worker uses a unique durable fence token. A stale request or worker cannot update a row after another worker takes its lease. diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 9b48272..4a28b2a 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -249,7 +249,11 @@ async function verify(req: Request, secret: string): Promise { ### Delivery semantics -- **Fire-once.** No retries. If a webhook returns non-2xx or the request fails, the gateway logs and moves on. The consumer's webhook handler SHOULD idempotently re-fetch state via `tasks/get` rather than rely on at-least-once delivery. +- **Fire-once.** The gateway claims each terminal `(task, config)` delivery with the durable task store before it sends the webhook. + Concurrent workers therefore produce at most one attempt for that terminal state. + There are no retries after a failed attempt. + If a webhook returns non-2xx or the request fails, the gateway logs and moves on. + The consumer's webhook handler SHOULD idempotently re-fetch state via `tasks/get` rather than rely on at-least-once delivery. - **No partial-state deliveries.** Only terminal transitions fire push. `input-required` is NOT terminal — it's a pause, not an end. - **Fire even on cancel + fail.** Consumers want to know the task ended for any reason, not just success. diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 50b0658..53bc650 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -118,6 +118,11 @@ interface TaskPaymentRecoveryMarker { id: string } +interface TaskPushDeliveryClaims { + version: 1 + claims: Record +} + interface TaskOriginBinding { version: 1 agentId: string @@ -1564,6 +1569,7 @@ const FINALIZING_METADATA_KEY = 'gatewayFinalizing' const PAYMENT_RELEASE_METADATA_KEY = 'gatewayPaymentRelease' const PAYMENT_RECOVERY_METADATA_KEY = 'gatewayPaymentRecovery' const EXECUTION_RECOVERY_METADATA_KEY = 'gatewayExecutionRecovery' +const PUSH_DELIVERY_METADATA_KEY = 'gatewayPushDelivery' const FINALIZATION_LEASE_MS = 5 * 60 * 1000 const PAYMENT_RELEASE_LEASE_MS = 5 * 60 * 1000 @@ -2572,13 +2578,20 @@ async function readJsonBody(request: Request): Promise { async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise { if (!deps.pushStore || !TERMINAL_STATES.has(task.status.state)) return try { + const deliveryTask = clearPushDeliveryClaims(task) await deliverPushNotifications({ - task, + task: deliveryTask, store: deps.pushStore, webhookSecret: deps.config.a2a?.webhookSecret, fetcher: deps.config.a2a?.pushFetcher, urlValidator: deps.config.a2a?.pushUrlValidator, requireUrlValidator: !deps.config.x402.demoMode, + claimDelivery: (taskId, configId, terminalState) => claimTaskPushDelivery( + deps.taskStore, + taskId, + configId, + terminalState, + ), onDelivery: (result) => { if (!result.ok) { deps.state.obs?.onStreamError?.( @@ -2599,3 +2612,58 @@ async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise ) } } + +async function claimTaskPushDelivery( + taskStore: TaskStore, + taskId: string, + configId: string, + terminalState: Task['status']['state'], +): Promise { + if (!TERMINAL_STATES.has(terminalState)) return false + for (let attempt = 0; attempt < 16; attempt += 1) { + const current = await taskStore.get(taskId) + if (!current || current.status.state !== terminalState) return false + const existing = readPushDeliveryClaims(current) + if (existing?.claims[configId] === terminalState) return false + const next: Task = { + ...current, + metadata: { + ...(current.metadata ?? {}), + [PUSH_DELIVERY_METADATA_KEY]: { + version: 1, + claims: { + ...(existing?.claims ?? {}), + [configId]: terminalState, + }, + } satisfies TaskPushDeliveryClaims, + }, + } + if (await compareAndSetTask(taskStore, current, next)) return true + } + throw new Error(`A2A push delivery claim changed too many times for task '${taskId}'`) +} + +function readPushDeliveryClaims(task: Task): TaskPushDeliveryClaims | undefined { + const raw = task.metadata?.[PUSH_DELIVERY_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const record = raw as Partial + if (!record.claims || typeof record.claims !== 'object') return undefined + const claims = Object.fromEntries( + Object.entries(record.claims).filter(([, state]) => + typeof state === 'string' && TERMINAL_STATES.has(state as Task['status']['state']), + ), + ) as Record + return record.version === 1 ? { version: 1, claims } : undefined +} + +function clearPushDeliveryClaims(task: Task): Task { + if (!task.metadata || !(PUSH_DELIVERY_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[PUSH_DELIVERY_METADATA_KEY] + return Object.keys(metadata).length > 0 + ? { ...task, metadata } + : (() => { + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata + })() +} diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index a16c72d..c4d0180 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -298,6 +298,12 @@ export async function deliverPushNotifications(args: { task: Task store: PushNotificationStore webhookSecret: string | undefined + /** Atomically claim one terminal delivery before its external side effect. */ + claimDelivery?: ( + taskId: string, + configId: string, + terminalState: Task['status']['state'], + ) => Promise /** Inject for tests. Defaults to global `fetch`. */ fetcher?: typeof fetch /** Optional DNS-aware host policy for production deployments. */ @@ -320,6 +326,12 @@ export async function deliverPushNotifications(args: { const results: PushDeliveryResult[] = [] for (const config of configs) { + if ( + args.claimDelivery && + !await args.claimDelivery(args.task.id, config.id, args.task.status.state) + ) { + continue + } const headers: Record = { 'Content-Type': 'application/json' } if (config.token) headers['X-A2A-Notification-Token'] = config.token if (signature) headers['X-A2A-Signature'] = signature diff --git a/src/dispatch.ts b/src/dispatch.ts index ae88da1..b8e4c3b 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -588,9 +588,16 @@ export async function claimPayment( } let operation: PaymentOperation | undefined if (config.x402.authorizePayment) { + if (config.x402.paymentProtocolVersion !== 2 && !config.x402.demoMode) { + throw new Error( + 'production x402 version 1 cannot use authorizePayment; ' + + 'use paymentProtocolVersion: 2 with paymentOperations', + ) + } // Version 1 has no durable operation to release if another request wins - // the shared nonce while this callback is still running. Claim first so - // an external reserve or charge cannot happen for a losing request. + // the shared nonce while this callback is still running. This callback + // remains only for explicit demo-mode compatibility; production callers + // must use the durable version 2 operation lifecycle. const legacyClaimed = config.x402.paymentProtocolVersion !== 2 && authz.paymentNonceKey ? await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) : undefined diff --git a/src/middleware.ts b/src/middleware.ts index c87c1c0..3bca144 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -87,6 +87,16 @@ export function createAgentGateway(inputConfig: GatewayConfig) { if (config.x402.paymentOperations && config.x402.paymentProtocolVersion === undefined) { throw new Error('createAgentGateway: paymentProtocolVersion must be explicit when durable payment operations are configured') } + if ( + config.x402.authorizePayment && + config.x402.paymentProtocolVersion !== 2 && + !config.x402.demoMode + ) { + throw new Error( + 'createAgentGateway: production x402 version 1 cannot use authorizePayment; ' + + 'use paymentProtocolVersion: 2 with paymentOperations for durable payment ownership', + ) + } if (config.x402.paymentProtocolVersion === 2 && (!config.x402.paymentOperations || config.x402.paymentOperations.protocolVersion !== 2)) { throw new Error('createAgentGateway: payment protocol version 2 requires durable payment operations') diff --git a/src/types.ts b/src/types.ts index 1442c26..c089616 100644 --- a/src/types.ts +++ b/src/types.ts @@ -101,7 +101,7 @@ export interface X402Config { rpcUrl?: string /** Demo mode: skip signature verification (default: false). NEVER enable in production. */ demoMode?: boolean - /** Protocol version for new durable payment operations. Version 1 remains supported for mixed deploys. */ + /** Protocol version for new durable payment operations. Production version 1 is read-only. */ paymentProtocolVersion?: 1 | 2 /** * Production signature verification. This callback must not reserve, claim, @@ -114,7 +114,9 @@ export interface X402Config { /** * Claim the verified payment after all request checks pass and immediately * before sandbox work starts. Version 2 returns durable operation ownership. - * A boolean return is the version 1 mixed-deploy compatibility path. + * A boolean return is the version 1 demo-only compatibility path. + * Production version 1 must omit this callback because it has no durable + * provider operation or recovery identity. */ authorizePayment?: ( payload: Record, diff --git a/tests/middleware.test.ts b/tests/middleware.test.ts index f44ff13..e13e955 100644 --- a/tests/middleware.test.ts +++ b/tests/middleware.test.ts @@ -1117,6 +1117,7 @@ describe('POST /:slug/chat/completions — authorizeConsumer hook', () => { x402: { operatorAddress, chainId: 3799, + demoMode: true, verifySigner: async () => true, authorizePayment: async () => { paymentAuthorizations += 1 @@ -1145,6 +1146,7 @@ describe('POST /:slug/chat/completions — authorizeConsumer hook', () => { x402: { operatorAddress, chainId: 3799, + demoMode: true, verifySigner: async () => true, authorizePayment: async () => { paymentAuthorizations += 1 diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index 7277af0..c471158 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -3,6 +3,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { SqlTaskStore, type SqlAdapter } from '../src/a2a/task-store-sql' +import { InMemoryPushNotificationStore } from '../src/a2a/push-notifications' import { createAgentGateway } from '../src/middleware' import { dispatchSandboxStreamRich, requiredX402Amount } from '../src/dispatch' import { MemoryNonceStore, claimStoredNonce, type NonceStore } from '../src/nonce-store' @@ -34,14 +35,14 @@ afterEach(() => { vi.useRealTimers() }) -function paymentHeader(nonce: string): string { +function paymentHeader(nonce: string, expiry = Math.floor(Date.now() / 1000) + 600): string { return JSON.stringify({ commitment, signature: '0xsig', operator: operatorAddress, amount: '1000000000', nonce, - expiry: String(Math.floor(Date.now() / 1000) + 600), + expiry: String(expiry), }) } @@ -88,6 +89,182 @@ function durableConfig( } describe('PR #11 production regressions', () => { + it('claims one terminal webhook when cancellation races settlement on two workers', async () => { + const taskStore = new InMemoryTaskStore() + const pushStore = new InMemoryPushNotificationStore() + const nonceStore = new MemoryNonceStore() + const recoveryStore = new MemoryPaymentRecoveryStore() + let settlements = 0 + const operations = new MemoryPaymentOperations({ + onSettle: async () => { settlements += 1 }, + onReclaim: async () => undefined, + }) + let sandboxEntered!: () => void + const sandboxReady = new Promise((resolve) => { sandboxEntered = resolve }) + let releaseSandbox!: () => void + const sandboxReleased = new Promise((resolve) => { releaseSandbox = resolve }) + let deliveries = 0 + const receivedTaskIds: string[] = [] + const webhook = new Hono() + webhook.post('/terminal', async (context) => { + deliveries += 1 + const body = await context.req.json() as { taskId?: string } + if (body.taskId) receivedTaskIds.push(body.taskId) + return context.text('ok') + }) + const pushFetcher = async ( + _input: string | URL | Request, + init?: RequestInit, + ): Promise => webhook.fetch(new Request('https://receiver.local/terminal', init)) + + const config = (sandbox: GatewayConfig['getSandbox']): GatewayConfig => ({ + resolveAgent: async () => agent, + getSandbox: sandbox, + recordUsage: async () => undefined, + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore, + paymentRecovery: { store: recoveryStore }, + a2a: { + taskStore, + pushStore, + pushFetcher: pushFetcher as typeof fetch, + authorizeTaskAccess: async () => true, + }, + }) + + const runner = new Hono() + runner.route('/v1/agents', createAgentGateway(config(async () => ({ + async *streamPrompt() { + sandboxEntered() + yield { type: 'sandbox.usage', data: { usage: usage() } } + await sandboxReleased + }, + })))) + const canceler = new Hono() + canceler.route('/v1/agents', createAgentGateway(config(async () => ({ + async *streamPrompt() { + throw new Error('canceler must not execute the sandbox') + }, + })))) + + const running = runner.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9011'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'pr11-push-race', + contextId: 'pr11-push-context', + messageId: 'pr11-push-message', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + await sandboxReady + await pushStore.set('pr11-push-race', { id: 'terminal', url: 'https://hook.example/terminal' }) + + const cancel = await canceler.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/cancel', + params: { id: 'pr11-push-race' }, + }), + }) + expect(cancel.status).toBe(200) + expect((await cancel.json() as { result?: { status?: { state?: string } } }) + .result?.status?.state).toBe('canceled') + + releaseSandbox() + const runnerResponse = await running + expect(runnerResponse.status).toBe(200) + expect((await runnerResponse.json() as { result?: { status?: { state?: string } } }) + .result?.status?.state).toBe('canceled') + expect(settlements).toBe(1) + expect(deliveries).toBe(1) + expect(receivedTaskIds).toEqual(['pr11-push-race']) + expect((await taskStore.get('pr11-push-race'))?.status.state).toBe('canceled') + }) + + it('rejects production v1 authorization callbacks without durable recovery', () => { + expect(() => createAgentGateway(durableConfig({ + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 1, + authorizePayment: async () => true, + }, + }))).toThrow(/production x402 version 1 cannot use authorizePayment/) + }) + + it('recovers a timed-out v2 reserve by operation id without consuming the nonce', async () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-15T12:00:00.000Z')) + let providerReservations = 0 + let providerRecoveries = 0 + const operations = new MemoryPaymentOperations({ + onClaim: async () => { + providerReservations += 1 + throw new Error('provider timeout after reserve') + }, + onReclaim: async () => { providerRecoveries += 1 }, + }) + const nonceStore = new MemoryNonceStore() + const recoveryStore = new MemoryPaymentRecoveryStore() + const config = durableConfig({ + x402: { + operatorAddress, + chainId: 1, + demoMode: true, + paymentProtocolVersion: 2, + paymentOperations: operations, + }, + nonceStore, + paymentRecovery: { store: recoveryStore, retryDelayMs: 1 }, + }) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const response = await app.request('/v1/agents/pr11/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9012', Math.floor(Date.now() / 1000) + 1), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'reserve' }] }), + }) + await response.text() + + const operationId = `x402:${commitment}:9012` + expect(response.status).toBe(402) + expect(providerReservations).toBe(1) + expect(operations.get(operationId)?.state).toBe('claiming') + expect(await nonceStore.hasSeen(`${commitment}:9012`)).toBe(false) + + vi.advanceTimersByTime(2_000) + const recovered = await recoverPayment(operationId, config, { force: true }) + expect(recovered?.state).toBe('reconciled') + expect(operations.get(operationId)?.state).toBe('reclaimed') + expect(providerRecoveries).toBe(1) + }) + it('does not mark payment execution before sandbox acquisition succeeds', async () => { const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) let sandboxReady = false From 00984b628f67eb8a672ae30439b78b262ec09c4b Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 17:57:36 -0600 Subject: [PATCH 29/32] fix(gateway): close migration and delivery races --- README.md | 7 +- docs/a2a-long-horizon.md | 66 +++++- src/a2a/execution-fence.ts | 53 ++++- src/a2a/handler.ts | 79 ++++++- src/a2a/push-notifications.ts | 98 +++++--- src/a2a/task-store-sql.ts | 150 ++++++++++++- src/a2a/task-store.ts | 31 +++ src/dispatch.ts | 16 ++ src/index.ts | 2 + src/middleware.ts | 7 + src/types.ts | 5 +- tests/a2a-atomicity.test.ts | 15 +- tests/a2a-durability.test.ts | 248 ++++++++++++++++++-- tests/a2a-lifecycle.test.ts | 13 ++ tests/a2a-payment-races.test.ts | 6 + tests/integration-x402.test.ts | 25 ++- tests/pr11-regressions.test.ts | 386 +++++++++++++++++++++++++++++--- 17 files changed, 1073 insertions(+), 134 deletions(-) diff --git a/README.md b/README.md index c5fe633..2b000f7 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,7 @@ Use version 2's `authorizePayment` to reserve or claim funds after rate limits, For production version 2, set `x402.paymentProtocolVersion: 2`, provide `paymentOperations`, and return its operation from `authorizePayment`. Production version 2 also requires a durable `paymentRecovery.store`. Production version 1 is read-only and must not configure `authorizePayment`. +Production x402 version 1 also rejects the legacy `settlePayment` callback before it consumes a nonce. Use version 2 whenever authorization can reserve, charge, or otherwise mutate external funds. Run `recoverPayments(config)` from a private scheduled worker. Every live request and worker uses a unique durable fence token. @@ -112,13 +113,17 @@ Wire protocol handlers only translate their request and response shapes. The gateway speaks Google's A2A protocol alongside its OpenAI-compatible surface: discovery via `.well-known/agent.json`, JSON-RPC 2.0 dispatch for `message/send`, `message/stream`, `tasks/get`, `tasks/cancel`, `tasks/resubscribe`, and the four `tasks/pushNotificationConfig/*` methods. Long-horizon agents — durable tasks across worker restarts, webhook delivery on terminal state, `input-required` pauses with multi-turn continuation — are documented in [`docs/a2a-long-horizon.md`](./docs/a2a-long-horizon.md). Production A2A task control requires `a2a.authorizeTaskAccess`; explicit demo mode is the local-test exception. -Custom production task stores must implement atomic `createIfAbsent` and `compareAndSet` methods. +Custom production task stores must implement atomic `createIfAbsent`, `compareAndSet`, and `compareAndSetExecution` methods. +`compareAndSetExecution` must reject a renewal when the stored owner lease has expired. Task stores must retain payment recovery metadata until reconciliation clears it. The short-lived `gatewaySubmission` marker is not a payment recovery record and may expire with its task. The bundled memory and SQL stores enforce this rule even after the normal task TTL. Push destinations must use HTTPS without URL credentials. Push delivery does not follow redirects. Production push delivery also requires `a2a.pushUrlValidator` to reject private DNS destinations. +Production push delivery requires `a2a.webhookSecret` so every webhook has an HMAC signature. +The exported `deliverPushNotifications` function also requires a non-empty secret. +Use `deliverDemoPushNotifications` only for explicit local demo mode. Tasks created before this release have no recorded origin and fail closed; migrate them with a verified owner binding or let them expire. The payment claim keeps its submission lease until the atomic submitted-to-working transition. An expired execution lease fails the working task and preserves its payment recovery markers. diff --git a/docs/a2a-long-horizon.md b/docs/a2a-long-horizon.md index 4a28b2a..225250e 100644 --- a/docs/a2a-long-horizon.md +++ b/docs/a2a-long-horizon.md @@ -16,7 +16,8 @@ Durable storage and push notifications are enabled through `GatewayConfig.a2a`. By default `GatewayConfig.a2a.taskStore` is in-memory: fast, zero-config, fine for tests and single-machine deployments. Production deployments swap in `SqlTaskStore` against any SQL store — D1, postgres, sqlite, libSQL, Turso — via a `SqlAdapter` shim. -Custom production task stores must implement both atomic methods, `createIfAbsent` and `compareAndSet`. +Custom production task stores must implement `createIfAbsent`, `compareAndSet`, and `compareAndSetExecution`. +`compareAndSetExecution` must reject a renewal when the stored owner lease has expired. The gateway keeps the OpenAI surface available when an older store lacks either method. It returns `503` for A2A until the store is upgraded, rather than using a cross-worker unsafe fallback. Use `SqlTaskStore` or another atomic adapter for multi-worker production deployments. @@ -53,27 +54,72 @@ Tasks stored before this release do not have a `gatewayOrigin` binding and canno Migrate those records with a verified owner binding, or allow them to expire before enabling production control methods. Push registration rejects reserved IP literals and private hostname suffixes. Production push delivery also requires `a2a.pushUrlValidator`, which should apply the deployment's DNS-aware private-network policy. +The exported `deliverPushNotifications` function requires a non-empty HMAC secret. +Use `deliverDemoPushNotifications` only for explicit local demo mode. ### D1 (Cloudflare Workers) +Cloudflare does not invoke arbitrary `migrate` exports on a module Worker. +Create a D1 migration file and apply it before serving traffic. + +```sql +-- migrations/0001_a2a.sql +CREATE TABLE IF NOT EXISTS a2a_tasks ( + id TEXT PRIMARY KEY, + context_id TEXT NOT NULL, + state TEXT NOT NULL, + payload TEXT NOT NULL, + updated_at INTEGER NOT NULL, + execution_request_id TEXT, + execution_lease_expires_at REAL +); +CREATE INDEX IF NOT EXISTS idx_a2a_tasks_context + ON a2a_tasks (context_id, updated_at); +CREATE TABLE IF NOT EXISTS a2a_push_configs ( + task_id TEXT NOT NULL, + config_id TEXT NOT NULL, + url TEXT NOT NULL, + token TEXT, + authentication TEXT, + PRIMARY KEY (task_id, config_id) +); +``` + +Apply it with `npx wrangler d1 migrations apply DB_NAME --remote`. +This creates both tables before the Worker uses them. +For an existing `a2a_tasks` table, apply a separate upgrade migration with these two statements before deployment: + +```sql +ALTER TABLE a2a_tasks ADD COLUMN execution_request_id TEXT; +ALTER TABLE a2a_tasks ADD COLUMN execution_lease_expires_at REAL; +``` + +Do not serve traffic until this upgrade is applied. +The upgrade leaves both columns `NULL` for rows created by the old schema. +On the first renewal of a working legacy task, `SqlTaskStore` seeds both columns from the valid `gatewayExecution` payload marker in one atomic compare-and-set. +The compare-and-set requires the old payload, `state = 'working'`, and both columns to remain `NULL`, so only one worker can seed the row. +If the marker is malformed or only one column is populated, renewal fails closed. +Do not replace this contract with an unconditional payload backfill. + ```ts import { createAgentGateway, d1ToSqlAdapter, - InMemoryPushNotificationStore, + SqlPushNotificationStore, SqlTaskStore, } from '@tangle-network/agent-gateway' export default { async fetch(req: Request, env: { DB: D1Database }) { - const taskStore = new SqlTaskStore(d1ToSqlAdapter(env.DB)) - await taskStore.migrate() // run once at deploy; idempotent + const db = d1ToSqlAdapter(env.DB) + const taskStore = new SqlTaskStore(db) + const pushStore = new SqlPushNotificationStore(db) const gw = createAgentGateway({ // ... your existing config ... a2a: { taskStore, - pushStore: new InMemoryPushNotificationStore(), + pushStore, webhookSecret: env.A2A_WEBHOOK_SECRET, }, }) @@ -145,11 +191,17 @@ CREATE TABLE IF NOT EXISTS a2a_tasks ( context_id TEXT NOT NULL, state TEXT NOT NULL, payload TEXT NOT NULL, -- JSON Task envelope - updated_at INTEGER NOT NULL -- ms since epoch + updated_at INTEGER NOT NULL, -- ms since epoch + execution_request_id TEXT, -- nullable durable execution owner + execution_lease_expires_at REAL -- ms since epoch ); CREATE INDEX IF NOT EXISTS idx_a2a_tasks_context ON a2a_tasks (context_id, updated_at); ``` +The two execution columns are required by the current store schema. +For an existing table, apply the upgrade migration above before traffic. +Rows from the old schema keep `NULL` execution columns until their first valid renewal, which performs the guarded seed described above. + One table stores the JSON payload. TTL is enforced at read time and defaults to one hour. Configure it with `new SqlTaskStore(db, { ttlMs })`. @@ -169,7 +221,7 @@ When a task reaches a terminal state (`completed`, `canceled`, `failed`, `reject ```ts import { createAgentGateway, - InMemoryPushNotificationStore, + d1ToSqlAdapter, SqlPushNotificationStore, } from '@tangle-network/agent-gateway' diff --git a/src/a2a/execution-fence.ts b/src/a2a/execution-fence.ts index 3fdbe79..b310a39 100644 --- a/src/a2a/execution-fence.ts +++ b/src/a2a/execution-fence.ts @@ -7,12 +7,17 @@ export const TASK_EXECUTION_METADATA_KEY = 'gatewayExecution' const TASK_EXECUTION_VERSION = 1 as const const TASK_EXECUTION_LEASE_MS = 5 * 60 * 1000 -interface TaskExecutionMarker { +export interface TaskExecutionMarker { version: typeof TASK_EXECUTION_VERSION requestId: string lease: { id: string; expiresAt: number } } +export type TaskExecutionInspection = + | { state: 'absent' } + | { state: 'valid'; marker: TaskExecutionMarker } + | { state: 'malformed'; reason: string } + export class TaskExecutionCanceledError extends Error { constructor(taskId: string) { super(`A2A task '${taskId}' was canceled before sandbox execution`) @@ -32,7 +37,11 @@ export async function claimTaskExecution( if (!current || current.status.state !== 'working') { throw new TaskExecutionCanceledError(task.id) } - const existing = readTaskExecution(current) + const inspection = inspectTaskExecution(current) + if (inspection.state === 'malformed') { + throw new Error(`A2A task '${task.id}' has a malformed execution marker`) + } + const existing = inspection.state === 'valid' ? inspection.marker : undefined if (existing && existing.lease.expiresAt > now) { if (existing.requestId === requestId) return current throw new Error(`A2A task '${task.id}' is already executing`) @@ -48,21 +57,32 @@ export async function renewTaskExecution( store: TaskStore, taskId: string, requestId: string, - now = Date.now(), + now?: number, ): Promise { for (let attempt = 0; attempt < 8; attempt += 1) { const current = await store.get(taskId) - const marker = current ? readTaskExecution(current) : undefined + const inspection = current ? inspectTaskExecution(current) : { state: 'absent' as const } + if (inspection.state === 'malformed') { + throw new Error(`A2A task '${taskId}' has a malformed execution marker`) + } + const marker = inspection.state === 'valid' ? inspection.marker : undefined if (!current || current.status.state !== 'working' || marker?.requestId !== requestId) { throw new TaskExecutionCanceledError(taskId) } - const next = withTaskExecution(current, requestId, now) - if (store.compareAndSet && await store.compareAndSet(current, next)) return next + const renewalNow = now ?? Date.now() + if (marker.lease.expiresAt <= renewalNow) { + throw new TaskExecutionCanceledError(taskId) + } + const next = withTaskExecution(current, requestId, renewalNow) + if (!store.compareAndSetExecution) { + throw new Error('A2A task store does not provide atomic execution renewal') + } + if (await store.compareAndSetExecution(current, next, requestId, renewalNow)) return next } throw new Error(`A2A task '${taskId}' changed too many times while execution was active`) } -/** Cancellation is rejected only while a live execution fence is held. */ +/** Remote cancellation is rejected while a live execution fence is held. */ export function hasActiveTaskExecution(task: Task, now = Date.now()): boolean { const marker = readTaskExecution(task) return marker !== undefined && marker.lease.expiresAt > now @@ -74,6 +94,11 @@ export function hasExpiredTaskExecution(task: Task, now = Date.now()): boolean { return marker !== undefined && !hasActiveTaskExecution(task, now) } +/** A working task with an execution key that cannot be trusted. */ +export function hasMalformedTaskExecution(task: Task): boolean { + return inspectTaskExecution(task).state === 'malformed' +} + /** Remove the marker when the task reaches a terminal or paused state. */ export function clearTaskExecution(task: Task): Task { if (!task.metadata || !(TASK_EXECUTION_METADATA_KEY in task.metadata)) return task @@ -102,8 +127,16 @@ function withTaskExecution(task: Task, requestId: string, now: number): Task { } function readTaskExecution(task: Task): TaskExecutionMarker | undefined { + const inspection = inspectTaskExecution(task) + return inspection.state === 'valid' ? inspection.marker : undefined +} + +export function inspectTaskExecution(task: Task): TaskExecutionInspection { const raw = task.metadata?.[TASK_EXECUTION_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined + if (raw === undefined) return { state: 'absent' } + if (!raw || typeof raw !== 'object' || Array.isArray(raw)) { + return { state: 'malformed', reason: 'marker must be an object' } + } const marker = raw as Partial if ( marker.version !== TASK_EXECUTION_VERSION || @@ -114,6 +147,6 @@ function readTaskExecution(task: Task): TaskExecutionMarker | undefined { marker.lease.id.length === 0 || typeof marker.lease.expiresAt !== 'number' || !Number.isFinite(marker.lease.expiresAt) - ) return undefined - return marker as TaskExecutionMarker + ) return { state: 'malformed', reason: 'marker fields are invalid' } + return { state: 'valid', marker: marker as TaskExecutionMarker } } diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 53bc650..9dfddfd 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -44,13 +44,19 @@ import { buildAgentCard } from './agent-card' import { claimTaskExecution, clearTaskExecution, + hasActiveTaskExecution, hasExpiredTaskExecution, + inspectTaskExecution, + hasMalformedTaskExecution, renewTaskExecution, } from './execution-fence' import { fail, ok, parseEnvelope } from './jsonrpc' import { + deliverDemoPushNotifications, deliverPushNotifications, validatePushNotificationUrl, + type PushNotificationDeliveryOptions, + type PushDeliveryResult, type PushNotificationStore, type TaskPushNotificationConfig, } from './push-notifications' @@ -186,6 +192,11 @@ class CancelRegistry { return this.finalizing.has(taskId) } + has(taskId: string): boolean { + const controller = this.controllers.get(taskId) + return controller !== undefined && !controller.signal.aborted + } + cancel(taskId: string): boolean { if (this.finalizing.has(taskId)) return false const c = this.controllers.get(taskId) @@ -927,13 +938,16 @@ async function handleTasksCancel( if ( isTaskFinalizing(task) || - cancels.isFinalizing(task.id) + cancels.isFinalizing(task.id) || + (hasActiveTaskExecution(task) && !cancels.has(task.id)) ) { return c.json( fail( req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, - `task '${task.id}' is being finalized`, + hasActiveTaskExecution(task) + ? `task '${task.id}' has an active execution fence` + : `task '${task.id}' is being finalized`, ), ) } @@ -950,6 +964,11 @@ async function handleTasksCancel( fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' is being finalized`), ) } + if (hasActiveTaskExecution(candidate) && !cancels.has(candidate.id)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' has an active execution fence`), + ) + } const canceled = withStatus(candidate, 'canceled') if (await compareAndSetTask(deps.taskStore, candidate, canceled)) { stillActive = cancels.cancel(task.id) @@ -1861,12 +1880,15 @@ function clearPaymentRecoveryMarker(task: Task): Task { } function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): TaskStore { - if (typeof taskStore.createIfAbsent === 'function' && typeof taskStore.compareAndSet === 'function') { + const hasCreateIfAbsent = typeof taskStore.createIfAbsent === 'function' + const hasCompareAndSet = typeof taskStore.compareAndSet === 'function' + const hasCompareAndSetExecution = typeof taskStore.compareAndSetExecution === 'function' + if (hasCreateIfAbsent && hasCompareAndSet && hasCompareAndSetExecution) { return taskStore } if (!allowUnsafeFallback) { throw new Error( - 'A2A production task store must implement createIfAbsent and compareAndSet', + 'A2A production task store must implement createIfAbsent, compareAndSet, and compareAndSetExecution', ) } return { @@ -1874,16 +1896,37 @@ function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): put: (task) => taskStore.put(task), delete: (id) => taskStore.delete(id), async createIfAbsent(task) { + if (hasCreateIfAbsent) return taskStore.createIfAbsent!(task) if (await taskStore.get(task.id)) return false await taskStore.put(task) return true }, async compareAndSet(expected, next) { + if (hasCompareAndSet) return taskStore.compareAndSet!(expected, next) const current = await taskStore.get(expected.id) if (!current || JSON.stringify(current) !== JSON.stringify(expected)) return false await taskStore.put(next) return true }, + async compareAndSetExecution(expected, next, requestId, now) { + if (hasCompareAndSetExecution) { + return taskStore.compareAndSetExecution!(expected, next, requestId, now) + } + const current = await taskStore.get(expected.id) + const expectedMarker = inspectTaskExecution(current ?? expected) + const nextMarker = inspectTaskExecution(next) + if ( + !current || + JSON.stringify(current) !== JSON.stringify(expected) || + expectedMarker.state !== 'valid' || + nextMarker.state !== 'valid' || + expectedMarker.marker.requestId !== requestId || + nextMarker.marker.requestId !== requestId || + expectedMarker.marker.lease.expiresAt <= now + ) return false + await taskStore.put(next) + return true + }, } } @@ -2296,17 +2339,21 @@ async function recoverTaskIfNeeded( } async function recoverExpiredExecutionIfNeeded(task: Task, deps: A2AHandlerDeps): Promise { + const malformed = hasMalformedTaskExecution(task) if ( task.status.state !== 'working' || - isTaskFinalizing(task) || - !hasExpiredTaskExecution(task) + (isTaskFinalizing(task) && !malformed) || + (!malformed && !hasExpiredTaskExecution(task)) ) return task + const inspection = inspectTaskExecution(task) const failed: Task = { ...withStatus(task, 'failed'), metadata: { ...(clearTaskExecution(task).metadata ?? {}), [EXECUTION_RECOVERY_METADATA_KEY]: { - error: 'A2A execution lease expired before a task result was stored', + error: inspection.state === 'malformed' + ? `A2A execution marker was malformed: ${inspection.reason}` + : 'A2A execution lease expired before a task result was stored', }, }, } @@ -2577,12 +2624,17 @@ async function readJsonBody(request: Request): Promise { */ async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise { if (!deps.pushStore || !TERMINAL_STATES.has(task.status.state)) return + const webhookSecret = deps.config.a2a?.webhookSecret + const hasWebhookSecret = typeof webhookSecret === 'string' && webhookSecret.trim().length > 0 + if (!deps.config.x402.demoMode && !hasWebhookSecret) { + console.error(`[agent-gateway] production A2A push requires a webhookSecret for task ${task.id}`) + return + } try { const deliveryTask = clearPushDeliveryClaims(task) - await deliverPushNotifications({ + const deliveryArgs: Omit = { task: deliveryTask, store: deps.pushStore, - webhookSecret: deps.config.a2a?.webhookSecret, fetcher: deps.config.a2a?.pushFetcher, urlValidator: deps.config.a2a?.pushUrlValidator, requireUrlValidator: !deps.config.x402.demoMode, @@ -2592,7 +2644,7 @@ async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise configId, terminalState, ), - onDelivery: (result) => { + onDelivery: (result: PushDeliveryResult) => { if (!result.ok) { deps.state.obs?.onStreamError?.( { requestId: result.taskId, agentSlug: task.id, startMs: Date.now() }, @@ -2603,7 +2655,12 @@ async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise ) } }, - }) + } + if (hasWebhookSecret) { + await deliverPushNotifications({ ...deliveryArgs, webhookSecret }) + } else { + await deliverDemoPushNotifications(deliveryArgs) + } } catch (err) { // Catastrophic failure of the push pipeline itself (e.g. the store threw). // Logged but never escalated — a busted webhook MUST NOT fail the agent. diff --git a/src/a2a/push-notifications.ts b/src/a2a/push-notifications.ts index c4d0180..6c90b9e 100644 --- a/src/a2a/push-notifications.ts +++ b/src/a2a/push-notifications.ts @@ -223,16 +223,14 @@ export class SqlPushNotificationStore implements PushNotificationStore { async set(taskId: string, config: PushNotificationConfig): Promise { const auth = config.authentication ? JSON.stringify(config.authentication) : null - const updated = await this.db.exec( - `UPDATE ${this.table} SET url = ?, token = ?, authentication = ? WHERE task_id = ? AND config_id = ?`, - [config.url, config.token ?? null, auth, taskId, config.id], + await this.db.exec( + `INSERT INTO ${this.table} (task_id, config_id, url, token, authentication) VALUES (?, ?, ?, ?, ?) + ON CONFLICT (task_id, config_id) DO UPDATE SET + url = excluded.url, + token = excluded.token, + authentication = excluded.authentication`, + [taskId, config.id, config.url, config.token ?? null, auth], ) - if (updated.rowsAffected === 0) { - await this.db.exec( - `INSERT INTO ${this.table} (task_id, config_id, url, token, authentication) VALUES (?, ?, ?, ?, ?)`, - [taskId, config.id, config.url, config.token ?? null, auth], - ) - } } async get(taskId: string, configId: string): Promise { @@ -285,19 +283,10 @@ export class SqlPushNotificationStore implements PushNotificationStore { } } -/** - * Send the webhook for each registered config on a task. Signs the body with - * HMAC-SHA256 against `webhookSecret` so the consumer can verify authenticity. - * Fire-and-forget per the design note above — the function awaits delivery - * (so observability hooks see the result) but does not retry on failure. - * - * The caller decides *when* to deliver — typically on terminal-state - * transitions emitted from `message/send` and `message/stream`. - */ -export async function deliverPushNotifications(args: { +interface PushDeliveryOptions { task: Task store: PushNotificationStore - webhookSecret: string | undefined + webhookSecret?: string /** Atomically claim one terminal delivery before its external side effect. */ claimDelivery?: ( taskId: string, @@ -312,7 +301,39 @@ export async function deliverPushNotifications(args: { requireUrlValidator?: boolean /** Optional callback so the gateway's observer can log delivery outcomes. */ onDelivery?: (result: PushDeliveryResult) => void -}): Promise { +} + +export type PushNotificationDeliveryOptions = Omit & { + webhookSecret: string +} + +/** + * Send signed webhooks for a terminal task. + * + * A non-empty HMAC secret is mandatory. This public production path cannot + * send an unsigned request. + */ +export async function deliverPushNotifications( + args: PushNotificationDeliveryOptions, +): Promise { + if (typeof args.webhookSecret !== 'string' || args.webhookSecret.trim().length === 0) { + throw new Error('deliverPushNotifications requires a non-empty webhookSecret') + } + return deliverPushNotificationsInternal(args) +} + +/** + * Deliver unsigned webhooks only for explicit local demo mode. + * + * Production callers must use `deliverPushNotifications`. + */ +export async function deliverDemoPushNotifications( + args: Omit, +): Promise { + return deliverPushNotificationsInternal(args) +} + +async function deliverPushNotificationsInternal(args: PushDeliveryOptions): Promise { const fetcher = args.fetcher ?? fetch const configs = await args.store.list(args.task.id) const body = JSON.stringify({ @@ -326,6 +347,31 @@ export async function deliverPushNotifications(args: { const results: PushDeliveryResult[] = [] for (const config of configs) { + // Validate before claiming so policy rejection remains retryable. + try { + const url = validatePushNotificationUrl(config.url) + if (!url) { + throw new Error('push notification URL is not a safe HTTPS destination') + } + if (args.requireUrlValidator && !args.urlValidator) { + throw new Error('push notification URL validation is not configured') + } + if (args.urlValidator && !await args.urlValidator(url)) { + throw new Error('push notification URL was rejected by host policy') + } + } catch (err) { + const result: PushDeliveryResult = { + taskId: args.task.id, + configId: config.id, + url: config.url, + ok: false, + error: err instanceof Error ? err.message : String(err), + } + args.onDelivery?.(result) + results.push(result) + continue + } + if ( args.claimDelivery && !await args.claimDelivery(args.task.id, config.id, args.task.status.state) @@ -342,16 +388,6 @@ export async function deliverPushNotifications(args: { let result: PushDeliveryResult try { - const url = new URL(config.url) - if (!validatePushNotificationUrl(config.url)) { - throw new Error('push notification URL is not a safe HTTPS destination') - } - if (args.requireUrlValidator && !args.urlValidator) { - throw new Error('push notification URL validation is not configured') - } - if (args.urlValidator && !await args.urlValidator(url)) { - throw new Error('push notification URL was rejected by host policy') - } const res = await fetcher(config.url, { method: 'POST', headers, diff --git a/src/a2a/task-store-sql.ts b/src/a2a/task-store-sql.ts index dcf1108..6c8228a 100644 --- a/src/a2a/task-store-sql.ts +++ b/src/a2a/task-store-sql.ts @@ -40,6 +40,7 @@ * await store.migrate() */ +import { inspectTaskExecution } from './execution-fence' import { hasPendingPaymentRecovery, type TaskStore } from './task-store' import type { Task } from './types' @@ -97,7 +98,9 @@ const TASKS_TABLE_DDL = (table: string) => ` context_id TEXT NOT NULL, state TEXT NOT NULL, payload TEXT NOT NULL, - updated_at INTEGER NOT NULL + updated_at INTEGER NOT NULL, + execution_request_id TEXT, + execution_lease_expires_at REAL ) ` const CTX_INDEX_DDL = (table: string) => ` @@ -126,13 +129,27 @@ export class SqlTaskStore implements TaskStore { private async readRow(id: string): Promise<{ payload: string updatedAt: number + executionRequestId: string | null + executionLeaseExpiresAt: number | null } | undefined> { - const rows = await this.db.query<{ payload: string; updated_at: number }>( - `SELECT payload, updated_at FROM ${this.table} WHERE id = ?`, + const rows = await this.db.query<{ + payload: string + updated_at: number + execution_request_id?: string | null + execution_lease_expires_at?: number | null + }>( + `SELECT payload, updated_at, execution_request_id, execution_lease_expires_at FROM ${this.table} WHERE id = ?`, [id], ) const row = rows[0] - return row ? { payload: row.payload, updatedAt: row.updated_at } : undefined + return row + ? { + payload: row.payload, + updatedAt: row.updated_at, + executionRequestId: row.execution_request_id ?? null, + executionLeaseExpiresAt: row.execution_lease_expires_at ?? null, + } + : undefined } private isExpired(updatedAt: number, task: Task): boolean { @@ -154,6 +171,17 @@ export class SqlTaskStore implements TaskStore { /** Idempotent. Call once at deploy. */ async migrate(): Promise { await this.db.exec(TASKS_TABLE_DDL(this.table)) + for (const column of [ + 'execution_request_id TEXT', + 'execution_lease_expires_at REAL', + ]) { + try { + await this.db.exec(`ALTER TABLE ${this.table} ADD COLUMN ${column}`) + } catch (error) { + const message = error instanceof Error ? error.message.toLowerCase() : String(error).toLowerCase() + if (!message.includes('duplicate column') && !message.includes('already exists')) throw error + } + } await this.db.exec(CTX_INDEX_DDL(this.table)) } @@ -172,9 +200,19 @@ export class SqlTaskStore implements TaskStore { private async insert(task: Task): Promise { const payload = JSON.stringify(task) + const [executionRequestId, executionLeaseExpiresAt] = executionColumns(task) const result = await this.db.exec( - `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at) VALUES (?, ?, ?, ?, ?)`, - [task.id, task.contextId, task.status.state, payload, Date.now()], + `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at, execution_request_id, execution_lease_expires_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + [ + task.id, + task.contextId, + task.status.state, + payload, + Date.now(), + executionRequestId, + executionLeaseExpiresAt, + ], ) return result.rowsAffected } @@ -182,17 +220,38 @@ export class SqlTaskStore implements TaskStore { async put(task: Task): Promise { const payload = JSON.stringify(task) const updatedAt = Date.now() + const [executionRequestId, executionLeaseExpiresAt] = executionColumns(task) // Adapter-agnostic upsert: try update, fall back to insert if no row // existed. Avoids needing ON CONFLICT (postgres) vs INSERT OR REPLACE // (sqlite/libSQL) divergence at the SQL layer. const updated = await this.db.exec( - `UPDATE ${this.table} SET context_id = ?, state = ?, payload = ?, updated_at = ? WHERE id = ?`, - [task.contextId, task.status.state, payload, updatedAt, task.id], + `UPDATE ${this.table} + SET context_id = ?, state = ?, payload = ?, updated_at = ?, + execution_request_id = ?, execution_lease_expires_at = ? + WHERE id = ?`, + [ + task.contextId, + task.status.state, + payload, + updatedAt, + executionRequestId, + executionLeaseExpiresAt, + task.id, + ], ) if (updated.rowsAffected === 0) { await this.db.exec( - `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at) VALUES (?, ?, ?, ?, ?)`, - [task.id, task.contextId, task.status.state, payload, updatedAt], + `INSERT INTO ${this.table} (id, context_id, state, payload, updated_at, execution_request_id, execution_lease_expires_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + [ + task.id, + task.contextId, + task.status.state, + payload, + updatedAt, + executionRequestId, + executionLeaseExpiresAt, + ], ) } } @@ -223,9 +282,69 @@ export class SqlTaskStore implements TaskStore { async compareAndSet(expected: Task, next: Task): Promise { const expectedPayload = JSON.stringify(expected) const payload = JSON.stringify(next) + const [executionRequestId, executionLeaseExpiresAt] = executionColumns(next) + const result = await this.db.exec( + `UPDATE ${this.table} + SET context_id = ?, state = ?, payload = ?, updated_at = ?, + execution_request_id = ?, execution_lease_expires_at = ? + WHERE id = ? AND payload = ?`, + [ + next.contextId, + next.status.state, + payload, + Date.now(), + executionRequestId, + executionLeaseExpiresAt, + expected.id, + expectedPayload, + ], + ) + return result.rowsAffected === 1 + } + + async compareAndSetExecution( + expected: Task, + next: Task, + requestId: string, + now: number, + ): Promise { + const expectedMarker = inspectTaskExecution(expected) + const nextMarker = inspectTaskExecution(next) + if ( + expectedMarker.state !== 'valid' || + nextMarker.state !== 'valid' || + expectedMarker.marker.requestId !== requestId || + nextMarker.marker.requestId !== requestId || + expectedMarker.marker.lease.expiresAt <= now || + nextMarker.marker.lease.expiresAt <= now + ) return false + const expectedPayload = JSON.stringify(expected) + const payload = JSON.stringify(next) + const updatedAt = Date.now() + // Both NULL columns identify a legacy row whose payload still owns the fence. const result = await this.db.exec( - `UPDATE ${this.table} SET context_id = ?, state = ?, payload = ?, updated_at = ? WHERE id = ? AND payload = ?`, - [next.contextId, next.status.state, payload, Date.now(), expected.id, expectedPayload], + `UPDATE ${this.table} + SET context_id = ?, state = ?, payload = ?, updated_at = ?, + execution_request_id = ?, execution_lease_expires_at = ? + WHERE id = ? + AND payload = ? + AND state = 'working' + AND ( + (execution_request_id = ? AND execution_lease_expires_at > ?) + OR (execution_request_id IS NULL AND execution_lease_expires_at IS NULL) + )`, + [ + next.contextId, + next.status.state, + payload, + updatedAt, + nextMarker.marker.requestId, + nextMarker.marker.lease.expiresAt, + expected.id, + expectedPayload, + requestId, + now, + ], ) return result.rowsAffected === 1 } @@ -259,3 +378,10 @@ export class SqlTaskStore implements TaskStore { .map(({ task }) => task) } } + +function executionColumns(task: Task): [string | null, number | null] { + const inspection = inspectTaskExecution(task) + return inspection.state === 'valid' + ? [inspection.marker.requestId, inspection.marker.lease.expiresAt] + : [null, null] +} diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index cae2f88..7e6d219 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -6,6 +6,7 @@ */ import type { Task } from './types' +import { inspectTaskExecution } from './execution-fence' export interface TaskStore { get(id: string): Promise @@ -14,6 +15,13 @@ export interface TaskStore { createIfAbsent?(task: Task): Promise /** Replace only when the stored task still equals `expected`. Required for production races. */ compareAndSet?(expected: Task, next: Task): Promise + /** Replace an execution marker only while its owner lease is still live. */ + compareAndSetExecution?( + expected: Task, + next: Task, + requestId: string, + now: number, + ): Promise delete(id: string): Promise } @@ -62,6 +70,29 @@ export class InMemoryTaskStore implements TaskStore { return true } + async compareAndSetExecution( + expected: Task, + next: Task, + requestId: string, + now: number, + ): Promise { + this.gc() + const entry = this.entries.get(expected.id) + if (!entry || JSON.stringify(entry.task) !== JSON.stringify(expected)) return false + const expectedMarker = inspectTaskExecution(expected) + const nextMarker = inspectTaskExecution(next) + if ( + expectedMarker.state !== 'valid' || + nextMarker.state !== 'valid' || + expectedMarker.marker.requestId !== requestId || + nextMarker.marker.requestId !== requestId || + expectedMarker.marker.lease.expiresAt <= now || + nextMarker.marker.lease.expiresAt <= now + ) return false + this.entries.set(expected.id, { task: clone(next), expiresAt: Date.now() + this.ttlMs }) + return true + } + async delete(id: string): Promise { const entry = this.entries.get(id) if (entry && hasPendingPaymentRecovery(entry.task)) return diff --git a/src/dispatch.ts b/src/dispatch.ts index b8e4c3b..2ebf25f 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -578,6 +578,7 @@ export async function claimPayment( state: GatewayState, hooks: PaymentClaimHooks = {}, ): Promise { + assertX402V1SettlementSafe(authz, config) if (authz.paymentMethod === 'x402' && authz.paymentPayload) { const context = paymentAuthorizationContext(authz) if (config.x402.paymentProtocolVersion === 2) { @@ -1433,6 +1434,7 @@ export async function settleAndRecord( obs: GatewayObserver | undefined, options: SettleAndRecordOptions = {}, ): Promise { + assertX402V1SettlementSafe(authz, config) const settlementBasis = options.settlementBasis ?? 'usage-receipt' await markRecoverySettling(authz, usage, settlementBasis, config) if (options.usageAlreadyRecorded) await markRecoveryUsageRecorded(authz, config) @@ -1543,6 +1545,20 @@ export async function settleAndRecord( } } +function assertX402V1SettlementSafe(authz: AuthorizedRequest, config: GatewayConfig): void { + if ( + authz.paymentMethod === 'x402' && + config.x402.paymentProtocolVersion !== 2 && + !config.x402.demoMode && + config.settlePayment + ) { + throw new Error( + 'production x402 version 1 cannot use settlePayment; ' + + 'use paymentProtocolVersion: 2 with paymentOperations', + ) + } +} + async function markRecoverySettling( authz: AuthorizedRequest, usage: SandboxUsageReceipt, diff --git a/src/index.ts b/src/index.ts index 514cb64..2d77c58 100644 --- a/src/index.ts +++ b/src/index.ts @@ -131,10 +131,12 @@ export { SqlTaskStore, } from './a2a/task-store-sql' export { + deliverDemoPushNotifications, deliverPushNotifications, InMemoryPushNotificationStore, validatePushNotificationUrl, type PushDeliveryResult, + type PushNotificationDeliveryOptions, type PushNotificationAuthentication, type PushNotificationConfig, type PushNotificationStore, diff --git a/src/middleware.ts b/src/middleware.ts index 3bca144..3be0efc 100644 --- a/src/middleware.ts +++ b/src/middleware.ts @@ -104,6 +104,13 @@ export function createAgentGateway(inputConfig: GatewayConfig) { if (config.x402.paymentProtocolVersion === 1 && config.x402.paymentOperations) { throw new Error('createAgentGateway: version 1 cannot be combined with version 2 payment operations') } + if ( + config.a2a?.pushStore && + !config.x402.demoMode && + (!config.a2a.webhookSecret || config.a2a.webhookSecret.trim().length === 0) + ) { + throw new Error('createAgentGateway: production A2A push requires a webhookSecret') + } const mppMethod = (config.mpp?.method ?? 'blueprintevm').toLowerCase() if (config.mpp?.authenticateCredential !== undefined && typeof config.mpp.authenticateCredential !== 'function') { diff --git a/src/types.ts b/src/types.ts index c089616..b619ac0 100644 --- a/src/types.ts +++ b/src/types.ts @@ -316,6 +316,7 @@ export interface GatewayConfig { /** * Settle a legacy payment after usage attribution is recorded. * Version 2 x402 operations use `x402.paymentOperations` instead. + * Production x402 version 1 rejects this callback before nonce claim. * For API keys, deduct from the spending limit. * Default: no-op in explicit demo mode. */ @@ -416,8 +417,8 @@ export interface GatewayConfig { * Shared HMAC secret used to sign webhook deliveries (`X-A2A-Signature: * sha256=`). The consumer's webhook verifies the body against this * secret to confirm the call originated from this gateway. Required when - * `pushStore` is set; without it, deliveries fire unsigned and a - * malicious party that knows the webhook URL can forge deliveries. + * `pushStore` is set in production. Explicit demo mode may omit it for + * local tests; production deliveries never run unsigned. */ webhookSecret?: string /** diff --git a/tests/a2a-atomicity.test.ts b/tests/a2a-atomicity.test.ts index 3abc7b7..cbe3ae3 100644 --- a/tests/a2a-atomicity.test.ts +++ b/tests/a2a-atomicity.test.ts @@ -98,11 +98,21 @@ class TwoWorkerCreateBarrier implements TaskStore { compareAndSet(expected: Task, next: Task): Promise { return this.inner.compareAndSet(expected, next) } + + compareAndSetExecution( + expected: Task, + next: Task, + requestId: string, + now: number, + ): Promise { + return this.inner.compareAndSetExecution(expected, next, requestId, now) + } } function atomicityConfig( taskStore: TaskStore, counters: { runs: number; records: number; settlements: number }, + demoMode = false, ): GatewayConfig { const sandbox: SandboxBox = { async *streamPrompt() { @@ -119,6 +129,7 @@ function atomicityConfig( x402: { operatorAddress, chainId: 1, + demoMode, paymentProtocolVersion: 1, verifySigner: async () => true, }, @@ -173,8 +184,8 @@ describe('A2A task atomicity and restart recovery', () => { const counters = { runs: 0, records: 0, settlements: 0 } const first = new Hono() const second = new Hono() - first.route('/v1/agents', createAgentGateway(atomicityConfig(taskStore, counters))) - second.route('/v1/agents', createAgentGateway(atomicityConfig(taskStore, counters))) + first.route('/v1/agents', createAgentGateway(atomicityConfig(taskStore, counters, true))) + second.route('/v1/agents', createAgentGateway(atomicityConfig(taskStore, counters, true))) const request = (app: Hono, nonce: string, text: string) => app.request( `/v1/agents/${agent.slug}`, diff --git a/tests/a2a-durability.test.ts b/tests/a2a-durability.test.ts index 5f52df1..3b524aa 100644 --- a/tests/a2a-durability.test.ts +++ b/tests/a2a-durability.test.ts @@ -1,12 +1,12 @@ /** * SqlTaskStore + SqlPushNotificationStore — exercises both stores * end-to-end against a fake-but-honest SqlAdapter that interprets the exact - * statements the stores issue. We do not use sqlite/D1/postgres in unit - * tests; consumers wire the real driver in deployment. + * statements the stores issue, plus real SQLite migration and lease-race tests. * * The fake adapter is intentionally minimal — if a store's SQL drifts, these * tests break loudly, which is the point. */ +import { DatabaseSync } from 'node:sqlite' import { describe, expect, it } from 'vitest' import { @@ -14,6 +14,7 @@ import { type PushNotificationConfig, SqlPushNotificationStore, } from '../src/a2a/push-notifications' +import { claimTaskExecution, renewTaskExecution } from '../src/a2a/execution-fence' import { type SqlAdapter, SqlTaskStore } from '../src/a2a/task-store-sql' import type { Task } from '../src/a2a/types' @@ -23,6 +24,8 @@ interface FakeTaskRow { state: string payload: string updated_at: number + execution_request_id: string | null + execution_lease_expires_at: number | null } interface FakePushRow { task_id: string @@ -40,34 +43,64 @@ function makeTaskAdapter( rows, async exec(sql, params = []) { const s = sql.trim() - if (s.startsWith('CREATE TABLE') || s.startsWith('CREATE INDEX')) { + if ( + s.startsWith('CREATE TABLE') || + s.startsWith('CREATE INDEX') || + s.startsWith('ALTER TABLE') + ) { return { rowsAffected: 0 } } if (s.startsWith('UPDATE')) { - const [contextId, state, payload, updatedAt, id] = params as [ + const [contextId, state, payload, updatedAt, executionRequestId, executionLeaseExpiresAt, id] = params as [ string, string, string, number, + string | null, + number | null, string, ] const row = rows.get(id) if (!row) return { rowsAffected: 0 } + if (s.includes('AND payload = ?') && row.payload !== params[7]) { + return { rowsAffected: 0 } + } + if ( + s.includes('AND execution_request_id = ?') && + (row.state !== 'working' || + row.execution_request_id !== params[8] || + row.execution_lease_expires_at === null || + row.execution_lease_expires_at <= Number(params[9])) + ) { + return { rowsAffected: 0 } + } row.context_id = contextId row.state = state row.payload = payload row.updated_at = updatedAt + row.execution_request_id = executionRequestId + row.execution_lease_expires_at = executionLeaseExpiresAt return { rowsAffected: 1 } } if (s.startsWith('INSERT INTO')) { - const [id, contextId, state, payload, updatedAt] = params as [ + const [id, contextId, state, payload, updatedAt, executionRequestId, executionLeaseExpiresAt] = params as [ string, string, string, string, number, + string | null, + number | null, ] - rows.set(id, { id, context_id: contextId, state, payload, updated_at: updatedAt }) + rows.set(id, { + id, + context_id: contextId, + state, + payload, + updated_at: updatedAt, + execution_request_id: executionRequestId, + execution_lease_expires_at: executionLeaseExpiresAt, + }) return { rowsAffected: 1 } } if (s.startsWith('DELETE')) { @@ -89,20 +122,44 @@ function makeTaskAdapter( if (s.includes('WHERE id =')) { const [id] = params as [string] const row = rows.get(id) - return (row ? [{ payload: row.payload, updated_at: row.updated_at }] : []) as TRow[] + return (row + ? [{ + payload: row.payload, + updated_at: row.updated_at, + execution_request_id: row.execution_request_id, + execution_lease_expires_at: row.execution_lease_expires_at, + }] + : []) as TRow[] } if (s.includes('WHERE context_id =')) { const [ctx] = params as [string] return [...rows.values()] .filter((r) => r.context_id === ctx) .sort((a, b) => b.updated_at - a.updated_at) - .map((r) => ({ payload: r.payload, updated_at: r.updated_at })) as TRow[] + .map((r) => ({ + payload: r.payload, + updated_at: r.updated_at, + execution_request_id: r.execution_request_id, + execution_lease_expires_at: r.execution_lease_expires_at, + })) as TRow[] } throw new Error(`unrecognised query SQL: ${s}`) }, } } +function sqliteAdapter(db: DatabaseSync): SqlAdapter { + return { + async exec(sql, params = []) { + const result = db.prepare(sql).run(...params as never[]) + return { rowsAffected: Number(result.changes) } + }, + async query(sql, params = []) { + return db.prepare(sql).all(...params as never[]) as TRow[] + }, + } +} + function makeTask(id: string, state: Task['status']['state'] = 'submitted'): Task { return { kind: 'task', @@ -121,6 +178,132 @@ function makeTask(id: string, state: Task['status']['state'] = 'submitted'): Tas } describe('SqlTaskStore', () => { + it('rejects a stale SQL renewal while expiry recovery claims the same row', async () => { + const db = new DatabaseSync(':memory:') + try { + const adapter = sqliteAdapter(db) + const staleWorker = new SqlTaskStore(adapter) + const recoveryWorker = new SqlTaskStore(adapter) + await staleWorker.migrate() + + const task: Task = { + ...makeTask('sql-execution-race', 'working'), + metadata: { + gatewayExecution: { + version: 1, + requestId: 'worker-a', + lease: { id: 'worker-a', expiresAt: 1_000 }, + }, + }, + } + await staleWorker.put(task) + const observed = await staleWorker.get(task.id) + expect(observed).toEqual(task) + await expect(renewTaskExecution(staleWorker, task.id, 'worker-a', 1_001)) + .rejects.toThrow("was canceled before sandbox execution") + + const staleRenewal = { + ...task, + metadata: { + ...task.metadata, + gatewayExecution: { + ...task.metadata!.gatewayExecution as Record, + lease: { id: 'worker-a', expiresAt: 5_000 }, + }, + }, + } + const [staleWon, recovered] = await Promise.all([ + staleWorker.compareAndSetExecution( + observed!, + staleRenewal, + 'worker-a', + 1_001, + ), + claimTaskExecution( + recoveryWorker, + observed!, + 'worker-b', + 1_000, + ), + ]) + expect(staleWon).toBe(false) + expect(recovered.metadata?.gatewayExecution).toMatchObject({ requestId: 'worker-b' }) + + expect(await staleWorker.compareAndSetExecution( + observed!, + staleRenewal, + 'worker-a', + 1_001, + )).toBe(false) + expect((await staleWorker.get(task.id))?.metadata?.gatewayExecution) + .toMatchObject({ requestId: 'worker-b' }) + } finally { + db.close() + } + }) + + it('seeds NULL execution columns on the first renewal after an old-schema migration', async () => { + const db = new DatabaseSync(':memory:') + try { + db.exec(` + CREATE TABLE a2a_tasks ( + id TEXT PRIMARY KEY, + context_id TEXT NOT NULL, + state TEXT NOT NULL, + payload TEXT NOT NULL, + updated_at INTEGER NOT NULL + ) + `) + const now = Date.now() + const task: Task = { + ...makeTask('legacy-execution-renewal', 'working'), + metadata: { + gatewayExecution: { + version: 1, + requestId: 'legacy-worker', + lease: { id: 'legacy-worker', expiresAt: now + 60_000 }, + }, + }, + } + db.prepare( + `INSERT INTO a2a_tasks (id, context_id, state, payload, updated_at) + VALUES (?, ?, ?, ?, ?)`, + ).run( + task.id, + task.contextId, + task.status.state, + JSON.stringify(task), + now, + ) + + const store = new SqlTaskStore(sqliteAdapter(db)) + await store.migrate() + expect(db.prepare( + `SELECT execution_request_id, execution_lease_expires_at + FROM a2a_tasks WHERE id = ?`, + ).get(task.id)).toEqual({ + execution_request_id: null, + execution_lease_expires_at: null, + }) + + const renewed = await renewTaskExecution(store, task.id, 'legacy-worker', now) + + expect(renewed.metadata?.gatewayExecution).toMatchObject({ + requestId: 'legacy-worker', + lease: { expiresAt: now + 5 * 60 * 1000 }, + }) + expect(db.prepare( + `SELECT execution_request_id, execution_lease_expires_at + FROM a2a_tasks WHERE id = ?`, + ).get(task.id)).toEqual({ + execution_request_id: 'legacy-worker', + execution_lease_expires_at: now + 5 * 60 * 1000, + }) + } finally { + db.close() + } + }) + it('round-trips a task through put → get with state + history preserved', async () => { const db = makeTaskAdapter() const store = new SqlTaskStore(db) @@ -259,30 +442,25 @@ describe('SqlTaskStore', () => { // ── Push store ────────────────────────────────────────────────────────────── -function makePushAdapter(): SqlAdapter & { rows: Map } { +function makePushAdapter( + beforeUpsert?: () => Promise, +): SqlAdapter & { rows: Map; readonly upsertCalls: number } { const rows = new Map() const key = (taskId: string, configId: string) => `${taskId}::${configId}` + let upsertCalls = 0 return { rows, + get upsertCalls() { + return upsertCalls + }, async exec(sql, params = []) { const s = sql.trim() if (s.startsWith('CREATE TABLE')) return { rowsAffected: 0 } if (s.startsWith('UPDATE')) { - const [url, token, auth, taskId, configId] = params as [ - string, - string | null, - string | null, - string, - string, - ] - const row = rows.get(key(taskId, configId)) - if (!row) return { rowsAffected: 0 } - row.url = url - row.token = token - row.authentication = auth - return { rowsAffected: 1 } + throw new Error('push config writes must use an atomic upsert') } - if (s.startsWith('INSERT INTO')) { + if (s.startsWith('INSERT INTO') && s.includes('ON CONFLICT')) { + upsertCalls += 1 const [taskId, configId, url, token, auth] = params as [ string, string, @@ -290,6 +468,7 @@ function makePushAdapter(): SqlAdapter & { rows: Map } { string | null, string | null, ] + await beforeUpsert?.() rows.set(key(taskId, configId), { task_id: taskId, config_id: configId, @@ -366,6 +545,29 @@ describe('SqlPushNotificationStore', () => { expect((await store.get('t1', 'cfg1'))?.url).toBe('https://b.example/h') }) + it('uses one atomic upsert when two workers register the same config concurrently', async () => { + let entered = 0 + let release!: () => void + const bothEntered = new Promise((resolve) => { release = resolve }) + const db = makePushAdapter(async () => { + entered += 1 + if (entered === 2) release() + await bothEntered + }) + const first = new SqlPushNotificationStore(db) + const second = new SqlPushNotificationStore(db) + + await Promise.all([ + first.set('t1', cfg({ url: 'https://a.example/h' })), + second.set('t1', cfg({ url: 'https://b.example/h' })), + ]) + + expect(db.upsertCalls).toBe(2) + expect(db.rows.size).toBe(1) + expect(['https://a.example/h', 'https://b.example/h']) + .toContain((await first.get('t1', 'cfg1'))?.url) + }) + it('list returns all configs for one task; other tasks isolated', async () => { const db = makePushAdapter() const store = new SqlPushNotificationStore(db) diff --git a/tests/a2a-lifecycle.test.ts b/tests/a2a-lifecycle.test.ts index 30fcbe8..ce981cd 100644 --- a/tests/a2a-lifecycle.test.ts +++ b/tests/a2a-lifecycle.test.ts @@ -112,6 +112,7 @@ describe('A2A lifecycle recovery and ownership', () => { x402: { operatorAddress, chainId: 1, + demoMode: true, paymentProtocolVersion: 1, verifySigner: async () => true, }, @@ -198,6 +199,15 @@ describe('A2A lifecycle recovery and ownership', () => { return this.inner.compareAndSet(expected, next) } + compareAndSetExecution( + expected: Task, + next: Task, + requestId: string, + now: number, + ) { + return this.inner.compareAndSetExecution(expected, next, requestId, now) + } + async delete(id: string) { return this.inner.delete(id) } @@ -279,6 +289,9 @@ describe('A2A lifecycle recovery and ownership', () => { get(id: string) { return this.inner.get(id) } put(task: Task) { return this.inner.put(task) } compareAndSet(expected: Task, next: Task) { return this.inner.compareAndSet(expected, next) } + compareAndSetExecution(expected: Task, next: Task, requestId: string, now: number) { + return this.inner.compareAndSetExecution(expected, next, requestId, now) + } delete(id: string) { return this.inner.delete(id) } async createIfAbsent(task: Task) { diff --git a/tests/a2a-payment-races.test.ts b/tests/a2a-payment-races.test.ts index c1817b1..457cca4 100644 --- a/tests/a2a-payment-races.test.ts +++ b/tests/a2a-payment-races.test.ts @@ -476,6 +476,8 @@ describe('A2A payment ownership races', () => { } return innerStore.compareAndSet(expected, next) }, + compareAndSetExecution: (expected, next, requestId, now) => + innerStore.compareAndSetExecution(expected, next, requestId, now), } let settlements = 0 const operations = new MemoryPaymentOperations({ @@ -557,6 +559,8 @@ describe('A2A payment ownership races', () => { } return transitioned }, + compareAndSetExecution: (expected, next, requestId, now) => + innerStore.compareAndSetExecution(expected, next, requestId, now), } let sandboxDrained!: () => void const sandboxReady = new Promise((resolve) => { sandboxDrained = resolve }) @@ -783,6 +787,8 @@ describe('A2A payment ownership races', () => { } return innerStore.compareAndSet(expected, next) }, + compareAndSetExecution: (expected, next, requestId, now) => + innerStore.compareAndSetExecution(expected, next, requestId, now), } let settlementAttempts = 0 let recoveryAttempts = 0 diff --git a/tests/integration-x402.test.ts b/tests/integration-x402.test.ts index 9bd4f7d..b533ccb 100644 --- a/tests/integration-x402.test.ts +++ b/tests/integration-x402.test.ts @@ -15,7 +15,7 @@ * EIP-712 signer address and asserts it matches the commitment * 4. Gateway resolves the agent, rate-limits the consumer, filters for * injection, gets a sandbox, streams the response, records usage, - * and fires settlePayment. + * and settles through the durable version 2 payment operation. * 5. Consumer parses the SSE stream back to a string. * * Every step uses real code — real signatures (not fixtures), real Hono @@ -41,6 +41,8 @@ import type { } from '../src/types' import { MemoryNonceStore } from '../src/nonce-store' import { MemoryRateLimitStore } from '../src/rate-limit' +import { MemoryPaymentOperations } from '../src/payment-operations' +import { MemoryPaymentRecoveryStore } from '../src/payment-recovery' // ----- Domain constants (mirror the Tangle ShieldedCredits contract shape) ----- @@ -184,6 +186,19 @@ function buildHarness(chunks = ['Hello', ', ', 'world!']): Harness { const usage: GatewayUsageEvent[] = [] const settlements: Array<{ payment: PaymentResult; cost: number }> = [] let verifyCalls = 0 + const paymentOperations = new MemoryPaymentOperations({ + onSettle: async (operation, input) => { + settlements.push({ + payment: { + method: 'x402', + consumerId: operation.authorizationId, + requestId: operation.acquiredByRequestId, + }, + cost: input.totalCostUsd, + }) + }, + onReclaim: async () => undefined, + }) const agent: AgentMeta = { id: 'agent_production', @@ -202,17 +217,19 @@ function buildHarness(chunks = ['Hello', ', ', 'world!']): Harness { resolveAgent: async (slug) => (slug === agent.slug ? agent : null), getSandbox: async () => new ReplySandbox(chunks), recordUsage: async (evt) => { usage.push(evt) }, - settlePayment: async (payment, cost) => { settlements.push({ payment, cost }) }, x402: { operatorAddress: OPERATOR_ADDRESS, chainId: CHAIN_ID, creditsAddress: CREDITS_ADDRESS, demoMode: false, // PRODUCTION PATH — verifySigner is authoritative + paymentProtocolVersion: 2, + paymentOperations, verifySigner: async (payload) => { verifyCalls += 1 return verifySignerOnChain(payload) }, }, + paymentRecovery: { store: new MemoryPaymentRecoveryStore() }, nonceStore: new MemoryNonceStore(), rateLimitStore: new MemoryRateLimitStore(), }) @@ -265,7 +282,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo beforeEach(() => { harness = buildHarness() }) - it('happy path: consumer signs → gateway verifies signer address → sandbox streams → settlement fires', async () => { + it('happy path: consumer signs → gateway verifies signer address → sandbox streams → durable settlement fires', async () => { const spendAuth = await signSpendAuth({ consumerPrivateKey: harness.consumerPrivateKey, amount: FUNDED_REQUEST_AMOUNT, @@ -304,7 +321,7 @@ describe('x402 end-to-end — real EIP-712 signatures, real gateway, real sandbo expect(harness.usage[0].ownerEarnedUsd).toBeCloseTo(harness.usage[0].totalCostUsd * 0.8, 10) expect(harness.usage[0].platformFeeUsd).toBeCloseTo(harness.usage[0].totalCostUsd * 0.2, 10) - // Settlement callback fired with the x402 method + // Durable settlement fired with the x402 method expect(harness.settlements).toHaveLength(1) expect(harness.settlements[0].payment.method).toBe('x402') expect(harness.settlements[0].cost).toBe(harness.usage[0].totalCostUsd) diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index c471158..2cee600 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -12,7 +12,7 @@ import { MemoryPaymentRecoveryStore, type PaymentRecoveryRecord } from '../src/p import { recoverPayment } from '../src/payment-recovery-worker' import { verifyMpp } from '../src/verify' import type { AgentMeta, GatewayConfig, SandboxStreamEvent } from '../src/types' -import type { Task } from '../src/a2a/types' +import { A2A_ERROR_CODES, type Task } from '../src/a2a/types' import type { MppConfig } from '../src/types' const operatorAddress = '0x1111111111111111111111111111111111111111' @@ -89,7 +89,7 @@ function durableConfig( } describe('PR #11 production regressions', () => { - it('claims one terminal webhook when cancellation races settlement on two workers', async () => { + it('claims one terminal webhook when cancellation races fenced settlement on two workers', async () => { const taskStore = new InMemoryTaskStore() const pushStore = new InMemoryPushNotificationStore() const nonceStore = new MemoryNonceStore() @@ -189,18 +189,18 @@ describe('PR #11 production regressions', () => { }), }) expect(cancel.status).toBe(200) - expect((await cancel.json() as { result?: { status?: { state?: string } } }) - .result?.status?.state).toBe('canceled') + expect((await cancel.json() as { error?: { code?: number } }).error?.code) + .toBe(A2A_ERROR_CODES.TASK_NOT_CANCELABLE) releaseSandbox() const runnerResponse = await running expect(runnerResponse.status).toBe(200) expect((await runnerResponse.json() as { result?: { status?: { state?: string } } }) - .result?.status?.state).toBe('canceled') + .result?.status?.state).toBe('completed') expect(settlements).toBe(1) expect(deliveries).toBe(1) expect(receivedTaskIds).toEqual(['pr11-push-race']) - expect((await taskStore.get('pr11-push-race'))?.status.state).toBe('canceled') + expect((await taskStore.get('pr11-push-race'))?.status.state).toBe('completed') }) it('rejects production v1 authorization callbacks without durable recovery', () => { @@ -265,6 +265,136 @@ describe('PR #11 production regressions', () => { expect(providerRecoveries).toBe(1) }) + it('rejects production v1 settlement before nonce claim or provider mutation', async () => { + const nonceStore = new MemoryNonceStore() + let sandboxRuns = 0 + let settlementCalls = 0 + const app = new Hono() + app.route('/v1/agents', createAgentGateway({ + ...durableConfig({ nonceStore }), + getSandbox: async () => ({ + async *streamPrompt() { + sandboxRuns += 1 + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + settlePayment: async () => { + settlementCalls += 1 + throw new Error('provider acknowledgement lost') + }, + x402: { + operatorAddress, + chainId: 1, + verifySigner: async () => true, + paymentProtocolVersion: 1, + }, + paymentRecovery: undefined, + })) + + const response = await app.request('/v1/agents/pr11/chat/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9013'), + }, + body: JSON.stringify({ messages: [{ role: 'user', content: 'must not run' }] }), + }) + + expect(response.status).toBe(402) + expect(await nonceStore.hasSeen(`${commitment}:9013`)).toBe(false) + expect(sandboxRuns).toBe(0) + expect(settlementCalls).toBe(0) + }) + + it('requires a webhook HMAC secret for production push delivery', () => { + expect(() => createAgentGateway(durableConfig({ + x402: { + operatorAddress, + chainId: 1, + demoMode: false, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations({ onReclaim: async () => undefined }), + }, + a2a: { + pushStore: new InMemoryPushNotificationStore(), + authorizeTaskAccess: async () => true, + }, + }))).toThrow(/production A2A push requires a webhookSecret/) + }) + + it('does not send an unsigned webhook if a production secret disappears at runtime', async () => { + const taskStore = new InMemoryTaskStore() + const pushStore = new InMemoryPushNotificationStore() + let sandboxStarted!: () => void + const sandboxReady = new Promise((resolve) => { sandboxStarted = resolve }) + let releaseSandbox!: () => void + const sandboxReleased = new Promise((resolve) => { releaseSandbox = resolve }) + let deliveries = 0 + const config = durableConfig({ + getSandbox: async () => ({ + async *streamPrompt() { + sandboxStarted() + await sandboxReleased + yield { type: 'sandbox.usage', data: { usage: usage() } } + }, + }), + x402: { + operatorAddress, + chainId: 1, + demoMode: false, + verifySigner: async () => true, + paymentProtocolVersion: 2, + paymentOperations: new MemoryPaymentOperations({ onReclaim: async () => undefined }), + }, + a2a: { + taskStore, + pushStore, + webhookSecret: 'runtime-secret', + pushFetcher: async () => { + deliveries += 1 + return new Response('ok') + }, + pushUrlValidator: async () => true, + authorizeTaskAccess: async () => true, + }, + }) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(config)) + const request = app.request('/v1/agents/pr11', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Payment-Signature': paymentHeader('9014'), + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'message/send', + params: { + message: { + kind: 'message', + role: 'user', + taskId: 'runtime-secret-task', + contextId: 'runtime-secret-context', + messageId: 'runtime-secret-message', + parts: [{ kind: 'text', text: 'run' }], + }, + }, + }), + }) + await sandboxReady + await pushStore.set('runtime-secret-task', { id: 'cfg', url: 'https://hook.example/terminal' }) + config.a2a!.webhookSecret = undefined + releaseSandbox() + + const response = await request + expect(response.status).toBe(200) + expect((await response.json() as { result?: { status?: { state?: string } } }) + .result?.status?.state).toBe('completed') + expect(deliveries).toBe(0) + }) + it('does not mark payment execution before sandbox acquisition succeeds', async () => { const operations = new MemoryPaymentOperations({ onReclaim: async () => undefined }) let sandboxReady = false @@ -427,20 +557,41 @@ describe('PR #11 production regressions', () => { expect((await taskStore.get('pr11-cancel-race'))?.status.state).toBe('canceled') }) - it('cancels an execution already fenced by another worker', async () => { - const taskStore = new InMemoryTaskStore() - let executionClaimed!: () => void - const executionReady = new Promise((resolve) => { executionClaimed = resolve }) - let releaseExecution!: () => void - const executionReleased = new Promise((resolve) => { releaseExecution = resolve }) + it('rejects cross-worker cancellation after the durable fence and before provider start', async () => { + const innerStore = new InMemoryTaskStore() + let fenceWritten!: () => void + const fenceReady = new Promise((resolve) => { fenceWritten = resolve }) + let releaseFence!: () => void + const fenceReleased = new Promise((resolve) => { releaseFence = resolve }) + let blocked = false + const taskStore: TaskStore = { + get: (id) => innerStore.get(id), + put: (task) => innerStore.put(task), + createIfAbsent: (task) => innerStore.createIfAbsent(task), + delete: (id) => innerStore.delete(id), + async compareAndSet(expected, next) { + const won = await innerStore.compareAndSet(expected, next) + if (won && !blocked && next.metadata?.gatewayExecution !== undefined) { + blocked = true + fenceWritten() + await fenceReleased + } + return won + }, + compareAndSetExecution: (expected, next, requestId, now) => + innerStore.compareAndSetExecution(expected, next, requestId, now), + } + let providerStarted = false + let releaseProvider!: () => void + const providerReleased = new Promise((resolve) => { releaseProvider = resolve }) const makeConfig = (worker: 'runner' | 'canceler'): GatewayConfig => ({ resolveAgent: async () => agent, getSandbox: async () => ({ async *streamPrompt() { if (worker === 'runner') { - executionClaimed() - await executionReleased + providerStarted = true + await providerReleased } yield { type: 'sandbox.usage', data: { usage: usage() } } }, @@ -484,7 +635,7 @@ describe('PR #11 production regressions', () => { }, }), }) - await executionReady + await fenceReady const cancel = await canceler.request('/v1/agents/pr11', { method: 'POST', @@ -496,14 +647,17 @@ describe('PR #11 production regressions', () => { params: { id: 'pr11-active-cancel-race' }, }), }) - const cancelBody = await cancel.json() as { result?: { status?: { state?: string } } } + const cancelBody = await cancel.json() as { error?: { code?: number } } expect(cancel.status).toBe(200) - expect(cancelBody.result?.status?.state).toBe('canceled') + expect(cancelBody.error?.code).toBe(A2A_ERROR_CODES.TASK_NOT_CANCELABLE) + expect(providerStarted).toBe(false) + expect((await taskStore.get('pr11-active-cancel-race'))?.status.state).toBe('working') - releaseExecution() + releaseFence() + releaseProvider() const runningResponse = await running const runningBody = await runningResponse.json() as { result?: { status?: { state?: string } } } - expect(runningBody.result?.status?.state).toBe('canceled') + expect(runningBody.result?.status?.state).toBe('completed') expect((await taskStore.get('pr11-active-cancel-race'))?.metadata?.gatewayExecution) .toBeUndefined() }) @@ -528,6 +682,8 @@ describe('PR #11 production regressions', () => { } return innerStore.compareAndSet(expected, next) }, + compareAndSetExecution: (expected, next, requestId, now) => + innerStore.compareAndSetExecution(expected, next, requestId, now), } const app = new Hono() app.route('/v1/agents', createAgentGateway(durableConfig({ @@ -578,7 +734,7 @@ describe('PR #11 production regressions', () => { .result?.status?.state).toBe('completed') }) - it('retries cancellation after a heartbeat changes the task row', async () => { + it('rejects cancellation when a heartbeat wins the cancellation CAS race', async () => { const innerStore = new InMemoryTaskStore() let cancelAttempts = 0 const taskStore: TaskStore = { @@ -603,6 +759,8 @@ describe('PR #11 production regressions', () => { } return innerStore.compareAndSet(expected, next) }, + compareAndSetExecution: (expected, next, requestId, now) => + innerStore.compareAndSetExecution(expected, next, requestId, now), } await taskStore.put({ kind: 'task', @@ -611,11 +769,6 @@ describe('PR #11 production regressions', () => { status: { state: 'working', timestamp: new Date().toISOString() }, metadata: { gatewayOrigin: { version: 1, agentId: agent.id, agentSlug: agent.slug }, - gatewayExecution: { - version: 1, - requestId: 'heartbeat-owner', - lease: { id: 'heartbeat-owner', expiresAt: Date.now() + 60_000 }, - }, }, }) const app = new Hono() @@ -636,9 +789,10 @@ describe('PR #11 production regressions', () => { const body = await response.json() as { result?: Task; error?: unknown } expect(response.status).toBe(200) - expect(body.error).toBeUndefined() - expect(body.result?.status.state).toBe('canceled') - expect(cancelAttempts).toBe(2) + expect((body.error as { code?: number } | undefined)?.code) + .toBe(A2A_ERROR_CODES.TASK_NOT_CANCELABLE) + expect(body.result).toBeUndefined() + expect(cancelAttempts).toBe(1) }) it('quotes retained A2A history before charging a continuation', async () => { @@ -1179,7 +1333,7 @@ describe('PR #11 production regressions', () => { get: async () => undefined, delete: async () => undefined, }, - webhookSecret: undefined, + webhookSecret: 'test-webhook-secret', fetcher: fetcher as unknown as typeof fetch, }) @@ -1188,6 +1342,176 @@ describe('PR #11 production regressions', () => { expect(fetcher.mock.calls[0]?.[1]).toMatchObject({ redirect: 'manual' }) }) + it('does not consume a rejected push claim and delivers on a valid retry', async () => { + const { deliverPushNotifications } = await import('../src/a2a/push-notifications') + const task: Task = { + kind: 'task', + id: 'push-validation-retry', + contextId: 'push-validation-context', + status: { state: 'completed', timestamp: new Date().toISOString() }, + } + const pushStore = new InMemoryPushNotificationStore() + await pushStore.set(task.id, { id: 'retry', url: 'https://receiver.example/hook' }) + const webhook = new Hono() + let deliveries = 0 + webhook.post('/hook', async (context) => { + deliveries += 1 + await context.req.json() + return context.text('ok') + }) + const fetcher = async ( + input: string | URL | Request, + init?: RequestInit, + ): Promise => webhook.fetch( + input instanceof Request ? new Request(input, init) : new Request(input.toString(), init), + ) + let policyAllows = false + let claims = 0 + + const deliver = () => deliverPushNotifications({ + task, + store: pushStore, + webhookSecret: 'test-webhook-secret', + requireUrlValidator: true, + urlValidator: async () => policyAllows, + claimDelivery: async () => { + claims += 1 + return true + }, + fetcher, + }) + + const rejected = await deliver() + expect(rejected[0]).toMatchObject({ ok: false, error: 'push notification URL was rejected by host policy' }) + expect(claims).toBe(0) + expect(deliveries).toBe(0) + + policyAllows = true + const delivered = await deliver() + expect(delivered[0]).toMatchObject({ ok: true, status: 200 }) + expect(claims).toBe(1) + expect(deliveries).toBe(1) + }) + + it('rejects unsigned direct push delivery before reading configs or fetching', async () => { + const { deliverPushNotifications } = await import('../src/a2a/push-notifications') + const list = vi.fn(async () => [{ id: 'unsigned', url: 'https://hook.example/terminal' }]) + const fetcher = vi.fn(async () => new Response('unexpected', { status: 200 })) + + await expect(deliverPushNotifications({ + task: { + kind: 'task', + id: 'unsigned-task', + contextId: 'unsigned-context', + status: { state: 'completed', timestamp: new Date().toISOString() }, + }, + store: { + list, + set: async () => undefined, + get: async () => undefined, + delete: async () => undefined, + }, + webhookSecret: ' ', + fetcher: fetcher as unknown as typeof fetch, + })).rejects.toThrow('non-empty webhookSecret') + + expect(list).not.toHaveBeenCalled() + expect(fetcher).not.toHaveBeenCalled() + }) + + it('closes a malformed execution marker through tasks/get and preserves payment recovery', async () => { + const taskStore = new InMemoryTaskStore() + const task: Task = { + kind: 'task', + id: 'malformed-execution-get', + contextId: 'malformed-execution-context', + status: { state: 'working', timestamp: new Date().toISOString() }, + metadata: { + gatewayExecution: { + version: 1, + requestId: 'worker-a', + lease: { id: 'worker-a' }, + }, + gatewayPaymentRecovery: { version: 1, id: 'payment-recovery-get' }, + }, + } + await taskStore.put(task) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + a2a: { taskStore }, + }))) + + const response = await app.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'tasks/get', + params: { id: task.id }, + }), + }) + const body = await response.json() as { result?: Task; error?: unknown } + + expect(response.status).toBe(200) + expect(body.error).toBeUndefined() + expect(body.result?.status.state).toBe('failed') + expect(body.result?.metadata?.gatewayExecution).toBeUndefined() + expect(body.result?.metadata?.gatewayExecutionRecovery).toMatchObject({ + error: expect.stringContaining('malformed'), + }) + expect(body.result?.metadata?.gatewayPaymentRecovery) + .toEqual({ version: 1, id: 'payment-recovery-get' }) + expect((await taskStore.get(task.id))?.status.state).toBe('failed') + }) + + it('closes a malformed execution marker through tasks/resubscribe and emits final state', async () => { + const taskStore = new InMemoryTaskStore() + const task: Task = { + kind: 'task', + id: 'malformed-execution-resubscribe', + contextId: 'malformed-execution-context', + status: { state: 'working', timestamp: new Date().toISOString() }, + metadata: { + gatewayExecution: { version: 1, requestId: 'worker-b', lease: null }, + gatewayPaymentRecovery: { version: 1, id: 'payment-recovery-resubscribe' }, + }, + } + await taskStore.put(task) + const app = new Hono() + app.route('/v1/agents', createAgentGateway(durableConfig({ + a2a: { taskStore }, + }))) + + const response = await app.request('/v1/agents/pr11', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'tasks/resubscribe', + params: { id: task.id }, + }), + }) + const eventLine = (await response.text()).split('\n') + .find((line) => line.startsWith('data: ')) + const event = eventLine + ? JSON.parse(eventLine.slice('data: '.length)) as { result?: { + final?: boolean + status?: Task['status'] + taskId?: string + } } + : undefined + + expect(response.status).toBe(200) + expect(event?.result?.taskId).toBe(task.id) + expect(event?.result?.status?.state).toBe('failed') + expect(event?.result?.final).toBe(true) + expect((await taskStore.get(task.id))?.metadata?.gatewayExecution).toBeUndefined() + expect((await taskStore.get(task.id))?.metadata?.gatewayPaymentRecovery) + .toEqual({ version: 1, id: 'payment-recovery-resubscribe' }) + }) + it('rejects direct private push destinations before fetch', async () => { const { deliverPushNotifications } = await import('../src/a2a/push-notifications') const fetcher = vi.fn(async () => new Response('unexpected', { status: 200 })) @@ -1204,7 +1528,7 @@ describe('PR #11 production regressions', () => { get: async () => undefined, delete: async () => undefined, }, - webhookSecret: undefined, + webhookSecret: 'test-webhook-secret', fetcher: fetcher as unknown as typeof fetch, }) @@ -1237,7 +1561,7 @@ describe('PR #11 production regressions', () => { get: async () => undefined, delete: async () => undefined, }, - webhookSecret: undefined, + webhookSecret: 'test-webhook-secret', fetcher: fetcher as unknown as typeof fetch, }) From b09a6c52863a5fafafd1525348c18d833fdf5ce4 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 18:23:17 -0600 Subject: [PATCH 30/32] refactor(gateway): split dispatch ownership modules --- src/dispatch-authorization.ts | 437 +++++++ src/dispatch-payment-recovery.ts | 248 ++++ src/dispatch-payment.ts | 425 +++++++ src/dispatch-pricing.ts | 108 ++ src/dispatch-sandbox.ts | 422 +++++++ src/dispatch-settlement.ts | 139 +++ src/dispatch-types.ts | 81 ++ src/dispatch.ts | 1831 +--------------------------- tests/dispatch-module-size.test.ts | 23 + 9 files changed, 1918 insertions(+), 1796 deletions(-) create mode 100644 src/dispatch-authorization.ts create mode 100644 src/dispatch-payment-recovery.ts create mode 100644 src/dispatch-payment.ts create mode 100644 src/dispatch-pricing.ts create mode 100644 src/dispatch-sandbox.ts create mode 100644 src/dispatch-settlement.ts create mode 100644 src/dispatch-types.ts create mode 100644 tests/dispatch-module-size.test.ts diff --git a/src/dispatch-authorization.ts b/src/dispatch-authorization.ts new file mode 100644 index 0000000..a3713b2 --- /dev/null +++ b/src/dispatch-authorization.ts @@ -0,0 +1,437 @@ +import type { Context } from 'hono' + +import { filterConsumerMessagesStrict } from './filter' +import { type RequestContext, generateRequestId } from './observer' +import { type GatewayState, type AuthorizedRequest } from './dispatch-types' +import { checkRateLimit } from './rate-limit' +import { + maximumBillableInputTokens, + requiredX402Amount, +} from './dispatch-pricing' +import type { + ApiKeyInfo, + ChatMessage, + GatewayConfig, + PaymentMethod, + SandboxExecutionBudget, +} from './types' +import { + defaultVerifyApiKey, + isApiKeyAuthEnabled, + isMppAuthEnabled, + mppPaymentPayload, + mppPaymentCredential, + verifyMppCredential, + verifyX402, +} from './verify' + +/** + * Resolve the agent, then run the full pre-dispatch pipeline: payment + + * rate-limit + injection filter + user-message extraction + optional + * `authorizeConsumer` hook. Returns the success record on the happy path + * or a fully-formed `Response` (402/404/429/400/403) on any short-circuit. + * + * Body parsing is the caller's responsibility — different wire formats + * (OpenAI chat completions vs A2A JSON-RPC) have different envelopes; both + * still ultimately produce a `ChatMessage[]`. + */ +export async function authenticateAndGuard( + c: Context, + slug: string, + messages: ChatMessage[], + config: GatewayConfig, + state: GatewayState, + requestedMaxOutputTokens?: number, +): Promise { + const startMs = Date.now() + const requestId = generateRequestId() + const ctx: RequestContext = { requestId, agentSlug: slug, startMs } + await state.obs?.onRequestStart?.(ctx) + + const agent = await config.resolveAgent(slug) + if (!agent || !agent.enabled) { + return c.json({ error: { message: 'Agent not found', type: 'not_found' } }, 404) + } + if (!messages?.length) { + return c.json( + { error: { message: 'messages array required', type: 'invalid_request' } }, + 400, + ) + } + + const maxOutputTokens = requestedMaxOutputTokens ?? state.defaultOutputTokens + if ( + !Number.isInteger(maxOutputTokens) || + maxOutputTokens <= 0 || + maxOutputTokens > state.maxOutputTokens + ) { + return c.json( + { + error: { + message: `max_tokens must be an integer between 1 and ${state.maxOutputTokens}`, + type: 'invalid_request', + code: 'invalid_max_tokens', + }, + }, + 400, + ) + } + + // Quote the maximum UTF-8 input plus every hidden provider cost before + // verification. The verifier must remain read-only at this point. + const { messages: filtered, injectionWarnings } = filterConsumerMessagesStrict( + messages, + state.maxLen, + ) + const userMessage = filtered + .filter((m) => m.role === 'user') + .map((m) => m.content) + .join('\n\n') + if (!userMessage) { + return c.json( + { error: { message: 'No user message provided', type: 'invalid_request' } }, + 400, + ) + } + + let requiredPaymentAmount: bigint + const messageInputBound = maximumBillableInputTokens(agent, filtered) + let maxInputTokens = messageInputBound + if (config.inputTokenBound) { + let configuredBound: number + try { + configuredBound = config.inputTokenBound({ agent, messages: filtered }) + } catch { + return c.json( + { + error: { + message: 'Agent input token bound is unavailable', + type: 'server_error', + code: 'input_token_bound_unavailable', + }, + }, + 503, + ) + } + if (!Number.isSafeInteger(configuredBound) || configuredBound < messageInputBound) { + return c.json( + { + error: { + message: 'Agent input token bound is invalid', + type: 'server_error', + code: 'invalid_input_token_bound', + }, + }, + 503, + ) + } + maxInputTokens = configuredBound + } + const maxReasoningTokens = state.maxReasoningTokens + const maxToolTokens = state.maxToolTokens + const maxToolCalls = state.maxToolCalls + const maxProviderCostUsd = state.maxProviderCostUsd ?? + (maxInputTokens + maxOutputTokens + maxReasoningTokens + maxToolTokens) * agent.pricePerTokenUsd + const executionBudget: SandboxExecutionBudget = { + maxInputTokens, + maxOutputTokens, + maxReasoningTokens, + maxToolTokens, + maxToolCalls, + maxProviderCostUsd, + } + try { + requiredPaymentAmount = requiredX402Amount( + agent.pricePerTokenUsd, + maxInputTokens, + maxOutputTokens, + config.x402.currencyDecimals, + maxReasoningTokens, + maxToolTokens, + maxProviderCostUsd, + ) + } catch { + return c.json( + { + error: { + message: 'Agent payment configuration is invalid', + type: 'server_error', + code: 'invalid_payment_configuration', + }, + }, + 503, + ) + } + + // Payment / auth. + const spendAuthHeader = c.req.header('X-Payment-Signature') + const authHeader = c.req.header('Authorization') ?? '' + let consumerId: string | null = null + let paymentMethod: PaymentMethod = 'none' + let keyInfo: ApiKeyInfo | null = null + let x402Payload: Record | null = null + let paymentNonceKey: string | undefined + let mppMethod: string | undefined + let mppCredential: string | undefined + let mppPaymentIdentity: string | undefined + + if (spendAuthHeader) { + const signer = await verifyX402( + spendAuthHeader, + config.x402, + state.nonceStore, + requiredPaymentAmount, + false, + ) + if (!signer) { + await state.obs?.onAuthFailure?.(ctx, { + method: 'x402', + code: 'invalid_spend_auth', + httpStatus: 402, + }) + return c.json( + { + error: { + message: 'Invalid X-Payment-Signature', + type: 'authentication_error', + code: 'invalid_spend_auth', + required_amount: requiredPaymentAmount.toString(), + currency_decimals: config.x402.currencyDecimals ?? 6, + }, + }, + { + status: 402, + headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': requestId }, + }, + ) + } + x402Payload = JSON.parse(spendAuthHeader) as Record + paymentNonceKey = `${String(x402Payload.commitment).toLowerCase()}:${BigInt(String(x402Payload.nonce)).toString()}` + consumerId = signer + paymentMethod = 'x402' + } else if (isMppAuthEnabled(config) && authHeader.toLowerCase().startsWith('payment ')) { + const authenticated = await verifyMppCredential( + authHeader, + config.mpp!, + config.x402, + state.nonceStore, + requiredPaymentAmount, + false, + ) + if (!authenticated) { + const realm = config.mpp!.realm + const method = config.mpp!.method ?? 'blueprintevm' + await state.obs?.onAuthFailure?.(ctx, { + method: 'mpp', + code: 'invalid_mpp_credential', + httpStatus: 401, + }) + return c.json( + { + error: { + message: 'Invalid Payment credential', + type: 'authentication_error', + code: 'invalid_mpp_credential', + }, + }, + { + status: 401, + headers: { + 'WWW-Authenticate': `Payment realm="${realm}", method="${method}"`, + 'X-Request-Id': requestId, + }, + }, + ) + } + consumerId = authenticated.consumerId + paymentMethod = 'mpp' + mppMethod = authHeader.match(/^Payment\s+(\S+)\s+/i)?.[1]?.toLowerCase() + mppCredential = mppPaymentCredential(authHeader) + mppPaymentIdentity = authenticated.paymentIdentity + x402Payload = mppPaymentPayload(authHeader) ?? null + paymentNonceKey = authenticated.replayKey + } else if (authHeader.startsWith('Bearer ')) { + const verify = config.verifyApiKey ?? (config.x402.demoMode ? defaultVerifyApiKey : null) + if (!verify || !isApiKeyAuthEnabled(config)) { + await state.obs?.onAuthFailure?.(ctx, { + method: 'apikey', + code: 'api_keys_not_configured', + httpStatus: 401, + }) + return c.json( + { error: { message: 'API key authentication is not configured', type: 'authentication_error' } }, + { status: 401, headers: { 'X-Request-Id': requestId } }, + ) + } + const key = await verify(authHeader) + if (!key) { + await state.obs?.onAuthFailure?.(ctx, { + method: 'apikey', + code: 'invalid_api_key', + httpStatus: 401, + }) + return c.json( + { error: { message: 'Invalid API key', type: 'authentication_error' } }, + { status: 401, headers: { 'X-Request-Id': requestId } }, + ) + } + if (key.scopes && key.scopes.length > 0 && !key.scopes.includes(state.requiredScope)) { + await state.obs?.onAuthFailure?.(ctx, { + method: 'apikey', + code: 'insufficient_scope', + httpStatus: 403, + }) + return c.json( + { + error: { + message: `API key missing required scope: ${state.requiredScope}`, + type: 'forbidden', + code: 'insufficient_scope', + }, + }, + { status: 403, headers: { 'X-Request-Id': requestId } }, + ) + } + consumerId = key.consumerId + paymentMethod = 'apikey' + keyInfo = key + } else { + await state.obs?.onAuthFailure?.(ctx, { + method: 'none', + code: 'payment_required', + httpStatus: 402, + }) + const methods: string[] = ['x402'] + if (isMppAuthEnabled(config)) methods.push('mpp') + if (isApiKeyAuthEnabled(config)) methods.push('api_key') + const headers: Record = { + 'X-Payment-Required': methods.join(', '), + 'X-Request-Id': requestId, + } + if (isMppAuthEnabled(config) && config.mpp) { + headers['WWW-Authenticate'] = + `Payment realm="${config.mpp.realm}", method="${config.mpp.method ?? 'blueprintevm'}"` + } + return c.json( + { + error: { + message: 'Payment required', + type: 'payment_required', + payment_methods: methods, + x402: { + operator: config.x402.operatorAddress, + chain_id: config.x402.chainId, + credits_address: config.x402.creditsAddress, + required_amount: requiredPaymentAmount.toString(), + currency_decimals: config.x402.currencyDecimals ?? 6, + max_output_tokens: maxOutputTokens, + }, + ...(isMppAuthEnabled(config) && config.mpp + ? { mpp: { realm: config.mpp.realm, method: config.mpp.method ?? 'blueprintevm' } } + : {}), + ...(isApiKeyAuthEnabled(config) + ? { + api_key: { + purchase_url: config.baseUrl + ? `${config.baseUrl}/agents/${slug}/api-keys` + : undefined, + }, + } + : {}), + }, + }, + { status: 402, headers }, + ) + } + + // Rate limit. + const effectiveRateLimit = keyInfo?.rateLimitPerMinute + ? { limit: keyInfo.rateLimitPerMinute, windowSeconds: 60 } + : state.globalRateLimit + const rl = await checkRateLimit(consumerId, effectiveRateLimit, state.rateLimitStore) + if (!rl.allowed) { + await state.obs?.onRateLimited?.(ctx, { + consumerId: consumerId, + retryAfterSeconds: rl.retryAfterSeconds ?? 60, + }) + return c.json( + { + error: { + message: 'Rate limit exceeded', + type: 'rate_limit_error', + retry_after: rl.retryAfterSeconds, + }, + }, + { + status: 429, + headers: { + 'Retry-After': String(rl.retryAfterSeconds ?? 60), + 'X-Request-Id': requestId, + }, + }, + ) + } + + // Reject or report injection only after authentication so observer events + // retain the authenticated consumer identity. + if (injectionWarnings.length > 0) { + await state.obs?.onInjectionDetected?.(ctx, { + consumerId: consumerId, + patterns: injectionWarnings, + blocked: !!config.blockInjection, + }) + if (config.blockInjection) { + return c.json( + { + error: { + message: 'Request rejected: potential prompt injection detected', + type: 'content_policy_violation', + }, + }, + { status: 400, headers: { 'X-Request-Id': requestId } }, + ) + } + } + + if (config.authorizeConsumer) { + const authz = await config.authorizeConsumer(agent, { + method: paymentMethod, + consumerId: consumerId, + keyId: keyInfo?.keyId, + requestId, + }) + if (!authz.allow) { + return c.json( + { + error: { + message: authz.reason, + type: 'authorization_denied', + code: authz.code, + }, + }, + { status: 403, headers: { 'X-Request-Id': requestId } }, + ) + } + } + + return { + agent, + consumerId, + paymentMethod, + keyInfo, + userMessage, + rateLimitRemaining: rl.remaining, + requestId, + startMs, + maxOutputTokens, + executionBudget, + requiredPaymentAmount, + paymentPayload: x402Payload, + paymentNonceKey, + mppMethod, + mppCredential, + mppPaymentIdentity, + } +} + +export type { AuthorizedRequest, GatewayState } diff --git a/src/dispatch-payment-recovery.ts b/src/dispatch-payment-recovery.ts new file mode 100644 index 0000000..806ae2a --- /dev/null +++ b/src/dispatch-payment-recovery.ts @@ -0,0 +1,248 @@ +import { + PAYMENT_RECOVERY_VERSION, + PaymentRecoveryFenceError, + PaymentRecoveryReplayError, + recoveryTiming, + serializePaymentOperation, + updateOwnedPaymentRecovery, + type PaymentRecoveryRecord, + type PaymentRecoveryTarget, + type PaymentSettlementBasis, +} from './payment-recovery' +import type { GatewayConfig, SandboxUsageReceipt } from './types' +import type { + AuthorizedRequest, + PaymentClaimHooks, +} from './dispatch-types' + +export function paymentAuthorizationContext(authz: AuthorizedRequest) { + return { + requestId: authz.requestId, + agentId: authz.agent.id, + requiredAmount: authz.requiredPaymentAmount, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + } +} + +export async function preparePaymentRecovery( + authz: AuthorizedRequest, + config: GatewayConfig, + payment: PaymentRecoveryTarget, + hooks: PaymentClaimHooks, +): Promise { + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const now = Date.now() + const fenceId = globalThis.crypto.randomUUID() + const leaseExpiresAt = now + recoveryTiming(recovery).staleRequestMs + const record: PaymentRecoveryRecord = { + version: PAYMENT_RECOVERY_VERSION, + id: payment.operationId, + revision: 0, + state: 'claiming', + payment, + attribution: { + requestId: authz.requestId, + agentId: authz.agent.id, + agentSlug: authz.agent.slug, + consumerId: authz.consumerId, + paymentMethod: authz.paymentMethod, + startMs: authz.startMs, + pricePerTokenUsd: authz.agent.pricePerTokenUsd, + platformFeePercent: authz.agent.platformFeePercent, + requiredAmount: authz.requiredPaymentAmount.toString(), + currencyDecimals: config.x402.currencyDecimals ?? 6, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + }, + workStarted: false, + usageRecorded: false, + attempts: 0, + nextAttemptAt: leaseExpiresAt, + lease: { id: fenceId, expiresAt: leaseExpiresAt }, + createdAt: now, + updatedAt: now, + } + if (!await recovery.store.createIfAbsent(record)) { + if ((await recovery.store.get(record.id))?.state === 'reconciled') { + throw new PaymentRecoveryReplayError(record.id) + } + throw new Error('payment recovery identity was already claimed') + } + authz.paymentRecoveryId = record.id + authz.paymentRecoveryFence = fenceId + await hooks.onRecoveryPrepared?.(record.id) +} + +export async function markRecoveryClaimed( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const fenceId = requirePaymentRecoveryFence(authz) + const now = Date.now() + const leaseExpiresAt = now + recoveryTiming(recovery).staleRequestMs + await updateRecovery(authz, config, (record) => ({ + ...record, + state: 'claimed', + payment: recoveryTarget(authz, record.payment), + lease: { id: fenceId, expiresAt: leaseExpiresAt }, + nextAttemptAt: leaseExpiresAt, + }), now) +} + +export function requirePaymentRecoveryFence(authz: AuthorizedRequest): string { + if (!authz.paymentRecoveryFence) { + throw new Error('payment recovery fence is unavailable') + } + return authz.paymentRecoveryFence +} + +async function updateRecovery( + authz: AuthorizedRequest, + config: GatewayConfig, + update: (record: PaymentRecoveryRecord) => PaymentRecoveryRecord, + now = Date.now(), +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return undefined + return updateOwnedPaymentRecovery( + recovery.store, + authz.paymentRecoveryId, + requirePaymentRecoveryFence(authz), + update, + now, + ) +} + +export function recoveryTarget( + authz: AuthorizedRequest, + current: PaymentRecoveryTarget, +): PaymentRecoveryTarget { + if (authz.paymentOperation) { + return { + kind: 'x402', + operationId: authz.paymentOperation.operationId, + operation: serializePaymentOperation(authz.paymentOperation), + } + } + if (authz.mppChargeOperation) { + return { + kind: 'mpp-charge', + method: authz.mppChargeOperation.method, + operationId: authz.mppChargeOperation.operationId, + operation: authz.mppChargeOperation, + } + } + return current +} + +export async function relinquishPaymentRecovery( + authz: AuthorizedRequest, + config: GatewayConfig, + nextAttemptAt: number, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId || !authz.paymentRecoveryFence) return + try { + await updateOwnedPaymentRecovery( + recovery.store, + authz.paymentRecoveryId, + authz.paymentRecoveryFence, + (record) => ({ ...record, lease: undefined, nextAttemptAt }), + ) + } catch (error) { + if (!(error instanceof PaymentRecoveryFenceError)) throw error + } +} + +export async function markRecoveryReleasing( + authz: AuthorizedRequest, + config: GatewayConfig, + reason: string, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + await updateRecovery(authz, config, (record) => ({ + ...record, + state: 'releasing', + payment: recoveryTarget(authz, record.payment), + reason, + nextAttemptAt: Date.now(), + })) +} + +export async function markRecoveryReconciled( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + const now = Date.now() + await updateRecovery(authz, config, (record) => ({ + ...record, + state: 'reconciled', + payment: recoveryTarget(authz, record.payment), + lease: undefined, + lastError: undefined, + nextAttemptAt: Number.MAX_SAFE_INTEGER, + reconciledAt: now, + }), now) +} + +export function assertX402V1SettlementSafe(authz: AuthorizedRequest, config: GatewayConfig): void { + if ( + authz.paymentMethod === 'x402' && + config.x402.paymentProtocolVersion !== 2 && + !config.x402.demoMode && + config.settlePayment + ) { + throw new Error( + 'production x402 version 1 cannot use settlePayment; ' + + 'use paymentProtocolVersion: 2 with paymentOperations', + ) + } +} + +export async function markRecoverySettling( + authz: AuthorizedRequest, + usage: SandboxUsageReceipt, + settlementBasis: PaymentSettlementBasis, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + await updateRecovery(authz, config, (record) => { + const next: PaymentRecoveryRecord = { + ...record, + state: 'settling', + payment: recoveryTarget(authz, record.payment), + workStarted: true, + settlementBasis, + nextAttemptAt: Date.now(), + } + // A quoted-ceiling settlement has no provider receipt. Keep the durable + // basis and original amount, then rebuild the synthetic accounting input + // on each retry instead of persisting a lossy floating-point surrogate. + if (settlementBasis !== 'quoted-ceiling' || record.usage !== undefined) { + next.usage = usage + } else { + delete next.usage + } + return next + }) +} + +export async function markRecoveryUsageRecorded( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + const recovery = config.paymentRecovery + if (!recovery || !authz.paymentRecoveryId) return + await updateRecovery(authz, config, (record) => ({ + ...record, + usageRecorded: true, + })) +} diff --git a/src/dispatch-payment.ts b/src/dispatch-payment.ts new file mode 100644 index 0000000..0dc8d70 --- /dev/null +++ b/src/dispatch-payment.ts @@ -0,0 +1,425 @@ +import { + assertMppChargeOperation, + mppPaymentOperationId, +} from './mpp-payment' +import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' +import { + paymentNonceKey, + type PaymentOperation, + type PaymentOperationRecoveryResult, +} from './payment-operations' +import { recoveryTiming, updateOwnedPaymentRecovery } from './payment-recovery' +import type { + AuthorizedRequest, + GatewayState, + PaymentClaimHooks, +} from './dispatch-types' +import { + assertX402V1SettlementSafe, + markRecoveryClaimed, + markRecoveryReconciled, + markRecoveryReleasing, + paymentAuthorizationContext, + preparePaymentRecovery, + recoveryTarget, + relinquishPaymentRecovery, + requirePaymentRecoveryFence, +} from './dispatch-payment-recovery' +import type { GatewayConfig } from './types' + +/** Claim payment ownership after every request guard has accepted the call. */ +export async function claimPayment( + authz: AuthorizedRequest, + config: GatewayConfig, + state: GatewayState, + hooks: PaymentClaimHooks = {}, +): Promise { + assertX402V1SettlementSafe(authz, config) + if (authz.paymentMethod === 'x402' && authz.paymentPayload) { + const context = paymentAuthorizationContext(authz) + if (config.x402.paymentProtocolVersion === 2) { + await preparePaymentRecovery(authz, config, { + kind: 'x402', + operationId: `x402:${paymentNonceKey(authz.paymentPayload)}`, + }, hooks) + } + let operation: PaymentOperation | undefined + if (config.x402.authorizePayment) { + if (config.x402.paymentProtocolVersion !== 2 && !config.x402.demoMode) { + throw new Error( + 'production x402 version 1 cannot use authorizePayment; ' + + 'use paymentProtocolVersion: 2 with paymentOperations', + ) + } + // Version 1 has no durable operation to release if another request wins + // the shared nonce while this callback is still running. This callback + // remains only for explicit demo-mode compatibility; production callers + // must use the durable version 2 operation lifecycle. + const legacyClaimed = config.x402.paymentProtocolVersion !== 2 && authz.paymentNonceKey + ? await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + : undefined + if (legacyClaimed === false) throw new Error('payment nonce was already consumed') + const result = await config.x402.authorizePayment(authz.paymentPayload, context) + if (!result) throw new Error('payment authorization was rejected') + if (typeof result !== 'boolean') { + operation = result + } + else if (config.x402.paymentProtocolVersion === 2) { + throw new Error('version 2 payment authorization did not return an operation') + } else if (authz.paymentNonceKey && legacyClaimed === undefined) { + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + if (!claimed) throw new Error('payment nonce was already consumed') + } + } else if (config.x402.paymentOperations) { + operation = await config.x402.paymentOperations.claimPayment(authz.paymentPayload, context) + } else if (authz.paymentNonceKey) { + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) + if (!claimed) throw new Error('payment nonce was already consumed') + } + if (operation && operation.protocolVersion !== 2) { + throw new Error('payment operation protocol version mismatch') + } + if ( + operation && + config.x402.paymentProtocolVersion === 2 && + operation.operationId !== authz.paymentRecoveryId + ) { + throw new Error('x402 payment operation identity mismatch') + } + if (operation && !config.x402.paymentOperations) { + throw new Error('durable payment operations are required to settle a claimed operation') + } + if (operation) { + authz.paymentOperationAcquired = operation.acquiredByRequestId === context.requestId + if (!authz.paymentOperationAcquired) { + throw new Error('payment operation was already claimed') + } + // Attach owned state before the shared nonce claim. If that claim fails, + // the caller can persist release recovery after an ambiguous refund. + authz.paymentOperation = operation + await markRecoveryClaimed(authz, config) + } + if (operation && authz.paymentNonceKey) { + const claimed = await claimPaymentNonce( + state.nonceStore, + authz.paymentNonceKey, + authz.paymentPayload, + `${operation.operationId}:${context.requestId}`, + ) + if (!claimed) { + await releaseAfterNonceConflict(authz, config) + throw new Error('payment nonce was already consumed') + } + } + } else if (authz.paymentMethod === 'mpp') { + if (!authz.paymentNonceKey) { + throw new Error('MPP payment has no replay identity') + } + const mppMethod = (authz.mppMethod ?? config.mpp?.method ?? 'blueprintevm').toLowerCase() + const durablePayload = durableMppPaymentPayload(authz.paymentPayload) + // Only BlueprinTEVM carries x402 authorization fields. Other MPP methods + // use the isolated immediate-charge lifecycle below. + if (mppMethod === 'blueprintevm' && durablePayload && config.x402.paymentOperations) { + const context = paymentAuthorizationContext(authz) + await preparePaymentRecovery(authz, config, { + kind: 'x402', + operationId: `x402:${paymentNonceKey(durablePayload)}`, + }, hooks) + const operation = await config.x402.paymentOperations.claimPayment(durablePayload, context) + if (operation.protocolVersion !== 2) throw new Error('payment operation protocol version mismatch') + if (operation.operationId !== authz.paymentRecoveryId) { + throw new Error('x402 payment operation identity mismatch') + } + if (operation.acquiredByRequestId !== context.requestId) { + throw new Error('payment operation was already claimed') + } + authz.paymentPayload = durablePayload + authz.paymentOperation = operation + authz.paymentOperationAcquired = true + await markRecoveryClaimed(authz, config) + const claimed = await claimPaymentNonce( + state.nonceStore, + authz.paymentNonceKey, + durablePayload, + `${operation.operationId}:${context.requestId}`, + ) + if (!claimed) { + await releaseAfterNonceConflict(authz, config) + throw new Error('payment nonce was already consumed') + } + } else if (mppMethod === 'blueprintevm') { + const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) + if (!claimed) throw new Error('payment nonce was already consumed') + } else { + const lifecycle = config.mpp?.charge + if (!lifecycle || lifecycle.protocolVersion !== 1) { + throw new Error('MPP charge lifecycle is not configured') + } + if (!authz.mppCredential) throw new Error('MPP payment credential is unavailable') + if (!authz.mppPaymentIdentity) throw new Error('MPP payment identity is unavailable') + const operationId = await mppPaymentOperationId(mppMethod, authz.mppPaymentIdentity) + await preparePaymentRecovery(authz, config, { + kind: 'mpp-charge', + method: mppMethod, + operationId, + }, hooks) + const claimed = await claimPaymentNonce( + state.nonceStore, + authz.paymentNonceKey, + authz.paymentPayload ?? {}, + `${operationId}:${authz.requestId}`, + ) + if (!claimed) { + await markRecoveryReconciled(authz, config) + throw new Error('payment nonce was already consumed') + } + const operation = await lifecycle.confirmPayment({ + operationId, + requestId: authz.requestId, + agentId: authz.agent.id, + consumerId: authz.consumerId, + method: mppMethod, + credential: authz.mppCredential, + amount: authz.requiredPaymentAmount, + currencyDecimals: config.x402.currencyDecimals ?? 6, + }) + assertMppChargeOperation( + operation, + { operationId, requestId: authz.requestId, method: mppMethod }, + ['confirmed'], + false, + ) + authz.mppChargeOperation = operation + await markRecoveryClaimed(authz, config) + assertMppChargeOperation( + operation, + { operationId, requestId: authz.requestId, method: mppMethod }, + ['confirmed'], + ) + } + } + + try { + await state.obs?.onPaymentVerified?.( + { + requestId: authz.requestId, + agentSlug: authz.agent.slug, + startMs: authz.startMs, + }, + { + method: authz.paymentMethod, + consumerId: authz.consumerId, + keyId: authz.keyInfo?.keyId, + }, + ) + } catch (error) { + // Observability must not turn a durable claim into a stranded payment. + console.error( + '[agent-gateway] payment observer failed for ' + authz.requestId + ':', + error instanceof Error ? error.message : String(error), + ) + } +} + +/** Release an owned operation when execution cannot produce a valid receipt. */ +export async function releasePayment( + authz: AuthorizedRequest, + config: GatewayConfig, + reason: string, +): Promise { + const ownsX402 = authz.paymentOperation && + authz.paymentOperationAcquired === true && + config.x402.paymentOperations + const ownsMpp = authz.mppChargeOperation && config.mpp?.charge + if (!ownsX402 && !ownsMpp) { + await relinquishPaymentRecovery(authz, config, Date.now()) + return + } + let reconciled = false + try { + await markRecoveryReleasing(authz, config, reason) + if (ownsX402) { + authz.paymentOperation = await config.x402.paymentOperations!.releasePayment( + authz.paymentOperation!, + reason, + ) + } else { + const operation = await config.mpp!.charge!.releasePayment(authz.mppChargeOperation!, reason) + assertMppChargeOperation( + operation, + { + operationId: authz.mppChargeOperation!.operationId, + requestId: authz.requestId, + method: authz.mppChargeOperation!.method, + }, + ['released'], + false, + ) + authz.mppChargeOperation = operation + } + await markRecoveryReconciled(authz, config) + reconciled = true + } finally { + if (!reconciled) { + try { + await relinquishPaymentRecovery(authz, config, Date.now()) + } catch { + // Preserve the original provider or metadata error. A later worker + // retry still has the durable row when cleanup itself is unavailable. + } + } + } +} + +/** Mark a durable reservation active immediately before sandbox execution. */ +export async function beginPaymentExecution( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { + authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) + } +} + +/** Persist the sandbox handoff immediately before the adapter call. */ +export async function markPaymentExecutionStarted( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + await updateExecutionLease(authz, config, true) +} + +/** Renew the live execution lease while a provider stream is still open. */ +export async function renewPaymentExecution( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + await updateExecutionLease(authz, config, false) +} + +async function updateExecutionLease( + authz: AuthorizedRequest, + config: GatewayConfig, + markStarted: boolean, +): Promise { + if (!authz.paymentRecoveryId) return + const recovery = config.paymentRecovery + if (!recovery) throw new Error('durable payment recovery is not configured') + const fenceId = requirePaymentRecoveryFence(authz) + const now = Date.now() + const fallbackAt = now + recoveryTiming(recovery).receiptTimeoutMs + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ + ...record, + state: 'executing', + ...(markStarted + ? { payment: recoveryTarget(authz, record.payment), workStarted: true } + : {}), + fallbackAt, + lease: { id: fenceId, expiresAt: fallbackAt }, + nextAttemptAt: fallbackAt, + }), now) +} + +/** + * Release only when no sandbox work was observed. Once output or a receipt + * exists, retain the owner for settlement or background recovery. + */ +export async function releasePaymentAfterFailure( + authz: AuthorizedRequest, + config: GatewayConfig, + reason: string, + workObserved: boolean, +): Promise { + if (workObserved) { + const recovery = config.paymentRecovery + if (recovery && authz.paymentRecoveryId) { + const fenceId = requirePaymentRecoveryFence(authz) + await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => { + if (record.state === 'settling' && record.usage) { + return { ...record, lease: undefined, nextAttemptAt: Date.now() } + } + const fallbackAt = record.fallbackAt ?? + Date.now() + recoveryTiming(recovery).receiptTimeoutMs + return { + ...record, + state: 'retained', + payment: recoveryTarget(authz, record.payment), + workStarted: true, + fallbackAt, + reason, + lease: undefined, + nextAttemptAt: fallbackAt, + } + }) + } + if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { + authz.paymentOperation = await config.x402.paymentOperations.retainPayment(authz.paymentOperation, reason) + } + console.error( + `[agent-gateway] retaining payment ownership after sandbox work for ${authz.requestId}: ${reason}`, + ) + return + } + await releasePayment(authz, config, reason) +} + +export async function reclaimPayment( + operationId: string, + config: GatewayConfig, +): Promise { + if (!config.x402.paymentOperations) throw new Error('durable payment operations are not configured') + return config.x402.paymentOperations.reclaimPayment(operationId) +} + +async function releaseAfterNonceConflict( + authz: AuthorizedRequest, + config: GatewayConfig, +): Promise { + try { + await releasePayment(authz, config, 'shared payment nonce was already owned') + } catch (releaseError) { + console.error( + `[agent-gateway] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } +} + +async function claimPaymentNonce( + nonceStore: NonceStore, + nonceKey: string, + payload: Record, + ownerId?: string, +): Promise { + const expiry = payload.expiry === undefined + ? BigInt(Math.floor(Date.now() / 1000) + 3600) + : BigInt(String(payload.expiry)) + const ttl = nonceTtlSeconds(expiry) + if (ttl === undefined) return false + return claimStoredNonce(nonceStore, nonceKey, ttl, ownerId) +} + +function durableMppPaymentPayload( + payload: Record | null, +): Record | undefined { + if (!payload) return undefined + const commitment = payload.commitment ?? payload.from + const amount = payload.amount ?? payload.value + const nonce = payload.nonce + if (typeof commitment !== 'string' || commitment.length === 0) return undefined + if (amount === undefined || nonce === undefined) return undefined + const amountText = String(amount) + const nonceText = String(nonce) + if (!/^\d+$/.test(amountText) || !/^\d+$/.test(nonceText)) return undefined + const expiryText = payload.expiry === undefined + ? String(Math.floor(Date.now() / 1000) + 3600) + : String(payload.expiry) + if (!/^\d+$/.test(expiryText)) return undefined + return { + ...payload, + commitment, + amount: amountText, + nonce: nonceText, + expiry: expiryText, + } +} diff --git a/src/dispatch-pricing.ts b/src/dispatch-pricing.ts new file mode 100644 index 0000000..3b1fe58 --- /dev/null +++ b/src/dispatch-pricing.ts @@ -0,0 +1,108 @@ +import type { AgentMeta, ChatMessage } from './types' + +function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { + if (!Number.isFinite(value) || value < 0) { + throw new Error('agent pricePerTokenUsd must be a finite non-negative number') + } + const [mantissa, exponentText] = value.toString().toLowerCase().split('e') + const exponent = exponentText ? Number(exponentText) : 0 + const [whole, fraction = ''] = mantissa.split('.') + let numerator = BigInt(`${whole}${fraction}`) + let scale = fraction.length - exponent + if (scale < 0) { + numerator *= 10n ** BigInt(-scale) + scale = 0 + } + return { numerator, denominator: 10n ** BigInt(scale) } +} + +function amountForTokens( + pricePerTokenUsd: number, + tokenCount: number, + currencyDecimals: number, + providerCostUsd: number, +): bigint { + const { numerator, denominator } = decimalFraction(pricePerTokenUsd) + const scaled = BigInt(tokenCount) * numerator * 10n ** BigInt(currencyDecimals) + const tokenAmount = (scaled + denominator - 1n) / denominator + const provider = providerCostUsd === 0 + ? { numerator: 0n, denominator: 1n } + : decimalFraction(providerCostUsd) + const providerScaled = provider.numerator * 10n ** BigInt(currencyDecimals) + const providerAmount = (providerScaled + provider.denominator - 1n) / provider.denominator + return tokenAmount > providerAmount ? tokenAmount : providerAmount +} + +/** Exact base-unit reservation required to cover the request's token ceiling. */ +export function requiredX402Amount( + pricePerTokenUsd: number, + inputTokens: number, + maxOutputTokens: number, + currencyDecimals = 6, + maxReasoningTokens = 0, + maxToolTokens = 0, + maxProviderCostUsd = 0, +): bigint { + if (!Number.isSafeInteger(inputTokens) || inputTokens < 0) { + throw new Error('input token estimate must be a non-negative safe integer') + } + if (!Number.isSafeInteger(maxOutputTokens) || maxOutputTokens <= 0) { + throw new Error('max output tokens must be a positive safe integer') + } + if (!Number.isInteger(currencyDecimals) || currencyDecimals < 0 || currencyDecimals > 18) { + throw new Error('x402 currencyDecimals must be an integer between 0 and 18') + } + for (const [name, value] of [ + ['maxReasoningTokens', maxReasoningTokens], + ['maxToolTokens', maxToolTokens], + ] as const) { + if (!Number.isSafeInteger(value) || value < 0) throw new Error(`${name} must be a non-negative safe integer`) + } + if (!Number.isFinite(maxProviderCostUsd) || maxProviderCostUsd < 0) { + throw new Error('maxProviderCostUsd must be finite and non-negative') + } + const tokenCount = inputTokens + maxOutputTokens + maxReasoningTokens + maxToolTokens + if (!Number.isSafeInteger(tokenCount)) throw new Error('token budget exceeds safe integer range') + return amountForTokens(pricePerTokenUsd, tokenCount, currencyDecimals, maxProviderCostUsd) +} + +export function actualX402Amount( + pricePerTokenUsd: number, + inputTokens: number, + outputTokens: number, + reasoningTokens: number, + toolTokens: number, + currencyDecimals = 6, + providerCostUsd = 0, +): bigint { + return amountForTokens( + pricePerTokenUsd, + inputTokens + outputTokens + reasoningTokens + toolTokens, + currencyDecimals, + providerCostUsd, + ) +} + +/** Token estimate matching the existing chat-completions handler (4 chars ≈ 1 token). */ +export function estimateTokens(text: string): number { + return Math.ceil(text.length / 4) +} + +/** Include the host-owned system prompt because the provider bills it too. */ +export function estimateBillableInputTokens(agent: AgentMeta, userMessage: string): number { + return estimateTokens(userMessage) + estimateTokens(agent.systemPrompt ?? '') +} + +/** A tokenizer cannot emit more tokens than the UTF-8 bytes it consumes. */ +export function maximumBillableInputTokens(agent: AgentMeta, userMessage: string): number +export function maximumBillableInputTokens(agent: AgentMeta, messages: readonly ChatMessage[]): number +export function maximumBillableInputTokens( + agent: AgentMeta, + userMessageOrMessages: string | readonly ChatMessage[], +): number { + const encoder = new TextEncoder() + const prompt = typeof userMessageOrMessages === 'string' + ? userMessageOrMessages + : JSON.stringify(userMessageOrMessages) + return encoder.encode(prompt).byteLength + encoder.encode(agent.systemPrompt ?? '').byteLength +} diff --git a/src/dispatch-sandbox.ts b/src/dispatch-sandbox.ts new file mode 100644 index 0000000..bd31abd --- /dev/null +++ b/src/dispatch-sandbox.ts @@ -0,0 +1,422 @@ +import { redactSystemPromptFromOutput } from './filter' +import type { A2ADispatchEvent } from './dispatch-types' +import { + estimateTokens, + maximumBillableInputTokens, +} from './dispatch-pricing' +import type { + AgentMeta, + GatewayConfig, + SandboxExecutionBudget, + SandboxStreamEvent, + SandboxUsageReceipt, +} from './types' + +export async function* dispatchSandboxStream( + agent: AgentMeta, + userMessage: string, + consumerId: string, + config: GatewayConfig, + signal?: AbortSignal, + sessionId?: string, + maxOutputTokens?: number, +): AsyncIterable { + for await (const event of dispatchSandboxStreamRich( + agent, + userMessage, + consumerId, + config, + signal, + sessionId, + maxOutputTokens, + )) { + if (event.kind === 'text') yield event.delta + } +} + +/** + * Like `dispatchSandboxStream` but yields a discriminated union so callers can + * react to `input-required` signals from the sandbox. The sandbox opts in by + * emitting `{ type: 'input-required', data: { inputRequired: { prompt? } } }` + * (or by setting `data.inputRequired` on any event); sandboxes that don't emit + * such events see identical behavior. + * + * `sessionId` defaults to `consumer:` matching the existing single-turn + * path; multi-turn continuations pass an explicit `taskId` so the sandbox can + * keep per-task conversation memory. + */ +export async function* dispatchSandboxStreamRich( + agent: AgentMeta, + userMessage: string, + consumerId: string, + config: GatewayConfig, + signal?: AbortSignal, + sessionId?: string, + maxOutputTokens?: number, + onExecutionStart?: () => Promise, + requiresReceipt = config.x402.paymentOperations !== undefined, + onSandboxStart?: () => void | Promise, + maxInputTokens?: number, + onExecutionHeartbeat?: () => Promise, +): AsyncIterable { + if (signal?.aborted) return + const box = await config.getSandbox(agent) + if (signal?.aborted) return + const outputLimit = maxOutputTokens ?? config.defaultOutputTokens ?? 1024 + if (!Number.isSafeInteger(outputLimit) || outputLimit <= 0) { + throw new Error('max output tokens must be a positive safe integer') + } + let outputBytes = 0 + // Bound untrusted adapters while the final receipt is pending. The receipt + // remains authoritative for token count, so over-limit output is never sent. + const maxOutputBytes = outputLimit * 4 + if (!Number.isSafeInteger(maxOutputBytes)) { + throw new Error('max output token bound exceeds safe integer range') + } + const encoder = new TextEncoder() + let usageParts: Partial = {} + let observedReasoningTokens = 0 + let observedToolTokens = 0 + let observedToolCalls = 0 + let legacyOutputText = '' + 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 { + // 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, + }) + iterator = promptStream[Symbol.asyncIterator]() + const heartbeatMs = onExecutionHeartbeat + ? Math.max(100, Math.min( + Math.floor((config.paymentRecovery?.receiptTimeoutMs ?? 5 * 60_000) / 3), + 5_000, + )) + : 0 + if (onExecutionHeartbeat) { + heartbeatTimer = setInterval(() => { + if (heartbeatInFlight || heartbeatError !== undefined) return + heartbeatInFlight = onExecutionHeartbeat() + .catch((error: unknown) => { + heartbeatError = error + executionController.abort() + }) + .finally(() => { + heartbeatInFlight = undefined + }) + }, heartbeatMs) + } + while (true) { + const next = await readSandboxEvent(iterator, executionController.signal) + if (next === ABORTED_SANDBOX_READ) { + if (heartbeatError !== undefined) throw heartbeatError + return + } + if (next.done) break + const event = next.value + if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) + if (event.data?.reasoning?.tokens !== undefined) { + observedReasoningTokens += nonNegativeSafeInteger(event.data.reasoning.tokens, 'reasoning tokens') + yield { kind: 'activity' } + } + if (event.data?.tool) { + observedToolCalls += 1 + observedToolTokens += + nonNegativeSafeInteger(event.data.tool.inputTokens ?? 0, 'tool input tokens') + + nonNegativeSafeInteger(event.data.tool.outputTokens ?? 0, 'tool output tokens') + yield { kind: 'activity' } + } + enforceUsageBudget(withObservedUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + ), executionBudget) + if ( + event.type === 'message.part.updated' && + event.data?.part?.type === 'text' && + event.data.delta + ) { + const remainingBytes = maxOutputBytes - outputBytes + if (remainingBytes <= 0) throw new Error('sandbox exceeded max output tokens') + const bounded = truncateUtf8(event.data.delta, remainingBytes, encoder) + if (bounded.truncated) { + yield { kind: 'activity' } + throw new Error('sandbox exceeded max output tokens') + } + outputBytes += bounded.bytes + legacyOutputText += bounded.text + yield { kind: 'activity' } + yield { kind: 'text', delta: redactSystemPromptFromOutput(bounded.text, agent.systemPrompt) } + continue + } + if (event.type === 'input-required' || event.data?.inputRequired) { + const usage = completeUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + userMessage, + legacyOutputText, + executionBudget, + requiresReceipt, + ) + yield { kind: 'input-required', prompt: event.data?.inputRequired?.prompt } + // Terminal for the sandbox stream — sandbox SHOULD stop emitting until + // the gateway dispatches a continuation message with the new user input. + yield { kind: 'usage', usage } + return + } + } + const usage = completeUsage( + usageParts, + observedReasoningTokens, + observedToolTokens, + observedToolCalls, + userMessage, + legacyOutputText, + executionBudget, + requiresReceipt, + ) + yield { kind: 'usage', usage } + } finally { + if (heartbeatTimer !== undefined) clearInterval(heartbeatTimer) + const pendingHeartbeat = heartbeatInFlight + if (pendingHeartbeat) await pendingHeartbeat + signal?.removeEventListener('abort', forwardAbort) + if (iterator) await closeSandboxIterator(iterator) + if (heartbeatError !== undefined && !signal?.aborted) throw heartbeatError + } +} + +const ABORTED_SANDBOX_READ = Symbol('aborted-sandbox-read') + +async function readSandboxEvent( + iterator: AsyncIterator, + signal?: AbortSignal, +): Promise | typeof ABORTED_SANDBOX_READ> { + if (!signal) return iterator.next() + if (signal.aborted) return ABORTED_SANDBOX_READ + return new Promise((resolve, reject) => { + const onAbort = () => { + signal.removeEventListener('abort', onAbort) + resolve(ABORTED_SANDBOX_READ) + } + signal.addEventListener('abort', onAbort, { once: true }) + iterator.next().then( + (result) => { + signal.removeEventListener('abort', onAbort) + resolve(result) + }, + (error: unknown) => { + signal.removeEventListener('abort', onAbort) + reject(error) + }, + ) + }) +} + +const SANDBOX_CLEANUP_TIMEOUT_MS = 50 + +async function closeSandboxIterator(iterator: AsyncIterator): Promise { + let closing: PromiseLike | undefined + try { + const result = iterator.return?.() + if (result) closing = Promise.resolve(result) + } catch { + return + } + if (!closing) return + let timeout: ReturnType | undefined + try { + await Promise.race([ + Promise.resolve(closing).catch(() => undefined), + new Promise((resolve) => { + timeout = setTimeout(resolve, SANDBOX_CLEANUP_TIMEOUT_MS) + }), + ]) + } finally { + if (timeout !== undefined) clearTimeout(timeout) + } +} + +function mergeUsage( + current: Partial, + update: Partial, +): Partial { + const merged = { ...current, ...update } + for (const key of ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount', 'providerCostUsd'] as const) { + const value = update[key] + if (value !== undefined && (!Number.isFinite(value) || value < 0)) { + throw new Error(`sandbox usage field ${key} is invalid`) + } + if (value !== undefined && current[key] !== undefined) { + // Usage events are cumulative receipts. Never let a later partial or + // final event erase spend observed earlier in the same execution. + merged[key] = Math.max(current[key]!, value) + } + } + if (current.budgetEnforced === false || update.budgetEnforced === false) { + merged.budgetEnforced = false + } + return merged +} + +function withObservedUsage( + usage: Partial, + reasoningTokens: number, + toolTokens: number, + toolCallCount: number, +): Partial { + return { + ...usage, + ...(usage.reasoningTokens !== undefined || reasoningTokens > 0 + ? { reasoningTokens: Math.max(usage.reasoningTokens ?? 0, reasoningTokens) } + : {}), + ...(usage.toolTokens !== undefined || toolTokens > 0 + ? { toolTokens: Math.max(usage.toolTokens ?? 0, toolTokens) } + : {}), + ...(usage.toolCallCount !== undefined || toolCallCount > 0 + ? { toolCallCount: Math.max(usage.toolCallCount ?? 0, toolCallCount) } + : {}), + } +} + +function nonNegativeSafeInteger(value: number, name: string): number { + if (!Number.isSafeInteger(value) || value < 0) { + throw new Error(`sandbox ${name} is invalid`) + } + return value +} + +function enforceUsageBudget( + usage: Partial, + budget: SandboxExecutionBudget, +): void { + if (usage.inputTokens !== undefined && usage.inputTokens > budget.maxInputTokens) { + throw new Error('sandbox exceeded max input tokens') + } + if (usage.outputTokens !== undefined && usage.outputTokens > budget.maxOutputTokens) { + throw new Error('sandbox exceeded max output tokens') + } + if (usage.reasoningTokens !== undefined && usage.reasoningTokens > budget.maxReasoningTokens) { + throw new Error('sandbox exceeded max reasoning tokens') + } + if (usage.toolTokens !== undefined && usage.toolTokens > budget.maxToolTokens) { + throw new Error('sandbox exceeded max tool tokens') + } + if (usage.toolCallCount !== undefined && usage.toolCallCount > budget.maxToolCalls) { + throw new Error('sandbox exceeded max tool calls') + } + if (usage.providerCostUsd !== undefined && usage.providerCostUsd > budget.maxProviderCostUsd) { + throw new Error('sandbox exceeded max provider cost') + } +} + +function finalizeUsage( + parts: Partial, + budget: SandboxExecutionBudget, +): SandboxUsageReceipt { + const fields = ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount', 'providerCostUsd', 'budgetEnforced'] as const + if (fields.some((field) => parts[field] === undefined)) { + throw new Error('sandbox did not provide a complete usage receipt') + } + const usage = parts as SandboxUsageReceipt + for (const field of ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount'] as const) { + if (!Number.isSafeInteger(usage[field]) || usage[field] < 0) { + throw new Error(`sandbox usage field ${field} is invalid`) + } + } + if (!Number.isFinite(usage.providerCostUsd) || usage.providerCostUsd < 0) { + throw new Error('sandbox usage provider cost is invalid') + } + if (typeof usage.budgetEnforced !== 'boolean') { + throw new Error('sandbox usage budget flag is invalid') + } + if (!Number.isSafeInteger( + usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens, + )) { + throw new Error('sandbox usage token total exceeds safe integer range') + } + enforceUsageBudget(usage, budget) + if (!usage.budgetEnforced) throw new Error('sandbox did not enforce the execution budget') + return usage +} + +function completeUsage( + parts: Partial, + reasoningTokens: number, + toolTokens: number, + toolCallCount: number, + userMessage: string, + outputText: string, + budget: SandboxExecutionBudget, + requiresReceipt: boolean, +): SandboxUsageReceipt { + const observed = withObservedUsage(parts, reasoningTokens, toolTokens, toolCallCount) + if ( + !requiresReceipt && + Object.keys(parts).length === 0 && + reasoningTokens === 0 && + toolTokens === 0 && + toolCallCount === 0 + ) { + // Preserve the pre-receipt SandboxBox contract for legacy API-key + // adapters. Durable payment operations must use provider-enforced usage. + return { + inputTokens: estimateTokens(userMessage), + outputTokens: estimateTokens(outputText), + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0, + budgetEnforced: false, + } + } + return finalizeUsage(observed, budget) +} + +function truncateUtf8( + value: string, + maxBytes: number, + encoder: TextEncoder, +): { text: string; bytes: number; truncated: boolean } { + const bytes = encoder.encode(value).byteLength + if (bytes <= maxBytes) return { text: value, bytes, truncated: false } + let text = '' + let used = 0 + for (const character of value) { + const characterBytes = encoder.encode(character).byteLength + if (used + characterBytes > maxBytes) break + text += character + used += characterBytes + } + return { text, bytes: used, truncated: true } +} diff --git a/src/dispatch-settlement.ts b/src/dispatch-settlement.ts new file mode 100644 index 0000000..e25186e --- /dev/null +++ b/src/dispatch-settlement.ts @@ -0,0 +1,139 @@ +import { type GatewayObserver, type RequestContext } from './observer' +import { actualX402Amount } from './dispatch-pricing' +import { + assertX402V1SettlementSafe, + markRecoveryReconciled, + markRecoverySettling, + markRecoveryUsageRecorded, +} from './dispatch-payment-recovery' +import type { + AuthorizedRequest, + SettleAndRecordOptions, +} from './dispatch-types' +import type { AgentMeta, GatewayConfig, SandboxUsageReceipt } from './types' + +/** + * Record usage, settle payment, and invoke the observer. Both wire formats + * call this once their stream has drained, so settlement happens exactly once + * per request regardless of protocol. + */ +export async function settleAndRecord( + agent: AgentMeta, + authz: AuthorizedRequest, + usage: SandboxUsageReceipt, + config: GatewayConfig, + obs: GatewayObserver | undefined, + options: SettleAndRecordOptions = {}, +): Promise { + assertX402V1SettlementSafe(authz, config) + const settlementBasis = options.settlementBasis ?? 'usage-receipt' + await markRecoverySettling(authz, usage, settlementBasis, config) + if (options.usageAlreadyRecorded) await markRecoveryUsageRecorded(authz, config) + const tokenCost = ( + usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens + ) * agent.pricePerTokenUsd + const totalCost = Math.max(tokenCost, usage.providerCostUsd) + const ownerEarned = totalCost * (1 - agent.platformFeePercent) + const platformFee = totalCost * agent.platformFeePercent + const usageEvent = { + requestId: authz.requestId, + agentId: agent.id, + agentSlug: agent.slug, + consumerId: authz.consumerId, + paymentMethod: authz.paymentMethod, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + reasoningTokens: usage.reasoningTokens, + toolTokens: usage.toolTokens, + toolCallCount: usage.toolCallCount, + providerCostUsd: usage.providerCostUsd, + totalCostUsd: totalCost, + ownerEarnedUsd: ownerEarned, + platformFeeUsd: platformFee, + durationMs: Date.now() - authz.startMs, + settlementBasis, + } + const ctx: RequestContext = { + requestId: authz.requestId, + agentSlug: agent.slug, + startMs: authz.startMs, + } + try { + if (authz.paymentOperation && config.x402.paymentOperations) { + const amount = options.paymentAmount ?? actualX402Amount( + agent.pricePerTokenUsd, + usage.inputTokens, + usage.outputTokens, + usage.reasoningTokens, + usage.toolTokens, + config.x402.currencyDecimals, + usage.providerCostUsd, + ) + if (options.paymentAlreadySettled) { + if ( + authz.paymentOperation.state !== 'settled' || + authz.paymentOperation.settledAmount !== amount + ) { + throw new Error('authoritative payment state does not match finalization') + } + } else { + authz.paymentOperation = await config.x402.paymentOperations.settlePayment( + authz.paymentOperation, + { amount, totalCostUsd: totalCost, usage, basis: settlementBasis }, + ) + } + // Durable settlement happens first. If attribution storage is + // unavailable, recovery must never refund delivered work. + if (!options.usageAlreadyRecorded) { + await config.recordUsage(usageEvent) + await options.onUsageRecorded?.() + await markRecoveryUsageRecorded(authz, config) + } + } else if (authz.mppChargeOperation) { + // Generic MPP charge methods confirm before the response. Finalization + // records attribution only; it never invokes the legacy settlement hook. + if (!options.usageAlreadyRecorded) { + await config.recordUsage(usageEvent) + await options.onUsageRecorded?.() + await markRecoveryUsageRecorded(authz, config) + } + } else { + // Legacy adapters retain attribution-before-charge because their + // settlement callback may resolve that usage row. + if (!options.usageAlreadyRecorded) { + await config.recordUsage(usageEvent) + await options.onUsageRecorded?.() + } + if (config.settlePayment) { + await config.settlePayment( + { + method: authz.paymentMethod, + consumerId: authz.consumerId, + requestId: authz.requestId, + }, + totalCost, + ) + } + } + await markRecoveryReconciled(authz, config) + } catch (err) { + const msg = err instanceof Error ? err.message : String(err) + console.error(`[agent-gateway] settlement failed for ${authz.consumerId}: ${msg}`) + await obs?.onSettlementError?.(ctx, { + consumerId: authz.consumerId, + method: authz.paymentMethod, + errorMessage: msg, + }) + throw err + } + try { + await obs?.onRequestComplete?.(ctx, usageEvent) + } catch (error) { + console.error( + `[agent-gateway] completion observer failed for ${authz.requestId}:`, + error instanceof Error ? error.message : String(error), + ) + } +} + +export type { SettleAndRecordOptions } diff --git a/src/dispatch-types.ts b/src/dispatch-types.ts new file mode 100644 index 0000000..da64c54 --- /dev/null +++ b/src/dispatch-types.ts @@ -0,0 +1,81 @@ +import type { MppChargeOperation } from './mpp-payment' +import type { NonceStore } from './nonce-store' +import type { GatewayObserver } from './observer' +import type { RateLimitStore } from './rate-limit' +import type { PaymentOperation } from './payment-operations' +import type { PaymentSettlementBasis } from './payment-recovery' +import type { + AgentMeta, + ApiKeyInfo, + PaymentMethod, + SandboxExecutionBudget, + SandboxUsageReceipt, +} from './types' + +/** Long-lived state shared by all handlers created for one gateway. */ +export interface GatewayState { + rateLimitStore: RateLimitStore + nonceStore: NonceStore + globalRateLimit: { limit: number; windowSeconds: number } + requiredScope: string + maxLen: number + maxOutputTokens: number + defaultOutputTokens: number + maxReasoningTokens: number + maxToolTokens: number + maxToolCalls: number + maxProviderCostUsd?: number + obs?: GatewayObserver +} + +/** Successful output from the request authorization pipeline. */ +export interface AuthorizedRequest { + agent: AgentMeta + consumerId: string + paymentMethod: PaymentMethod + keyInfo: ApiKeyInfo | null + userMessage: string + rateLimitRemaining: number | undefined + requestId: string + startMs: number + maxOutputTokens: number + executionBudget: SandboxExecutionBudget + requiredPaymentAmount: bigint + paymentPayload: Record | null + paymentNonceKey?: string + mppMethod?: string + /** Live generic MPP credential. Never write it to the recovery store. */ + mppCredential?: string + /** Stable method identity. Persist only its digest. */ + mppPaymentIdentity?: string + mppChargeOperation?: MppChargeOperation + paymentOperation?: PaymentOperation + paymentOperationAcquired?: boolean + paymentRecoveryId?: string + /** Unique ownership fence for live or recovery transitions. */ + paymentRecoveryFence?: string +} + +export interface PaymentClaimHooks { + /** Persist the recovery identity before the provider can mutate payment state. */ + onRecoveryPrepared?: (recoveryId: string) => Promise +} + +export type A2ADispatchEvent = + | { kind: 'text'; delta: string } + | { kind: 'input-required'; prompt?: string } + | { kind: 'activity' } + | { kind: 'usage'; usage: SandboxUsageReceipt } + +export interface SettleAndRecordOptions { + /** Skip attribution after a durable finalization marker confirms it ran. */ + usageAlreadyRecorded?: boolean + /** Skip provider settlement after an authoritative read found it settled. */ + paymentAlreadySettled?: boolean + /** Persist the recovery marker after attribution succeeds. */ + onUsageRecorded?: () => Promise + /** Recovery uses the original quoted ceiling when no receipt arrives. */ + settlementBasis?: PaymentSettlementBasis + /** Exact base-unit charge selected by the recovery policy. */ + paymentAmount?: bigint +} diff --git a/src/dispatch.ts b/src/dispatch.ts index 2ebf25f..5788e31 100644 --- a/src/dispatch.ts +++ b/src/dispatch.ts @@ -1,1799 +1,38 @@ /** - * Shared inner pipeline used by every wire-format the gateway exposes - * (OpenAI-compatible chat completions, A2A JSON-RPC). Each handler parses its - * own protocol's request body into a canonical `messages[]` form + headers, - * then calls into here for auth → rate-limit → injection filter → - * authorize → sandbox stream → settle. Keeping the pipeline single-sourced - * means every protocol surface gets the same security and billing guarantees - * for free; bugs fixed here fix every wrapper. + * Stable dispatch surface shared by the OpenAI-compatible and A2A handlers. + * Each implementation lives in the module that owns its state transitions. */ -import type { Context } from 'hono' - -import { filterConsumerMessagesStrict, redactSystemPromptFromOutput } from './filter' -import { - assertMppChargeOperation, - mppPaymentOperationId, - type MppChargeOperation, -} from './mpp-payment' -import { type GatewayObserver, type RequestContext, generateRequestId } from './observer' -import { type RateLimitStore, checkRateLimit } from './rate-limit' -import { claimStoredNonce, nonceTtlSeconds, type NonceStore } from './nonce-store' -import { - paymentNonceKey, - type PaymentOperation, - type PaymentOperationRecoveryResult, -} from './payment-operations' -import { - PAYMENT_RECOVERY_VERSION, - PaymentRecoveryFenceError, - PaymentRecoveryReplayError, - recoveryTiming, - serializePaymentOperation, - updateOwnedPaymentRecovery, - type PaymentRecoveryRecord, - type PaymentRecoveryTarget, - type PaymentSettlementBasis, -} from './payment-recovery' -import type { - AgentMeta, - ApiKeyInfo, - ChatMessage, - GatewayConfig, - PaymentMethod, - SandboxExecutionBudget, - SandboxStreamEvent, - SandboxUsageReceipt, -} from './types' -import { - defaultVerifyApiKey, - isApiKeyAuthEnabled, - isMppAuthEnabled, - mppPaymentPayload, - mppPaymentCredential, - verifyMppCredential, - verifyX402, -} from './verify' - -/** Single bundle of long-lived gateway state shared across all handlers in one createAgentGateway call. */ -export interface GatewayState { - rateLimitStore: RateLimitStore - nonceStore: NonceStore - globalRateLimit: { limit: number; windowSeconds: number } - requiredScope: string - maxLen: number - maxOutputTokens: number - defaultOutputTokens: number - maxReasoningTokens: number - maxToolTokens: number - maxToolCalls: number - maxProviderCostUsd?: number - obs?: GatewayObserver -} - -/** Returned by {@link authenticateAndGuard} on the success path. */ -export interface AuthorizedRequest { - agent: AgentMeta - consumerId: string - paymentMethod: PaymentMethod - keyInfo: ApiKeyInfo | null - userMessage: string - rateLimitRemaining: number | undefined - requestId: string - startMs: number - maxOutputTokens: number - executionBudget: SandboxExecutionBudget - requiredPaymentAmount: bigint - paymentPayload: Record | null - paymentNonceKey?: string - mppMethod?: string - /** Live generic MPP credential. Never written to the recovery store. */ - mppCredential?: string - /** Stable method identity. The gateway persists only its digest. */ - mppPaymentIdentity?: string - mppChargeOperation?: MppChargeOperation - paymentOperation?: PaymentOperation - paymentOperationAcquired?: boolean - paymentRecoveryId?: string - /** Unique ownership fence for live or recovery transitions. */ - paymentRecoveryFence?: string -} - -export interface PaymentClaimHooks { - /** Persist the recovery identity before the provider can mutate payment state. */ - onRecoveryPrepared?: (recoveryId: string) => Promise -} - -function decimalFraction(value: number): { numerator: bigint; denominator: bigint } { - if (!Number.isFinite(value) || value < 0) { - throw new Error('agent pricePerTokenUsd must be a finite non-negative number') - } - const [mantissa, exponentText] = value.toString().toLowerCase().split('e') - const exponent = exponentText ? Number(exponentText) : 0 - const [whole, fraction = ''] = mantissa.split('.') - let numerator = BigInt(`${whole}${fraction}`) - let scale = fraction.length - exponent - if (scale < 0) { - numerator *= 10n ** BigInt(-scale) - scale = 0 - } - return { numerator, denominator: 10n ** BigInt(scale) } -} - -/** Exact base-unit reservation required to cover the request's token ceiling. */ -export function requiredX402Amount( - pricePerTokenUsd: number, - inputTokens: number, - maxOutputTokens: number, - currencyDecimals = 6, - maxReasoningTokens = 0, - maxToolTokens = 0, - maxProviderCostUsd = 0, -): bigint { - if (!Number.isSafeInteger(inputTokens) || inputTokens < 0) { - throw new Error('input token estimate must be a non-negative safe integer') - } - if (!Number.isSafeInteger(maxOutputTokens) || maxOutputTokens <= 0) { - throw new Error('max output tokens must be a positive safe integer') - } - if (!Number.isInteger(currencyDecimals) || currencyDecimals < 0 || currencyDecimals > 18) { - throw new Error('x402 currencyDecimals must be an integer between 0 and 18') - } - for (const [name, value] of [ - ['maxReasoningTokens', maxReasoningTokens], - ['maxToolTokens', maxToolTokens], - ] as const) { - if (!Number.isSafeInteger(value) || value < 0) throw new Error(`${name} must be a non-negative safe integer`) - } - if (!Number.isFinite(maxProviderCostUsd) || maxProviderCostUsd < 0) { - throw new Error('maxProviderCostUsd must be finite and non-negative') - } - const { numerator, denominator } = decimalFraction(pricePerTokenUsd) - const tokenCount = inputTokens + maxOutputTokens + maxReasoningTokens + maxToolTokens - if (!Number.isSafeInteger(tokenCount)) throw new Error('token budget exceeds safe integer range') - const tokenScaled = BigInt(tokenCount) * - numerator * 10n ** BigInt(currencyDecimals) - const tokenAmount = (tokenScaled + denominator - 1n) / denominator - const provider = maxProviderCostUsd === 0 - ? { numerator: 0n, denominator: 1n } - : decimalFraction(maxProviderCostUsd) - const providerScaled = provider.numerator * 10n ** BigInt(currencyDecimals) - const providerAmount = (providerScaled + provider.denominator - 1n) / provider.denominator - return tokenAmount > providerAmount ? tokenAmount : providerAmount -} - -/** - * Resolve the agent, then run the full pre-dispatch pipeline: payment + - * rate-limit + injection filter + user-message extraction + optional - * `authorizeConsumer` hook. Returns the success record on the happy path - * or a fully-formed `Response` (402/404/429/400/403) on any short-circuit. - * - * Body parsing is the caller's responsibility — different wire formats - * (OpenAI chat completions vs A2A JSON-RPC) have different envelopes; both - * still ultimately produce a `ChatMessage[]`. - */ -export async function authenticateAndGuard( - c: Context, - slug: string, - messages: ChatMessage[], - config: GatewayConfig, - state: GatewayState, - requestedMaxOutputTokens?: number, -): Promise { - const startMs = Date.now() - const requestId = generateRequestId() - const ctx: RequestContext = { requestId, agentSlug: slug, startMs } - await state.obs?.onRequestStart?.(ctx) - - const agent = await config.resolveAgent(slug) - if (!agent || !agent.enabled) { - return c.json({ error: { message: 'Agent not found', type: 'not_found' } }, 404) - } - if (!messages?.length) { - return c.json( - { error: { message: 'messages array required', type: 'invalid_request' } }, - 400, - ) - } - - const maxOutputTokens = requestedMaxOutputTokens ?? state.defaultOutputTokens - if ( - !Number.isInteger(maxOutputTokens) || - maxOutputTokens <= 0 || - maxOutputTokens > state.maxOutputTokens - ) { - return c.json( - { - error: { - message: `max_tokens must be an integer between 1 and ${state.maxOutputTokens}`, - type: 'invalid_request', - code: 'invalid_max_tokens', - }, - }, - 400, - ) - } - - // Quote the maximum UTF-8 input plus every hidden provider cost before - // verification. The verifier must remain read-only at this point. - const { messages: filtered, injectionWarnings } = filterConsumerMessagesStrict( - messages, - state.maxLen, - ) - const userMessage = filtered - .filter((m) => m.role === 'user') - .map((m) => m.content) - .join('\n\n') - if (!userMessage) { - return c.json( - { error: { message: 'No user message provided', type: 'invalid_request' } }, - 400, - ) - } - - let requiredPaymentAmount: bigint - const messageInputBound = maximumBillableInputTokens(agent, filtered) - let maxInputTokens = messageInputBound - if (config.inputTokenBound) { - let configuredBound: number - try { - configuredBound = config.inputTokenBound({ agent, messages: filtered }) - } catch { - return c.json( - { - error: { - message: 'Agent input token bound is unavailable', - type: 'server_error', - code: 'input_token_bound_unavailable', - }, - }, - 503, - ) - } - if (!Number.isSafeInteger(configuredBound) || configuredBound < messageInputBound) { - return c.json( - { - error: { - message: 'Agent input token bound is invalid', - type: 'server_error', - code: 'invalid_input_token_bound', - }, - }, - 503, - ) - } - maxInputTokens = configuredBound - } - const maxReasoningTokens = state.maxReasoningTokens - const maxToolTokens = state.maxToolTokens - const maxToolCalls = state.maxToolCalls - const maxProviderCostUsd = state.maxProviderCostUsd ?? - (maxInputTokens + maxOutputTokens + maxReasoningTokens + maxToolTokens) * agent.pricePerTokenUsd - const executionBudget: SandboxExecutionBudget = { - maxInputTokens, - maxOutputTokens, - maxReasoningTokens, - maxToolTokens, - maxToolCalls, - maxProviderCostUsd, - } - try { - requiredPaymentAmount = requiredX402Amount( - agent.pricePerTokenUsd, - maxInputTokens, - maxOutputTokens, - config.x402.currencyDecimals, - maxReasoningTokens, - maxToolTokens, - maxProviderCostUsd, - ) - } catch { - return c.json( - { - error: { - message: 'Agent payment configuration is invalid', - type: 'server_error', - code: 'invalid_payment_configuration', - }, - }, - 503, - ) - } - - // Payment / auth. - const spendAuthHeader = c.req.header('X-Payment-Signature') - const authHeader = c.req.header('Authorization') ?? '' - let consumerId: string | null = null - let paymentMethod: PaymentMethod = 'none' - let keyInfo: ApiKeyInfo | null = null - let x402Payload: Record | null = null - let paymentNonceKey: string | undefined - let mppMethod: string | undefined - let mppCredential: string | undefined - let mppPaymentIdentity: string | undefined - - if (spendAuthHeader) { - const signer = await verifyX402( - spendAuthHeader, - config.x402, - state.nonceStore, - requiredPaymentAmount, - false, - ) - if (!signer) { - await state.obs?.onAuthFailure?.(ctx, { - method: 'x402', - code: 'invalid_spend_auth', - httpStatus: 402, - }) - return c.json( - { - error: { - message: 'Invalid X-Payment-Signature', - type: 'authentication_error', - code: 'invalid_spend_auth', - required_amount: requiredPaymentAmount.toString(), - currency_decimals: config.x402.currencyDecimals ?? 6, - }, - }, - { - status: 402, - headers: { 'X-Payment-Required': 'spendauth', 'X-Request-Id': requestId }, - }, - ) - } - x402Payload = JSON.parse(spendAuthHeader) as Record - paymentNonceKey = `${String(x402Payload.commitment).toLowerCase()}:${BigInt(String(x402Payload.nonce)).toString()}` - consumerId = signer - paymentMethod = 'x402' - } else if (isMppAuthEnabled(config) && authHeader.toLowerCase().startsWith('payment ')) { - const authenticated = await verifyMppCredential( - authHeader, - config.mpp!, - config.x402, - state.nonceStore, - requiredPaymentAmount, - false, - ) - if (!authenticated) { - const realm = config.mpp!.realm - const method = config.mpp!.method ?? 'blueprintevm' - await state.obs?.onAuthFailure?.(ctx, { - method: 'mpp', - code: 'invalid_mpp_credential', - httpStatus: 401, - }) - return c.json( - { - error: { - message: 'Invalid Payment credential', - type: 'authentication_error', - code: 'invalid_mpp_credential', - }, - }, - { - status: 401, - headers: { - 'WWW-Authenticate': `Payment realm="${realm}", method="${method}"`, - 'X-Request-Id': requestId, - }, - }, - ) - } - consumerId = authenticated.consumerId - paymentMethod = 'mpp' - mppMethod = authHeader.match(/^Payment\s+(\S+)\s+/i)?.[1]?.toLowerCase() - mppCredential = mppPaymentCredential(authHeader) - mppPaymentIdentity = authenticated.paymentIdentity - x402Payload = mppPaymentPayload(authHeader) ?? null - paymentNonceKey = authenticated.replayKey - } else if (authHeader.startsWith('Bearer ')) { - const verify = config.verifyApiKey ?? (config.x402.demoMode ? defaultVerifyApiKey : null) - if (!verify || !isApiKeyAuthEnabled(config)) { - await state.obs?.onAuthFailure?.(ctx, { - method: 'apikey', - code: 'api_keys_not_configured', - httpStatus: 401, - }) - return c.json( - { error: { message: 'API key authentication is not configured', type: 'authentication_error' } }, - { status: 401, headers: { 'X-Request-Id': requestId } }, - ) - } - const key = await verify(authHeader) - if (!key) { - await state.obs?.onAuthFailure?.(ctx, { - method: 'apikey', - code: 'invalid_api_key', - httpStatus: 401, - }) - return c.json( - { error: { message: 'Invalid API key', type: 'authentication_error' } }, - { status: 401, headers: { 'X-Request-Id': requestId } }, - ) - } - if (key.scopes && key.scopes.length > 0 && !key.scopes.includes(state.requiredScope)) { - await state.obs?.onAuthFailure?.(ctx, { - method: 'apikey', - code: 'insufficient_scope', - httpStatus: 403, - }) - return c.json( - { - error: { - message: `API key missing required scope: ${state.requiredScope}`, - type: 'forbidden', - code: 'insufficient_scope', - }, - }, - { status: 403, headers: { 'X-Request-Id': requestId } }, - ) - } - consumerId = key.consumerId - paymentMethod = 'apikey' - keyInfo = key - } else { - await state.obs?.onAuthFailure?.(ctx, { - method: 'none', - code: 'payment_required', - httpStatus: 402, - }) - const methods: string[] = ['x402'] - if (isMppAuthEnabled(config)) methods.push('mpp') - if (isApiKeyAuthEnabled(config)) methods.push('api_key') - const headers: Record = { - 'X-Payment-Required': methods.join(', '), - 'X-Request-Id': requestId, - } - if (isMppAuthEnabled(config) && config.mpp) { - headers['WWW-Authenticate'] = - `Payment realm="${config.mpp.realm}", method="${config.mpp.method ?? 'blueprintevm'}"` - } - return c.json( - { - error: { - message: 'Payment required', - type: 'payment_required', - payment_methods: methods, - x402: { - operator: config.x402.operatorAddress, - chain_id: config.x402.chainId, - credits_address: config.x402.creditsAddress, - required_amount: requiredPaymentAmount.toString(), - currency_decimals: config.x402.currencyDecimals ?? 6, - max_output_tokens: maxOutputTokens, - }, - ...(isMppAuthEnabled(config) && config.mpp - ? { mpp: { realm: config.mpp.realm, method: config.mpp.method ?? 'blueprintevm' } } - : {}), - ...(isApiKeyAuthEnabled(config) - ? { - api_key: { - purchase_url: config.baseUrl - ? `${config.baseUrl}/agents/${slug}/api-keys` - : undefined, - }, - } - : {}), - }, - }, - { status: 402, headers }, - ) - } - - // Rate limit. - const effectiveRateLimit = keyInfo?.rateLimitPerMinute - ? { limit: keyInfo.rateLimitPerMinute, windowSeconds: 60 } - : state.globalRateLimit - const rl = await checkRateLimit(consumerId, effectiveRateLimit, state.rateLimitStore) - if (!rl.allowed) { - await state.obs?.onRateLimited?.(ctx, { - consumerId: consumerId, - retryAfterSeconds: rl.retryAfterSeconds ?? 60, - }) - return c.json( - { - error: { - message: 'Rate limit exceeded', - type: 'rate_limit_error', - retry_after: rl.retryAfterSeconds, - }, - }, - { - status: 429, - headers: { - 'Retry-After': String(rl.retryAfterSeconds ?? 60), - 'X-Request-Id': requestId, - }, - }, - ) - } - - // Reject or report injection only after authentication so observer events - // retain the authenticated consumer identity. - if (injectionWarnings.length > 0) { - await state.obs?.onInjectionDetected?.(ctx, { - consumerId: consumerId, - patterns: injectionWarnings, - blocked: !!config.blockInjection, - }) - if (config.blockInjection) { - return c.json( - { - error: { - message: 'Request rejected: potential prompt injection detected', - type: 'content_policy_violation', - }, - }, - { status: 400, headers: { 'X-Request-Id': requestId } }, - ) - } - } - - if (config.authorizeConsumer) { - const authz = await config.authorizeConsumer(agent, { - method: paymentMethod, - consumerId: consumerId, - keyId: keyInfo?.keyId, - requestId, - }) - if (!authz.allow) { - return c.json( - { - error: { - message: authz.reason, - type: 'authorization_denied', - code: authz.code, - }, - }, - { status: 403, headers: { 'X-Request-Id': requestId } }, - ) - } - } - - return { - agent, - consumerId, - paymentMethod, - keyInfo, - userMessage, - rateLimitRemaining: rl.remaining, - requestId, - startMs, - maxOutputTokens, - executionBudget, - requiredPaymentAmount, - paymentPayload: x402Payload, - paymentNonceKey, - mppMethod, - mppCredential, - mppPaymentIdentity, - } -} - -/** Claim payment ownership after every request guard has accepted the call. */ -export async function claimPayment( - authz: AuthorizedRequest, - config: GatewayConfig, - state: GatewayState, - hooks: PaymentClaimHooks = {}, -): Promise { - assertX402V1SettlementSafe(authz, config) - if (authz.paymentMethod === 'x402' && authz.paymentPayload) { - const context = paymentAuthorizationContext(authz) - if (config.x402.paymentProtocolVersion === 2) { - await preparePaymentRecovery(authz, config, { - kind: 'x402', - operationId: `x402:${paymentNonceKey(authz.paymentPayload)}`, - }, hooks) - } - let operation: PaymentOperation | undefined - if (config.x402.authorizePayment) { - if (config.x402.paymentProtocolVersion !== 2 && !config.x402.demoMode) { - throw new Error( - 'production x402 version 1 cannot use authorizePayment; ' + - 'use paymentProtocolVersion: 2 with paymentOperations', - ) - } - // Version 1 has no durable operation to release if another request wins - // the shared nonce while this callback is still running. This callback - // remains only for explicit demo-mode compatibility; production callers - // must use the durable version 2 operation lifecycle. - const legacyClaimed = config.x402.paymentProtocolVersion !== 2 && authz.paymentNonceKey - ? await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) - : undefined - if (legacyClaimed === false) throw new Error('payment nonce was already consumed') - const result = await config.x402.authorizePayment(authz.paymentPayload, context) - if (!result) throw new Error('payment authorization was rejected') - if (typeof result !== 'boolean') { - operation = result - } - else if (config.x402.paymentProtocolVersion === 2) { - throw new Error('version 2 payment authorization did not return an operation') - } else if (authz.paymentNonceKey && legacyClaimed === undefined) { - const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) - if (!claimed) throw new Error('payment nonce was already consumed') - } - } else if (config.x402.paymentOperations) { - operation = await config.x402.paymentOperations.claimPayment(authz.paymentPayload, context) - } else if (authz.paymentNonceKey) { - const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload) - if (!claimed) throw new Error('payment nonce was already consumed') - } - if (operation && operation.protocolVersion !== 2) { - throw new Error('payment operation protocol version mismatch') - } - if ( - operation && - config.x402.paymentProtocolVersion === 2 && - operation.operationId !== authz.paymentRecoveryId - ) { - throw new Error('x402 payment operation identity mismatch') - } - if (operation && !config.x402.paymentOperations) { - throw new Error('durable payment operations are required to settle a claimed operation') - } - if (operation) { - authz.paymentOperationAcquired = operation.acquiredByRequestId === context.requestId - if (!authz.paymentOperationAcquired) { - throw new Error('payment operation was already claimed') - } - // Attach owned state before the shared nonce claim. If that claim fails, - // the caller can persist release recovery after an ambiguous refund. - authz.paymentOperation = operation - await markRecoveryClaimed(authz, config) - } - if (operation && authz.paymentNonceKey) { - const claimed = await claimPaymentNonce( - state.nonceStore, - authz.paymentNonceKey, - authz.paymentPayload, - `${operation.operationId}:${context.requestId}`, - ) - if (!claimed) { - try { - await releasePayment(authz, config, 'shared payment nonce was already owned') - } catch (releaseError) { - console.error( - `[agent-gateway] payment release failed for ${authz.requestId}:`, - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) - } - throw new Error('payment nonce was already consumed') - } - } - } else if (authz.paymentMethod === 'mpp') { - if (!authz.paymentNonceKey) { - throw new Error('MPP payment has no replay identity') - } - const mppMethod = (authz.mppMethod ?? config.mpp?.method ?? 'blueprintevm').toLowerCase() - const durablePayload = durableMppPaymentPayload(authz.paymentPayload) - // Only BlueprinTEVM carries x402 authorization fields. Other MPP methods - // use the isolated immediate-charge lifecycle below. - if (mppMethod === 'blueprintevm' && durablePayload && config.x402.paymentOperations) { - const context = paymentAuthorizationContext(authz) - await preparePaymentRecovery(authz, config, { - kind: 'x402', - operationId: `x402:${paymentNonceKey(durablePayload)}`, - }, hooks) - const operation = await config.x402.paymentOperations.claimPayment(durablePayload, context) - if (operation.protocolVersion !== 2) throw new Error('payment operation protocol version mismatch') - if (operation.operationId !== authz.paymentRecoveryId) { - throw new Error('x402 payment operation identity mismatch') - } - if (operation.acquiredByRequestId !== context.requestId) { - throw new Error('payment operation was already claimed') - } - authz.paymentPayload = durablePayload - authz.paymentOperation = operation - authz.paymentOperationAcquired = true - await markRecoveryClaimed(authz, config) - const claimed = await claimPaymentNonce( - state.nonceStore, - authz.paymentNonceKey, - durablePayload, - `${operation.operationId}:${context.requestId}`, - ) - if (!claimed) { - try { - await releasePayment(authz, config, 'shared payment nonce was already owned') - } catch (releaseError) { - console.error( - `[agent-gateway] payment release failed for ${authz.requestId}:`, - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) - } - throw new Error('payment nonce was already consumed') - } - } else if (mppMethod === 'blueprintevm') { - const claimed = await claimPaymentNonce(state.nonceStore, authz.paymentNonceKey, authz.paymentPayload ?? {}) - if (!claimed) throw new Error('payment nonce was already consumed') - } else { - const lifecycle = config.mpp?.charge - if (!lifecycle || lifecycle.protocolVersion !== 1) { - throw new Error('MPP charge lifecycle is not configured') - } - if (!authz.mppCredential) throw new Error('MPP payment credential is unavailable') - if (!authz.mppPaymentIdentity) throw new Error('MPP payment identity is unavailable') - const operationId = await mppPaymentOperationId(mppMethod, authz.mppPaymentIdentity) - await preparePaymentRecovery(authz, config, { - kind: 'mpp-charge', - method: mppMethod, - operationId, - }, hooks) - const claimed = await claimPaymentNonce( - state.nonceStore, - authz.paymentNonceKey, - authz.paymentPayload ?? {}, - `${operationId}:${authz.requestId}`, - ) - if (!claimed) { - await markRecoveryReconciled(authz, config) - throw new Error('payment nonce was already consumed') - } - const operation = await lifecycle.confirmPayment({ - operationId, - requestId: authz.requestId, - agentId: authz.agent.id, - consumerId: authz.consumerId, - method: mppMethod, - credential: authz.mppCredential, - amount: authz.requiredPaymentAmount, - currencyDecimals: config.x402.currencyDecimals ?? 6, - }) - assertMppChargeOperation( - operation, - { operationId, requestId: authz.requestId, method: mppMethod }, - ['confirmed'], - false, - ) - authz.mppChargeOperation = operation - await markRecoveryClaimed(authz, config) - assertMppChargeOperation( - operation, - { operationId, requestId: authz.requestId, method: mppMethod }, - ['confirmed'], - ) - } - } - - try { - await state.obs?.onPaymentVerified?.( - { - requestId: authz.requestId, - agentSlug: authz.agent.slug, - startMs: authz.startMs, - }, - { - method: authz.paymentMethod, - consumerId: authz.consumerId, - keyId: authz.keyInfo?.keyId, - }, - ) - } catch (error) { - // Observability must not turn a durable claim into a stranded payment. - console.error( - '[agent-gateway] payment observer failed for ' + authz.requestId + ':', - error instanceof Error ? error.message : String(error), - ) - } -} - -function paymentAuthorizationContext(authz: AuthorizedRequest) { - return { - requestId: authz.requestId, - agentId: authz.agent.id, - requiredAmount: authz.requiredPaymentAmount, - maxOutputTokens: authz.maxOutputTokens, - executionBudget: authz.executionBudget, - } -} - -async function preparePaymentRecovery( - authz: AuthorizedRequest, - config: GatewayConfig, - payment: PaymentRecoveryTarget, - hooks: PaymentClaimHooks, -): Promise { - const recovery = config.paymentRecovery - if (!recovery) throw new Error('durable payment recovery is not configured') - const now = Date.now() - const fenceId = globalThis.crypto.randomUUID() - const leaseExpiresAt = now + recoveryTiming(recovery).staleRequestMs - const record: PaymentRecoveryRecord = { - version: PAYMENT_RECOVERY_VERSION, - id: payment.operationId, - revision: 0, - state: 'claiming', - payment, - attribution: { - requestId: authz.requestId, - agentId: authz.agent.id, - agentSlug: authz.agent.slug, - consumerId: authz.consumerId, - paymentMethod: authz.paymentMethod, - startMs: authz.startMs, - pricePerTokenUsd: authz.agent.pricePerTokenUsd, - platformFeePercent: authz.agent.platformFeePercent, - requiredAmount: authz.requiredPaymentAmount.toString(), - currencyDecimals: config.x402.currencyDecimals ?? 6, - maxOutputTokens: authz.maxOutputTokens, - executionBudget: authz.executionBudget, - }, - workStarted: false, - usageRecorded: false, - attempts: 0, - nextAttemptAt: leaseExpiresAt, - lease: { id: fenceId, expiresAt: leaseExpiresAt }, - createdAt: now, - updatedAt: now, - } - if (!await recovery.store.createIfAbsent(record)) { - if ((await recovery.store.get(record.id))?.state === 'reconciled') { - throw new PaymentRecoveryReplayError(record.id) - } - throw new Error('payment recovery identity was already claimed') - } - authz.paymentRecoveryId = record.id - authz.paymentRecoveryFence = fenceId - await hooks.onRecoveryPrepared?.(record.id) -} - -async function markRecoveryClaimed( - authz: AuthorizedRequest, - config: GatewayConfig, -): Promise { - const recovery = config.paymentRecovery - if (!recovery || !authz.paymentRecoveryId) return - const fenceId = requirePaymentRecoveryFence(authz) - const now = Date.now() - const leaseExpiresAt = now + recoveryTiming(recovery).staleRequestMs - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ - ...record, - state: 'claimed', - payment: recoveryTarget(authz, record.payment), - lease: { id: fenceId, expiresAt: leaseExpiresAt }, - nextAttemptAt: leaseExpiresAt, - }), now) -} - -function requirePaymentRecoveryFence(authz: AuthorizedRequest): string { - if (!authz.paymentRecoveryFence) { - throw new Error('payment recovery fence is unavailable') - } - return authz.paymentRecoveryFence -} - -function recoveryTarget( - authz: AuthorizedRequest, - current: PaymentRecoveryTarget, -): PaymentRecoveryTarget { - if (authz.paymentOperation) { - return { - kind: 'x402', - operationId: authz.paymentOperation.operationId, - operation: serializePaymentOperation(authz.paymentOperation), - } - } - if (authz.mppChargeOperation) { - return { - kind: 'mpp-charge', - method: authz.mppChargeOperation.method, - operationId: authz.mppChargeOperation.operationId, - operation: authz.mppChargeOperation, - } - } - return current -} - -/** Release an owned operation when execution cannot produce a valid receipt. */ -export async function releasePayment( - authz: AuthorizedRequest, - config: GatewayConfig, - reason: string, -): Promise { - const ownsX402 = authz.paymentOperation && - authz.paymentOperationAcquired === true && - config.x402.paymentOperations - const ownsMpp = authz.mppChargeOperation && config.mpp?.charge - if (!ownsX402 && !ownsMpp) { - await relinquishPaymentRecovery(authz, config, Date.now()) - return - } - let reconciled = false - try { - await markRecoveryReleasing(authz, config, reason) - if (ownsX402) { - authz.paymentOperation = await config.x402.paymentOperations!.releasePayment( - authz.paymentOperation!, - reason, - ) - } else { - const operation = await config.mpp!.charge!.releasePayment(authz.mppChargeOperation!, reason) - assertMppChargeOperation( - operation, - { - operationId: authz.mppChargeOperation!.operationId, - requestId: authz.requestId, - method: authz.mppChargeOperation!.method, - }, - ['released'], - false, - ) - authz.mppChargeOperation = operation - } - await markRecoveryReconciled(authz, config) - reconciled = true - } finally { - if (!reconciled) { - try { - await relinquishPaymentRecovery(authz, config, Date.now()) - } catch { - // Preserve the original provider or metadata error. A later worker - // retry still has the durable row when cleanup itself is unavailable. - } - } - } -} - -/** Mark a durable reservation active immediately before sandbox execution. */ -export async function beginPaymentExecution( - authz: AuthorizedRequest, - config: GatewayConfig, -): Promise { - if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { - authz.paymentOperation = await config.x402.paymentOperations.beginPaymentExecution(authz.paymentOperation) - } -} - -/** Persist the sandbox handoff immediately before the adapter call. */ -export async function markPaymentExecutionStarted( - authz: AuthorizedRequest, - config: GatewayConfig, -): Promise { - if (!authz.paymentRecoveryId) return - const recovery = config.paymentRecovery - if (!recovery) throw new Error('durable payment recovery is not configured') - const fenceId = requirePaymentRecoveryFence(authz) - const now = Date.now() - const fallbackAt = now + recoveryTiming(recovery).receiptTimeoutMs - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ - ...record, - state: 'executing', - payment: recoveryTarget(authz, record.payment), - workStarted: true, - fallbackAt, - lease: { id: fenceId, expiresAt: fallbackAt }, - nextAttemptAt: fallbackAt, - }), now) -} - -/** Renew the live execution lease while a provider stream is still open. */ -export async function renewPaymentExecution( - authz: AuthorizedRequest, - config: GatewayConfig, -): Promise { - if (!authz.paymentRecoveryId) return - const recovery = config.paymentRecovery - if (!recovery) throw new Error('durable payment recovery is not configured') - const fenceId = requirePaymentRecoveryFence(authz) - const now = Date.now() - const fallbackAt = now + recoveryTiming(recovery).receiptTimeoutMs - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ - ...record, - state: 'executing', - fallbackAt, - lease: { id: fenceId, expiresAt: fallbackAt }, - nextAttemptAt: fallbackAt, - }), now) -} - -/** - * Release only when no sandbox work was observed. Once output or a receipt - * exists, retain the owner for settlement or background recovery. - */ -export async function releasePaymentAfterFailure( - authz: AuthorizedRequest, - config: GatewayConfig, - reason: string, - workObserved: boolean, -): Promise { - if (workObserved) { - const recovery = config.paymentRecovery - if (recovery && authz.paymentRecoveryId) { - const fenceId = requirePaymentRecoveryFence(authz) - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => { - if (record.state === 'settling' && record.usage) { - return { ...record, lease: undefined, nextAttemptAt: Date.now() } - } - const fallbackAt = record.fallbackAt ?? - Date.now() + recoveryTiming(recovery).receiptTimeoutMs - return { - ...record, - state: 'retained', - payment: recoveryTarget(authz, record.payment), - workStarted: true, - fallbackAt, - reason, - lease: undefined, - nextAttemptAt: fallbackAt, - } - }) - } - if (authz.paymentOperation && authz.paymentOperationAcquired === true && config.x402.paymentOperations) { - authz.paymentOperation = await config.x402.paymentOperations.retainPayment(authz.paymentOperation, reason) - } - console.error( - `[agent-gateway] retaining payment ownership after sandbox work for ${authz.requestId}: ${reason}`, - ) - return - } - await releasePayment(authz, config, reason) -} - -async function relinquishPaymentRecovery( - authz: AuthorizedRequest, - config: GatewayConfig, - nextAttemptAt: number, -): Promise { - const recovery = config.paymentRecovery - if (!recovery || !authz.paymentRecoveryId || !authz.paymentRecoveryFence) return - try { - await updateOwnedPaymentRecovery( - recovery.store, - authz.paymentRecoveryId, - authz.paymentRecoveryFence, - (record) => ({ ...record, lease: undefined, nextAttemptAt }), - ) - } catch (error) { - if (!(error instanceof PaymentRecoveryFenceError)) throw error - } -} - -async function markRecoveryReleasing( - authz: AuthorizedRequest, - config: GatewayConfig, - reason: string, -): Promise { - const recovery = config.paymentRecovery - if (!recovery || !authz.paymentRecoveryId) return - const fenceId = requirePaymentRecoveryFence(authz) - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ - ...record, - state: 'releasing', - payment: recoveryTarget(authz, record.payment), - reason, - nextAttemptAt: Date.now(), - })) -} - -async function markRecoveryReconciled( - authz: AuthorizedRequest, - config: GatewayConfig, -): Promise { - const recovery = config.paymentRecovery - if (!recovery || !authz.paymentRecoveryId) return - const fenceId = requirePaymentRecoveryFence(authz) - const now = Date.now() - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ - ...record, - state: 'reconciled', - payment: recoveryTarget(authz, record.payment), - lease: undefined, - lastError: undefined, - nextAttemptAt: Number.MAX_SAFE_INTEGER, - reconciledAt: now, - }), now) -} - -export async function reclaimPayment( - operationId: string, - config: GatewayConfig, -): Promise { - if (!config.x402.paymentOperations) throw new Error('durable payment operations are not configured') - return config.x402.paymentOperations.reclaimPayment(operationId) -} - -async function claimPaymentNonce( - nonceStore: NonceStore, - nonceKey: string, - payload: Record, - ownerId?: string, -): Promise { - const expiry = payload.expiry === undefined - ? BigInt(Math.floor(Date.now() / 1000) + 3600) - : BigInt(String(payload.expiry)) - const ttl = nonceTtlSeconds(expiry) - if (ttl === undefined) return false - return claimStoredNonce(nonceStore, nonceKey, ttl, ownerId) -} - -function durableMppPaymentPayload( - payload: Record | null, -): Record | undefined { - if (!payload) return undefined - const commitment = payload.commitment ?? payload.from - const amount = payload.amount ?? payload.value - const nonce = payload.nonce - if (typeof commitment !== 'string' || commitment.length === 0) return undefined - if (amount === undefined || nonce === undefined) return undefined - const amountText = String(amount) - const nonceText = String(nonce) - if (!/^\d+$/.test(amountText) || !/^\d+$/.test(nonceText)) return undefined - const expiryText = payload.expiry === undefined - ? String(Math.floor(Date.now() / 1000) + 3600) - : String(payload.expiry) - if (!/^\d+$/.test(expiryText)) return undefined - return { - ...payload, - commitment, - amount: amountText, - nonce: nonceText, - expiry: expiryText, - } -} - -/** - * Yield the inner sandbox's response as text deltas, applying the - * system-prompt redaction filter on each delta so leakage of the agent's - * system prompt back through the model's output is suppressed identically - * whether the caller is on the OpenAI-compat path or A2A. - * - * Aborts when `signal` fires (used by A2A `tasks/cancel`). - */ -export async function* dispatchSandboxStream( - agent: AgentMeta, - userMessage: string, - consumerId: string, - config: GatewayConfig, - signal?: AbortSignal, - sessionId?: string, - maxOutputTokens?: number, -): AsyncIterable { - for await (const event of dispatchSandboxStreamRich( - agent, - userMessage, - consumerId, - config, - signal, - sessionId, - maxOutputTokens, - )) { - if (event.kind === 'text') yield event.delta - } -} - -/** - * A2A-shaped dispatch event. Distinguishes text deltas from sandbox-signalled - * pause-for-input events. The A2A handler uses this richer variant so it can - * emit `input-required` status updates; the OpenAI-compat path consumes the - * text-only `dispatchSandboxStream` adapter above. - */ -export type A2ADispatchEvent = - | { kind: 'text'; delta: string } - | { kind: 'input-required'; prompt?: string } - | { kind: 'activity' } - | { kind: 'usage'; usage: SandboxUsageReceipt } - -/** - * Like `dispatchSandboxStream` but yields a discriminated union so callers can - * react to `input-required` signals from the sandbox. The sandbox opts in by - * emitting `{ type: 'input-required', data: { inputRequired: { prompt? } } }` - * (or by setting `data.inputRequired` on any event); sandboxes that don't - * emit such events see identical behavior. - * - * `sessionId` defaults to `consumer:` matching the existing single-turn - * path; multi-turn continuations pass an explicit `taskId` so the sandbox can - * keep per-task conversation memory. - */ -export async function* dispatchSandboxStreamRich( - agent: AgentMeta, - userMessage: string, - consumerId: string, - config: GatewayConfig, - signal?: AbortSignal, - sessionId?: string, - maxOutputTokens?: number, - onExecutionStart?: () => Promise, - requiresReceipt = config.x402.paymentOperations !== undefined, - onSandboxStart?: () => void | Promise, - maxInputTokens?: number, - onExecutionHeartbeat?: () => Promise, -): AsyncIterable { - if (signal?.aborted) return - const box = await config.getSandbox(agent) - if (signal?.aborted) return - const outputLimit = maxOutputTokens ?? config.defaultOutputTokens ?? 1024 - if (!Number.isSafeInteger(outputLimit) || outputLimit <= 0) { - throw new Error('max output tokens must be a positive safe integer') - } - let outputBytes = 0 - // Bound untrusted adapters while the final receipt is pending. The receipt - // remains authoritative for token count, so over-limit output is never sent. - const maxOutputBytes = outputLimit * 4 - if (!Number.isSafeInteger(maxOutputBytes)) { - throw new Error('max output token bound exceeds safe integer range') - } - const encoder = new TextEncoder() - let usageParts: Partial = {} - let observedReasoningTokens = 0 - let observedToolTokens = 0 - let observedToolCalls = 0 - let legacyOutputText = '' - 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 { - // 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, - }) - iterator = promptStream[Symbol.asyncIterator]() - const heartbeatMs = onExecutionHeartbeat - ? Math.max(100, Math.min( - Math.floor((config.paymentRecovery?.receiptTimeoutMs ?? 5 * 60_000) / 3), - 5_000, - )) - : 0 - if (onExecutionHeartbeat) { - heartbeatTimer = setInterval(() => { - if (heartbeatInFlight || heartbeatError !== undefined) return - heartbeatInFlight = onExecutionHeartbeat() - .catch((error: unknown) => { - heartbeatError = error - executionController.abort() - }) - .finally(() => { - heartbeatInFlight = undefined - }) - }, heartbeatMs) - } - while (true) { - const next = await readSandboxEvent(iterator, executionController.signal) - if (next === ABORTED_SANDBOX_READ) { - if (heartbeatError !== undefined) throw heartbeatError - return - } - if (next.done) break - const event = next.value - if (event.data?.usage) usageParts = mergeUsage(usageParts, event.data.usage) - if (event.data?.reasoning?.tokens !== undefined) { - observedReasoningTokens += nonNegativeSafeInteger(event.data.reasoning.tokens, 'reasoning tokens') - yield { kind: 'activity' } - } - if (event.data?.tool) { - observedToolCalls += 1 - observedToolTokens += - nonNegativeSafeInteger(event.data.tool.inputTokens ?? 0, 'tool input tokens') + - nonNegativeSafeInteger(event.data.tool.outputTokens ?? 0, 'tool output tokens') - yield { kind: 'activity' } - } - enforceUsageBudget(withObservedUsage( - usageParts, - observedReasoningTokens, - observedToolTokens, - observedToolCalls, - ), executionBudget) - if ( - event.type === 'message.part.updated' && - event.data?.part?.type === 'text' && - event.data.delta - ) { - const remainingBytes = maxOutputBytes - outputBytes - if (remainingBytes <= 0) throw new Error('sandbox exceeded max output tokens') - const bounded = truncateUtf8(event.data.delta, remainingBytes, encoder) - if (bounded.truncated) { - yield { kind: 'activity' } - throw new Error('sandbox exceeded max output tokens') - } - outputBytes += bounded.bytes - legacyOutputText += bounded.text - yield { kind: 'activity' } - yield { kind: 'text', delta: redactSystemPromptFromOutput(bounded.text, agent.systemPrompt) } - continue - } - if (event.type === 'input-required' || event.data?.inputRequired) { - const usage = completeUsage( - usageParts, - observedReasoningTokens, - observedToolTokens, - observedToolCalls, - userMessage, - legacyOutputText, - executionBudget, - requiresReceipt, - ) - yield { kind: 'input-required', prompt: event.data?.inputRequired?.prompt } - // Terminal for the sandbox stream — sandbox SHOULD stop emitting until - // the gateway dispatches a continuation message with the new user input. - yield { kind: 'usage', usage } - return - } - } - const usage = completeUsage( - usageParts, - observedReasoningTokens, - observedToolTokens, - observedToolCalls, - userMessage, - legacyOutputText, - executionBudget, - requiresReceipt, - ) - yield { kind: 'usage', usage } - } finally { - if (heartbeatTimer !== undefined) clearInterval(heartbeatTimer) - const pendingHeartbeat = heartbeatInFlight - if (pendingHeartbeat) await pendingHeartbeat - signal?.removeEventListener('abort', forwardAbort) - if (iterator) await closeSandboxIterator(iterator) - if (heartbeatError !== undefined && !signal?.aborted) throw heartbeatError - } -} - -const ABORTED_SANDBOX_READ = Symbol('aborted-sandbox-read') - -async function readSandboxEvent( - iterator: AsyncIterator, - signal?: AbortSignal, -): Promise | typeof ABORTED_SANDBOX_READ> { - if (!signal) return iterator.next() - if (signal.aborted) return ABORTED_SANDBOX_READ - return new Promise((resolve, reject) => { - const onAbort = () => { - signal.removeEventListener('abort', onAbort) - resolve(ABORTED_SANDBOX_READ) - } - signal.addEventListener('abort', onAbort, { once: true }) - iterator.next().then( - (result) => { - signal.removeEventListener('abort', onAbort) - resolve(result) - }, - (error: unknown) => { - signal.removeEventListener('abort', onAbort) - reject(error) - }, - ) - }) -} - -const SANDBOX_CLEANUP_TIMEOUT_MS = 50 - -async function closeSandboxIterator(iterator: AsyncIterator): Promise { - let closing: PromiseLike | undefined - try { - const result = iterator.return?.() - if (result) closing = Promise.resolve(result) - } catch { - return - } - if (!closing) return - let timeout: ReturnType | undefined - try { - await Promise.race([ - Promise.resolve(closing).catch(() => undefined), - new Promise((resolve) => { - timeout = setTimeout(resolve, SANDBOX_CLEANUP_TIMEOUT_MS) - }), - ]) - } finally { - if (timeout !== undefined) clearTimeout(timeout) - } -} - -/** - * Record usage event + settle payment + invoke the observer. Both wire - * formats call this once their stream has drained, so settlement happens - * exactly once per request regardless of protocol. - */ -export interface SettleAndRecordOptions { - /** Skip attribution after a durable finalization marker confirms it ran. */ - usageAlreadyRecorded?: boolean - /** Skip provider settlement only after an authoritative read found it settled. */ - paymentAlreadySettled?: boolean - /** Persist the caller's recovery marker after attribution succeeds. */ - onUsageRecorded?: () => Promise - /** Recovery uses the original quoted ceiling when no receipt arrives. */ - settlementBasis?: PaymentSettlementBasis - /** Exact base-unit charge selected by the recovery policy. */ - paymentAmount?: bigint -} - -export async function settleAndRecord( - agent: AgentMeta, - authz: AuthorizedRequest, - usage: SandboxUsageReceipt, - config: GatewayConfig, - obs: GatewayObserver | undefined, - options: SettleAndRecordOptions = {}, -): Promise { - assertX402V1SettlementSafe(authz, config) - const settlementBasis = options.settlementBasis ?? 'usage-receipt' - await markRecoverySettling(authz, usage, settlementBasis, config) - if (options.usageAlreadyRecorded) await markRecoveryUsageRecorded(authz, config) - const tokenCost = ( - usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens - ) * agent.pricePerTokenUsd - const totalCost = Math.max(tokenCost, usage.providerCostUsd) - const ownerEarned = totalCost * (1 - agent.platformFeePercent) - const platformFee = totalCost * agent.platformFeePercent - const usageEvent = { - requestId: authz.requestId, - agentId: agent.id, - agentSlug: agent.slug, - consumerId: authz.consumerId, - paymentMethod: authz.paymentMethod, - inputTokens: usage.inputTokens, - outputTokens: usage.outputTokens, - reasoningTokens: usage.reasoningTokens, - toolTokens: usage.toolTokens, - toolCallCount: usage.toolCallCount, - providerCostUsd: usage.providerCostUsd, - totalCostUsd: totalCost, - ownerEarnedUsd: ownerEarned, - platformFeeUsd: platformFee, - durationMs: Date.now() - authz.startMs, - settlementBasis, - } - const ctx: RequestContext = { - requestId: authz.requestId, - agentSlug: agent.slug, - startMs: authz.startMs, - } - try { - if (authz.paymentOperation && config.x402.paymentOperations) { - const amount = options.paymentAmount ?? actualX402Amount( - agent.pricePerTokenUsd, - usage.inputTokens, - usage.outputTokens, - usage.reasoningTokens, - usage.toolTokens, - config.x402.currencyDecimals, - usage.providerCostUsd, - ) - if (options.paymentAlreadySettled) { - if ( - authz.paymentOperation.state !== 'settled' || - authz.paymentOperation.settledAmount !== amount - ) { - throw new Error('authoritative payment state does not match finalization') - } - } else { - authz.paymentOperation = await config.x402.paymentOperations.settlePayment( - authz.paymentOperation, - { amount, totalCostUsd: totalCost, usage, basis: settlementBasis }, - ) - } - // Durable settlement happens first. If attribution storage is - // unavailable, recovery must never refund delivered work. - if (!options.usageAlreadyRecorded) { - await config.recordUsage(usageEvent) - await options.onUsageRecorded?.() - await markRecoveryUsageRecorded(authz, config) - } - } else if (authz.mppChargeOperation) { - // Generic MPP charge methods confirm before the response. Finalization - // records attribution only; it never invokes the legacy settlement hook. - if (!options.usageAlreadyRecorded) { - await config.recordUsage(usageEvent) - await options.onUsageRecorded?.() - await markRecoveryUsageRecorded(authz, config) - } - } else { - // Legacy adapters retain attribution-before-charge because their - // settlement callback may resolve that usage row. - if (!options.usageAlreadyRecorded) { - await config.recordUsage(usageEvent) - await options.onUsageRecorded?.() - } - if (config.settlePayment) { - await config.settlePayment( - { - method: authz.paymentMethod, - consumerId: authz.consumerId, - requestId: authz.requestId, - }, - totalCost, - ) - } - } - await markRecoveryReconciled(authz, config) - } catch (err) { - const msg = err instanceof Error ? err.message : String(err) - console.error(`[agent-gateway] settlement failed for ${authz.consumerId}: ${msg}`) - await obs?.onSettlementError?.(ctx, { - consumerId: authz.consumerId, - method: authz.paymentMethod, - errorMessage: msg, - }) - throw err - } - try { - await obs?.onRequestComplete?.(ctx, usageEvent) - } catch (error) { - console.error( - `[agent-gateway] completion observer failed for ${authz.requestId}:`, - error instanceof Error ? error.message : String(error), - ) - } -} - -function assertX402V1SettlementSafe(authz: AuthorizedRequest, config: GatewayConfig): void { - if ( - authz.paymentMethod === 'x402' && - config.x402.paymentProtocolVersion !== 2 && - !config.x402.demoMode && - config.settlePayment - ) { - throw new Error( - 'production x402 version 1 cannot use settlePayment; ' + - 'use paymentProtocolVersion: 2 with paymentOperations', - ) - } -} - -async function markRecoverySettling( - authz: AuthorizedRequest, - usage: SandboxUsageReceipt, - settlementBasis: PaymentSettlementBasis, - config: GatewayConfig, -): Promise { - const recovery = config.paymentRecovery - if (!recovery || !authz.paymentRecoveryId) return - const fenceId = requirePaymentRecoveryFence(authz) - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => { - const next: PaymentRecoveryRecord = { - ...record, - state: 'settling', - payment: recoveryTarget(authz, record.payment), - workStarted: true, - settlementBasis, - nextAttemptAt: Date.now(), - } - // A quoted-ceiling settlement has no provider receipt. Keep the durable - // basis and original amount, then rebuild the synthetic accounting input - // on each retry instead of persisting a lossy floating-point surrogate. - if (settlementBasis !== 'quoted-ceiling' || record.usage !== undefined) { - next.usage = usage - } else { - delete next.usage - } - return next - }) -} - -async function markRecoveryUsageRecorded( - authz: AuthorizedRequest, - config: GatewayConfig, -): Promise { - const recovery = config.paymentRecovery - if (!recovery || !authz.paymentRecoveryId) return - const fenceId = requirePaymentRecoveryFence(authz) - await updateOwnedPaymentRecovery(recovery.store, authz.paymentRecoveryId, fenceId, (record) => ({ - ...record, - usageRecorded: true, - })) -} - -function mergeUsage( - current: Partial, - update: Partial, -): Partial { - const merged = { ...current, ...update } - for (const key of ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount', 'providerCostUsd'] as const) { - const value = update[key] - if (value !== undefined && (!Number.isFinite(value) || value < 0)) { - throw new Error(`sandbox usage field ${key} is invalid`) - } - if (value !== undefined && current[key] !== undefined) { - // Usage events are cumulative receipts. Never let a later partial or - // final event erase spend observed earlier in the same execution. - merged[key] = Math.max(current[key]!, value) - } - } - if (current.budgetEnforced === false || update.budgetEnforced === false) { - merged.budgetEnforced = false - } - return merged -} - -function withObservedUsage( - usage: Partial, - reasoningTokens: number, - toolTokens: number, - toolCallCount: number, -): Partial { - return { - ...usage, - ...(usage.reasoningTokens !== undefined || reasoningTokens > 0 - ? { reasoningTokens: Math.max(usage.reasoningTokens ?? 0, reasoningTokens) } - : {}), - ...(usage.toolTokens !== undefined || toolTokens > 0 - ? { toolTokens: Math.max(usage.toolTokens ?? 0, toolTokens) } - : {}), - ...(usage.toolCallCount !== undefined || toolCallCount > 0 - ? { toolCallCount: Math.max(usage.toolCallCount ?? 0, toolCallCount) } - : {}), - } -} - -function nonNegativeSafeInteger(value: number, name: string): number { - if (!Number.isSafeInteger(value) || value < 0) { - throw new Error(`sandbox ${name} is invalid`) - } - return value -} - -function enforceUsageBudget( - usage: Partial, - budget: SandboxExecutionBudget, -): void { - if (usage.inputTokens !== undefined && usage.inputTokens > budget.maxInputTokens) { - throw new Error('sandbox exceeded max input tokens') - } - if (usage.outputTokens !== undefined && usage.outputTokens > budget.maxOutputTokens) { - throw new Error('sandbox exceeded max output tokens') - } - if (usage.reasoningTokens !== undefined && usage.reasoningTokens > budget.maxReasoningTokens) { - throw new Error('sandbox exceeded max reasoning tokens') - } - if (usage.toolTokens !== undefined && usage.toolTokens > budget.maxToolTokens) { - throw new Error('sandbox exceeded max tool tokens') - } - if (usage.toolCallCount !== undefined && usage.toolCallCount > budget.maxToolCalls) { - throw new Error('sandbox exceeded max tool calls') - } - if (usage.providerCostUsd !== undefined && usage.providerCostUsd > budget.maxProviderCostUsd) { - throw new Error('sandbox exceeded max provider cost') - } -} - -function finalizeUsage( - parts: Partial, - budget: SandboxExecutionBudget, -): SandboxUsageReceipt { - const fields = ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount', 'providerCostUsd', 'budgetEnforced'] as const - if (fields.some((field) => parts[field] === undefined)) { - throw new Error('sandbox did not provide a complete usage receipt') - } - const usage = parts as SandboxUsageReceipt - for (const field of ['inputTokens', 'outputTokens', 'reasoningTokens', 'toolTokens', 'toolCallCount'] as const) { - if (!Number.isSafeInteger(usage[field]) || usage[field] < 0) { - throw new Error(`sandbox usage field ${field} is invalid`) - } - } - if (!Number.isFinite(usage.providerCostUsd) || usage.providerCostUsd < 0) { - throw new Error('sandbox usage provider cost is invalid') - } - if (typeof usage.budgetEnforced !== 'boolean') { - throw new Error('sandbox usage budget flag is invalid') - } - if (!Number.isSafeInteger( - usage.inputTokens + usage.outputTokens + usage.reasoningTokens + usage.toolTokens, - )) { - throw new Error('sandbox usage token total exceeds safe integer range') - } - enforceUsageBudget(usage, budget) - if (!usage.budgetEnforced) throw new Error('sandbox did not enforce the execution budget') - return usage -} - -function completeUsage( - parts: Partial, - reasoningTokens: number, - toolTokens: number, - toolCallCount: number, - userMessage: string, - outputText: string, - budget: SandboxExecutionBudget, - requiresReceipt: boolean, -): SandboxUsageReceipt { - const observed = withObservedUsage(parts, reasoningTokens, toolTokens, toolCallCount) - if ( - !requiresReceipt && - Object.keys(parts).length === 0 && - reasoningTokens === 0 && - toolTokens === 0 && - toolCallCount === 0 - ) { - // Preserve the pre-receipt SandboxBox contract for legacy API-key - // adapters. Durable payment operations must use provider-enforced usage. - return { - inputTokens: estimateTokens(userMessage), - outputTokens: estimateTokens(outputText), - reasoningTokens: 0, - toolTokens: 0, - toolCallCount: 0, - providerCostUsd: 0, - budgetEnforced: false, - } - } - return finalizeUsage(observed, budget) -} - -function actualX402Amount( - pricePerTokenUsd: number, - inputTokens: number, - outputTokens: number, - reasoningTokens: number, - toolTokens: number, - currencyDecimals = 6, - providerCostUsd = 0, -): bigint { - const { numerator, denominator } = decimalFraction(pricePerTokenUsd) - const scaled = BigInt(inputTokens + outputTokens + reasoningTokens + toolTokens) * - numerator * 10n ** BigInt(currencyDecimals) - const tokenAmount = (scaled + denominator - 1n) / denominator - const provider = providerCostUsd === 0 - ? { numerator: 0n, denominator: 1n } - : decimalFraction(providerCostUsd) - const providerScaled = provider.numerator * 10n ** BigInt(currencyDecimals) - const providerAmount = (providerScaled + provider.denominator - 1n) / provider.denominator - return tokenAmount > providerAmount ? tokenAmount : providerAmount -} - -function truncateUtf8( - value: string, - maxBytes: number, - encoder: TextEncoder, -): { text: string; bytes: number; truncated: boolean } { - const bytes = encoder.encode(value).byteLength - if (bytes <= maxBytes) return { text: value, bytes, truncated: false } - let text = '' - let used = 0 - for (const character of value) { - const characterBytes = encoder.encode(character).byteLength - if (used + characterBytes > maxBytes) break - text += character - used += characterBytes - } - return { text, bytes: used, truncated: true } -} - -/** Token estimate matching the existing chat-completions handler (4 chars ≈ 1 token). */ -export function estimateTokens(text: string): number { - return Math.ceil(text.length / 4) -} - -/** Include the host-owned system prompt because the provider bills it too. */ -export function estimateBillableInputTokens(agent: AgentMeta, userMessage: string): number { - return estimateTokens(userMessage) + estimateTokens(agent.systemPrompt ?? '') -} - -/** A tokenizer cannot emit more tokens than the UTF-8 bytes it consumes. */ -export function maximumBillableInputTokens(agent: AgentMeta, userMessage: string): number -export function maximumBillableInputTokens(agent: AgentMeta, messages: readonly ChatMessage[]): number -export function maximumBillableInputTokens(agent: AgentMeta, userMessageOrMessages: string | readonly ChatMessage[]): number { - const encoder = new TextEncoder() - const prompt = typeof userMessageOrMessages === 'string' - ? userMessageOrMessages - : JSON.stringify(userMessageOrMessages) - return encoder.encode(prompt).byteLength + encoder.encode(agent.systemPrompt ?? '').byteLength -} +export type { + A2ADispatchEvent, + AuthorizedRequest, + GatewayState, + PaymentClaimHooks, + SettleAndRecordOptions, +} from './dispatch-types' + +export { + estimateBillableInputTokens, + estimateTokens, + maximumBillableInputTokens, + requiredX402Amount, +} from './dispatch-pricing' + +export { authenticateAndGuard } from './dispatch-authorization' + +export { + beginPaymentExecution, + claimPayment, + markPaymentExecutionStarted, + reclaimPayment, + releasePayment, + releasePaymentAfterFailure, + renewPaymentExecution, +} from './dispatch-payment' + +export { + dispatchSandboxStream, + dispatchSandboxStreamRich, +} from './dispatch-sandbox' + +export { settleAndRecord } from './dispatch-settlement' diff --git a/tests/dispatch-module-size.test.ts b/tests/dispatch-module-size.test.ts new file mode 100644 index 0000000..0d0e3bc --- /dev/null +++ b/tests/dispatch-module-size.test.ts @@ -0,0 +1,23 @@ +import { readFileSync, readdirSync } from 'node:fs' +import { fileURLToPath } from 'node:url' +import { describe, expect, it } from 'vitest' + +const sourceDirectory = fileURLToPath(new URL('../src/', import.meta.url)) + +function lineCount(fileName: string): number { + const source = readFileSync(`${sourceDirectory}/${fileName}`, 'utf8') + return source.trimEnd().split(/\r?\n/).length +} + +describe('dispatch module boundaries', () => { + const files = readdirSync(sourceDirectory).filter( + (fileName) => fileName === 'dispatch.ts' || /^dispatch-.+\.ts$/.test(fileName), + ) + + for (const fileName of files) { + it(`${fileName} stays focused`, () => { + const limit = fileName === 'dispatch.ts' ? 100 : 500 + expect(lineCount(fileName)).toBeLessThanOrEqual(limit) + }) + } +}) From 0c801b7080e6b622c4e433e3eafb39a9cff415d6 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 18:37:59 -0600 Subject: [PATCH 31/32] refactor(a2a): split handler state machines --- src/a2a/execution-fence.ts | 16 +- src/a2a/handler.ts | 2309 ++------------------------- src/a2a/message-send-execution.ts | 241 +++ src/a2a/message-stream-execution.ts | 392 +++++ src/a2a/payment-recovery.ts | 431 +++++ src/a2a/push-config-methods.ts | 158 ++ src/a2a/task-cancellation.ts | 50 + src/a2a/task-finalization.ts | 451 ++++++ src/a2a/task-lifecycle.ts | 54 + src/a2a/task-methods.ts | 163 ++ src/a2a/task-push-delivery.ts | 119 ++ src/a2a/task-recovery.ts | 11 + src/a2a/task-state.ts | 99 ++ src/a2a/task-store.ts | 14 +- src/a2a/task-submission-recovery.ts | 178 +++ src/observer-types.ts | 63 + src/observer.ts | 66 +- src/payment-operations.ts | 6 +- src/payment-recovery.ts | 7 +- src/payment-types.ts | 48 + src/types.ts | 74 +- tests/a2a-state-machines.test.ts | 201 +++ tests/module-size.test.ts | 26 + 23 files changed, 2891 insertions(+), 2286 deletions(-) create mode 100644 src/a2a/message-send-execution.ts create mode 100644 src/a2a/message-stream-execution.ts create mode 100644 src/a2a/payment-recovery.ts create mode 100644 src/a2a/push-config-methods.ts create mode 100644 src/a2a/task-cancellation.ts create mode 100644 src/a2a/task-finalization.ts create mode 100644 src/a2a/task-lifecycle.ts create mode 100644 src/a2a/task-methods.ts create mode 100644 src/a2a/task-push-delivery.ts create mode 100644 src/a2a/task-recovery.ts create mode 100644 src/a2a/task-state.ts create mode 100644 src/a2a/task-submission-recovery.ts create mode 100644 src/observer-types.ts create mode 100644 src/payment-types.ts create mode 100644 tests/a2a-state-machines.test.ts create mode 100644 tests/module-size.test.ts diff --git a/src/a2a/execution-fence.ts b/src/a2a/execution-fence.ts index b310a39..2972e0f 100644 --- a/src/a2a/execution-fence.ts +++ b/src/a2a/execution-fence.ts @@ -1,5 +1,15 @@ import type { Task } from './types' -import type { TaskStore } from './task-store' + +interface ExecutionTaskStore { + get(id: string): Promise + compareAndSet?(expected: Task, next: Task): Promise + compareAndSetExecution?( + expected: Task, + next: Task, + requestId: string, + now: number, + ): Promise +} /** Durable marker that prevents cancellation from racing sandbox start. */ export const TASK_EXECUTION_METADATA_KEY = 'gatewayExecution' @@ -27,7 +37,7 @@ export class TaskExecutionCanceledError extends Error { /** Claim the right to start one task after its sandbox has been acquired. */ export async function claimTaskExecution( - store: TaskStore, + store: ExecutionTaskStore, task: Task, requestId: string, now = Date.now(), @@ -54,7 +64,7 @@ export async function claimTaskExecution( /** Renew both the task execution fence and its cancellation protection. */ export async function renewTaskExecution( - store: TaskStore, + store: ExecutionTaskStore, taskId: string, requestId: string, now?: number, diff --git a/src/a2a/handler.ts b/src/a2a/handler.ts index 9dfddfd..8cad9dc 100644 --- a/src/a2a/handler.ts +++ b/src/a2a/handler.ts @@ -13,68 +13,87 @@ import type { Context } from 'hono' import { - type A2ADispatchEvent, type AuthorizedRequest, type GatewayState, authenticateAndGuard, - beginPaymentExecution, - markPaymentExecutionStarted, - renewPaymentExecution, claimPayment, - dispatchSandboxStreamRich, - releasePayment, - releasePaymentAfterFailure, - settleAndRecord, } from '../dispatch' -import type { PaymentOperation } from '../payment-operations' -import { - deserializePaymentOperation, - serializePaymentOperation, - type SerializedPaymentOperation, -} from '../payment-recovery' -import { recoverPayment as recoverDurablePayment } from '../payment-recovery-worker' import type { ChatMessage, GatewayConfig, - PaymentMethod, - SandboxExecutionBudget, - SandboxUsageReceipt, } from '../types' import { buildAgentCard } from './agent-card' import { - claimTaskExecution, clearTaskExecution, - hasActiveTaskExecution, hasExpiredTaskExecution, - inspectTaskExecution, hasMalformedTaskExecution, - renewTaskExecution, + inspectTaskExecution, } from './execution-fence' +import { executeMessageSend } from './message-send-execution' +import { executeMessageStream } from './message-stream-execution' import { fail, ok, parseEnvelope } from './jsonrpc' import { - deliverDemoPushNotifications, - deliverPushNotifications, - validatePushNotificationUrl, - type PushNotificationDeliveryOptions, - type PushDeliveryResult, type PushNotificationStore, - type TaskPushNotificationConfig, } from './push-notifications' -import { hasPendingPaymentRecovery, type TaskStore } from './task-store' -import { extractTextFromMessage, responseTextToArtifact } from './translate' +import { type TaskStore } from './task-store' +import { extractTextFromMessage } from './translate' import { A2A_ERROR_CODES, - type Artifact, type JSONRPCRequest, type Message, type MessageSendParams, - type StreamingEvent, type Task, - type TaskArtifactUpdateEvent, - type TaskIdParams, - type TaskPushNotificationConfigGetParams, - type TaskStatusUpdateEvent, } from './types' +import { + attachPaymentRecoveryMarker, + hasPaymentReleaseRecovery, + preservePaymentRecoveryMarker, + releaseTaskPayment, + recoverPaymentMarkerIfNeeded, + recoverPaymentReleaseIfNeeded, + retainPaymentRecoveryMarker, +} from './payment-recovery' +import { + isTaskFinalizing, + recoverFinalizationIfNeeded, +} from './task-finalization' +import { + clearTaskSubmission, + readTaskOrigin, + recoverSubmissionIfNeeded, + withTaskOrigin, + withTaskSubmission, +} from './task-submission-recovery' +import { deliverTaskPush, type PushDeliveryDependencies } from './task-push-delivery' +import { + bindRequestAbort, + TaskCancellationRegistry, +} from './task-cancellation' +import { + createTaskLifecycle as buildTaskLifecycle, + type TaskLifecycle, +} from './task-lifecycle' +import { + handleTasksCancel, + handleTasksGet, + handleTasksResubscribe, + type TaskMethodDependencies, +} from './task-methods' +import { + handlePushDelete, + handlePushGet, + handlePushList, + handlePushSet, + type PushConfigMethodDependencies, +} from './push-config-methods' +import { + compareAndSetTask, + cryptoRandomId, + isTerminal, + shouldPreserveTask, + nowIso, + withStatus, +} from './task-state' export interface A2AHandlerDeps { config: GatewayConfig @@ -83,127 +102,44 @@ export interface A2AHandlerDeps { pushStore?: PushNotificationStore } -type FinalizationState = 'completed' | 'input-required' | 'canceled' - -interface FinalizationRecord { - version: 1 - lease: { id: string; expiresAt: number } - agentSlug: string - requestId: string - consumerId: string - paymentMethod: PaymentMethod - startMs: number - operationId: string | null - paymentOperation: SerializedPaymentOperation | null - receipt: SandboxUsageReceipt - artifact: Artifact | null - inputRequired: boolean - inputRequiredPrompt?: string - finalState?: FinalizationState - maxOutputTokens: number - executionBudget: SandboxExecutionBudget - usageRecorded: boolean - recoveryAttempts?: number - recoveryError?: string -} - -interface PaymentReleaseRecord { - version: 1 - lease: { id: string; expiresAt: number } - agentSlug: string - requestId: string - operationId: string - paymentOperation: SerializedPaymentOperation - reason: string - recoveryAttempts?: number - recoveryError?: string -} - -interface TaskPaymentRecoveryMarker { - version: 1 - id: string -} - -interface TaskPushDeliveryClaims { - version: 1 - claims: Record -} - -interface TaskOriginBinding { - version: 1 - agentId: string - agentSlug: string -} - -interface TaskSubmissionRecord { - version: 1 - lease: { id: string; expiresAt: number } - agentId: string - agentSlug: string - requestId: string - consumerId: string -} - -/** Terminal task states — fire-once push delivery occurs on these transitions. */ -const TERMINAL_STATES: ReadonlySet = new Set([ - 'completed', - 'canceled', - 'failed', - 'rejected', -]) - -const TASK_ORIGIN_METADATA_KEY = 'gatewayOrigin' -const TASK_SUBMISSION_METADATA_KEY = 'gatewaySubmission' -const TASK_SUBMISSION_RECOVERY_METADATA_KEY = 'gatewaySubmissionRecovery' -const TASK_SUBMISSION_LEASE_MS = 5 * 60 * 1000 const MAX_A2A_BODY_BYTES = 64 * 1024 class RequestBodyTooLargeError extends Error {} -/** - * Per-gateway in-process registry of cancellable runs. Keyed by task id; - * absent = task already terminal or never streamed. Cleared by the streaming - * handler on completion. Cancel is best-effort: a cancel arriving after the - * stream finished is reported as `TASK_NOT_CANCELABLE`. - */ -class CancelRegistry { - private readonly controllers = new Map() - private readonly finalizing = new Set() - - register(taskId: string): AbortController { - const c = new AbortController() - this.controllers.set(taskId, c) - return c - } - - clear(taskId: string): void { - this.controllers.delete(taskId) - this.finalizing.delete(taskId) - } - - beginFinalization(taskId: string): boolean { - const controller = this.controllers.get(taskId) - if (!controller || controller.signal.aborted || this.finalizing.has(taskId)) return false - this.finalizing.add(taskId) - return true - } - - isFinalizing(taskId: string): boolean { - return this.finalizing.has(taskId) - } +function createTaskLifecycle(deps: A2AHandlerDeps): TaskLifecycle { + return buildTaskLifecycle({ + taskStore: deps.taskStore, + config: deps.config, + state: deps.state, + deliverPush: (task) => maybeDeliverPush(task, deps), + }) +} - has(taskId: string): boolean { - const controller = this.controllers.get(taskId) - return controller !== undefined && !controller.signal.aborted +function buildTaskMethodDependencies( + deps: A2AHandlerDeps, + cancels: TaskCancellationRegistry, +): TaskMethodDependencies { + const lifecycle = createTaskLifecycle(deps) + return { + taskStore: deps.taskStore, + payment: lifecycle.payment, + cancels, + authorizeTaskAccess: (c, req, task) => authorizeTaskAccess(c, req, task, deps), + recoverTask: (task, requestedAgentSlug) => + recoverTaskIfNeeded(task, deps, requestedAgentSlug), + deliverPush: (task) => maybeDeliverPush(task, deps), } +} - cancel(taskId: string): boolean { - if (this.finalizing.has(taskId)) return false - const c = this.controllers.get(taskId) - if (!c) return false - c.abort() - this.controllers.delete(taskId) - return true +function buildPushConfigMethodDependencies( + deps: A2AHandlerDeps, +): PushConfigMethodDependencies { + return { + taskStore: deps.taskStore, + pushStore: deps.pushStore, + demoMode: deps.config.x402.demoMode === true, + urlValidator: deps.config.a2a?.pushUrlValidator, + authorizeTaskAccess: (c, req, task) => authorizeTaskAccess(c, req, task, deps), } } @@ -212,7 +148,7 @@ export function createA2AHandlers(deps: A2AHandlerDeps) { ...deps, taskStore: normalizeTaskStore(deps.taskStore, deps.config.x402.demoMode === true), } - const cancels = new CancelRegistry() + const cancels = new TaskCancellationRegistry() // GET /:slug/.well-known/agent.json const handleAgentCard = async (c: Context): Promise => { @@ -263,19 +199,19 @@ export function createA2AHandlers(deps: A2AHandlerDeps) { case 'message/stream': return handleMessageStream(c, slug, parsed, runtimeDeps, cancels) case 'tasks/get': - return handleTasksGet(c, parsed, runtimeDeps) + return handleTasksGet(c, parsed, buildTaskMethodDependencies(runtimeDeps, cancels)) case 'tasks/cancel': - return handleTasksCancel(c, parsed, runtimeDeps, cancels) + return handleTasksCancel(c, parsed, buildTaskMethodDependencies(runtimeDeps, cancels)) case 'tasks/resubscribe': - return handleTasksResubscribe(c, parsed, runtimeDeps) + return handleTasksResubscribe(c, parsed, buildTaskMethodDependencies(runtimeDeps, cancels)) case 'tasks/pushNotificationConfig/set': - return handlePushSet(c, parsed, runtimeDeps) + return handlePushSet(c, parsed, buildPushConfigMethodDependencies(runtimeDeps)) case 'tasks/pushNotificationConfig/get': - return handlePushGet(c, parsed, runtimeDeps) + return handlePushGet(c, parsed, buildPushConfigMethodDependencies(runtimeDeps)) case 'tasks/pushNotificationConfig/list': - return handlePushList(c, parsed, runtimeDeps) + return handlePushList(c, parsed, buildPushConfigMethodDependencies(runtimeDeps)) case 'tasks/pushNotificationConfig/delete': - return handlePushDelete(c, parsed, runtimeDeps) + return handlePushDelete(c, parsed, buildPushConfigMethodDependencies(runtimeDeps)) default: return c.json( fail(parsed.id, A2A_ERROR_CODES.METHOD_NOT_FOUND, `unknown method '${parsed.method}'`), @@ -293,7 +229,7 @@ async function handleMessageSend( slug: string, req: JSONRPCRequest, deps: A2AHandlerDeps, - cancels: CancelRegistry, + cancels: TaskCancellationRegistry, ): Promise { const guard = await guardMessageRequest(c, slug, req, deps) if (guard instanceof Response) return guard @@ -301,223 +237,22 @@ async function handleMessageSend( setPaymentResponseHeaders(c, authz) const controller = cancels.register(task.id) const detachRequestAbort = bindRequestAbort(c.req.raw.signal, controller) + const lifecycle = createTaskLifecycle(deps) try { - return await executeMessageSend(c, req, deps, authz, task, controller.signal) + return await executeMessageSend( + c, + req, + { taskStore: deps.taskStore, config: deps.config, ...lifecycle }, + authz, + task, + controller.signal, + ) } finally { detachRequestAbort() cancels.clear(task.id) } } -async function executeMessageSend( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, - authz: AuthorizedRequest, - task: Task, - signal: AbortSignal, -): Promise { - if (isTerminal(task.status.state)) return c.json(ok(req.id, task)) - const taskWithoutSubmission = clearTaskSubmission(task) - let workingTask: Task = task.status.state === 'working' - ? taskWithoutSubmission - : { ...taskWithoutSubmission, status: { state: 'working', timestamp: nowIso() } } - if ( - JSON.stringify(task) !== JSON.stringify(workingTask) && - !await compareAndSetTask(deps.taskStore, task, workingTask) - ) { - await releaseTaskPayment(authz, task, deps, 'A2A task changed before execution started', false) - if (signal.aborted) { - const canceled = await deps.taskStore.get(task.id) - if (canceled?.status.state === 'canceled') { - await maybeDeliverPush(canceled, deps) - return c.json(ok(req.id, canceled)) - } - } - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) - } - - let responseText = '' - let usage: SandboxUsageReceipt | undefined - let workObserved = false - let inputRequiredPrompt: string | undefined - let inputRequiredSeen = false - let finalizationLeaseId: string | undefined - try { - for await (const event of dispatchSandboxStreamRich( - authz.agent, - authz.userMessage, - authz.consumerId, - deps.config, - signal, - task.id, - authz.maxOutputTokens, - async () => { - workingTask = await claimTaskExecution(deps.taskStore, workingTask, authz.requestId) - await beginPaymentExecution(authz, deps.config) - }, - authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, - async () => { - workObserved = true - await markPaymentExecutionStarted(authz, deps.config) - }, - authz.executionBudget.maxInputTokens, - async () => { - workingTask = await renewTaskExecution(deps.taskStore, task.id, authz.requestId) - await renewPaymentExecution(authz, deps.config) - }, - )) { - if (event.kind === 'text') { - responseText += event.delta - workObserved = true - } else if (event.kind === 'activity') { - workObserved = true - } else if (event.kind === 'usage') { - usage = event.usage - } else { - inputRequiredSeen = true - inputRequiredPrompt = event.prompt - workObserved = true - } - } - } catch (err) { - const releasedTask = await releaseTaskPayment( - authz, - workingTask, - deps, - err instanceof Error ? err.message : String(err), - workObserved || usage !== undefined, - ) - const currentTask = await deps.taskStore.get(task.id) ?? releasedTask - const failed = shouldPreserveTask(currentTask) - ? currentTask - : withStatus(clearTaskSubmission(currentTask), 'failed') - try { - const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) - await maybeDeliverPush(persisted, deps) - } catch (taskError) { - console.error( - `[a2a] failed to persist failed task ${task.id}:`, - taskError instanceof Error ? taskError.message : String(taskError), - ) - } - return c.json( - fail( - req.id, - A2A_ERROR_CODES.INTERNAL_ERROR, - err instanceof Error ? err.message : String(err), - ), - ) - } - - if (signal.aborted) { - const canceled = await completeCanceledTask( - authz, - workingTask, - responseText, - usage, - workObserved, - deps, - ) - return c.json(ok(req.id, canceled)) - } - - // Settle for the work done so far before short-circuiting on input-required. - // The user has been charged for the partial response, which is the right - // commercial behavior — the sandbox produced tokens. - try { - if (!usage) throw new Error('sandbox did not provide a usage receipt') - const finalizationArtifact = responseText - ? responseTextToArtifact(responseText, `${task.id}-artifact-0`) - : task.artifacts?.[0] ?? null - const finalization = buildFinalizationRecord( - authz, - usage, - finalizationArtifact, - inputRequiredSeen, - inputRequiredPrompt, - ) - const finalizingTask = withFinalizationRecord(workingTask, finalization) - if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask?.status.state === 'canceled') { - const canceled = await completeCanceledTask( - authz, - currentTask, - responseText, - usage, - workObserved, - deps, - ) - return c.json(ok(req.id, canceled)) - } - throw new Error('A2A task changed before payment settlement') - } - finalizationLeaseId = finalization.lease.id - let usageRecordedTask = finalizingTask - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { - onUsageRecorded: async () => { - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - }, - }) - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - const settledBase = clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)) - const result = inputRequiredSeen - ? withStatus( - settledBase, - 'input-required', - inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, - responseText - ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] - : task.artifacts, - ) - : withStatus(settledBase, 'completed', undefined, [ - responseTextToArtifact(responseText, `${task.id}-artifact-0`), - ]) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, result)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask && (isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required')) { - return c.json(ok(req.id, currentTask)) - } - throw new Error('A2A task changed after payment settlement') - } - if (inputRequiredSeen) return c.json(ok(req.id, result)) - await maybeDeliverPush(result, deps) - return c.json(ok(req.id, result)) - } catch (err) { - const releasedTask = await releaseTaskPayment( - authz, - workingTask, - deps, - err instanceof Error ? err.message : String(err), - workObserved || usage !== undefined, - ) - if (finalizationLeaseId) { - await retainFinalizationForRecovery( - deps.taskStore, - task.id, - finalizationLeaseId, - asError(err), - ) - return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) - } - const currentTask = await deps.taskStore.get(task.id) ?? releasedTask - const failed = shouldPreserveTask(currentTask) - ? currentTask - : withStatus(clearTaskSubmission(currentTask), 'failed') - try { - const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) - await maybeDeliverPush(persisted, deps) - } catch (taskError) { - console.error( - `[a2a] failed to persist failed task ${task.id}:`, - taskError instanceof Error ? taskError.message : String(taskError), - ) - } - return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) - } -} - // ── message/stream (SSE) ────────────────────────────────────────────────── async function handleMessageStream( @@ -525,663 +260,38 @@ async function handleMessageStream( slug: string, req: JSONRPCRequest, deps: A2AHandlerDeps, - cancels: CancelRegistry, + cancels: TaskCancellationRegistry, ): Promise { const guard = await guardMessageRequest(c, slug, req, deps) if (guard instanceof Response) return guard const { authz, task } = guard setPaymentResponseHeaders(c, authz) - - const controller = cancels.register(task.id) - const detachRequestAbort = bindRequestAbort(c.req.raw.signal, controller) - const workingStatus: TaskStatusUpdateEvent = { - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: { state: 'working', timestamp: nowIso() }, - final: false, - } - if (isTerminal(task.status.state)) { - detachRequestAbort() - cancels.clear(task.id) - return c.json(ok(req.id, task)) - } - let workingTask: Task = task.status.state === 'working' - ? task - : { ...task, status: workingStatus.status } - if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { - detachRequestAbort() - cancels.clear(task.id) - const current = await releaseTaskPayment( - authz, - await deps.taskStore.get(task.id) ?? task, - deps, - 'A2A task changed before execution started', - false, - ) - if (current.status.state === 'canceled') return c.json(ok(req.id, current)) - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) - } - let responseText = '' - let usage: SandboxUsageReceipt | undefined - let workObserved = false - - const stream = new ReadableStream({ - start(ctrl) { - void (async () => { - const encoder = new TextEncoder() - const send = (event: StreamingEvent) => { - if (ctrl.desiredSize === null) return - try { - ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) - } catch { - // The client can cancel between the desiredSize check and enqueue. - } - } - - let inputRequiredPrompt: string | undefined - let inputRequiredSeen = false - let finalizationLeaseId: string | undefined - try { - send(workingStatus) - - for await (const event of dispatchSandboxStreamRich( - authz.agent, - authz.userMessage, - authz.consumerId, - deps.config, - controller.signal, - task.id, - authz.maxOutputTokens, - async () => { - workingTask = await claimTaskExecution(deps.taskStore, workingTask, authz.requestId) - await beginPaymentExecution(authz, deps.config) - }, - authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, - async () => { - workObserved = true - await markPaymentExecutionStarted(authz, deps.config) - }, - authz.executionBudget.maxInputTokens, - async () => { - workingTask = await renewTaskExecution(deps.taskStore, task.id, authz.requestId) - await renewPaymentExecution(authz, deps.config) - }, - )) { - if (event.kind === 'text') { - responseText += event.delta - workObserved = true - const artifactEvent: TaskArtifactUpdateEvent = { - kind: 'artifact-update', - taskId: task.id, - contextId: task.contextId, - artifact: { - artifactId: `${task.id}-artifact-0`, - name: 'response', - parts: [{ kind: 'text', text: event.delta }], - }, - append: true, - } - send(artifactEvent) - } else if (event.kind === 'activity') { - workObserved = true - } else if (event.kind === 'usage') { - usage = event.usage - } else { - inputRequiredSeen = true - inputRequiredPrompt = event.prompt - workObserved = true - } - } - - // Caller aborted via tasks/cancel. Charge a complete receipt if one - // exists; otherwise retain ownership when output or hidden work was - // observed because releasing would make paid work free. - if (controller.signal.aborted) { - const canceled = await completeCanceledTask( - authz, - workingTask, - responseText, - usage, - workObserved, - deps, - ) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: canceled.status, - final: true, - }) - return - } - - // Settle once for whatever the sandbox produced (full or partial). - if (!usage) throw new Error('sandbox did not provide a usage receipt') - const finalizationArtifact = responseTextToArtifact(responseText, `${task.id}-artifact-0`) - const finalization = buildFinalizationRecord( - authz, - usage, - finalizationArtifact, - inputRequiredSeen, - inputRequiredPrompt, - ) - // Let the durable task-store CAS decide the cancellation race. Mark - // the local registry only after that CAS wins, so cancel can replace a - // still-pending finalization instead of being rejected prematurely. - const finalizingTask = withFinalizationRecord(workingTask, finalization) - if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask?.status.state === 'canceled') { - const canceled = await completeCanceledTask( - authz, - currentTask, - responseText, - usage, - workObserved, - deps, - ) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: canceled.status, - final: true, - }) - return - } - await releaseTaskPayment( - authz, - task, - deps, - 'A2A task changed before payment settlement', - workObserved || usage !== undefined, - ) - return - } - finalizationLeaseId = finalization.lease.id - cancels.beginFinalization(task.id) - let usageRecordedTask = finalizingTask - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { - onUsageRecorded: async () => { - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - }, - }) - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - - if (inputRequiredSeen) { - const paused = withStatus( - clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), - 'input-required', - inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, - responseText - ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] - : task.artifacts, - ) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, paused)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask) { - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: currentTask.status, - final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', - }) - } - return - } - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: paused.status, - final: true, - }) - // input-required is non-terminal — do NOT deliver push notifications. - return - } - - // Final: persist the terminal task before emitting terminal events. - const completed = withStatus(clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), 'completed', undefined, [ - responseTextToArtifact(responseText, `${task.id}-artifact-0`), - ]) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, completed)) { - const currentTask = await deps.taskStore.get(task.id) - if (currentTask) { - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: currentTask.status, - final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', - }) - } - return - } - send({ - kind: 'artifact-update', - taskId: task.id, - contextId: task.contextId, - artifact: { - artifactId: `${task.id}-artifact-0`, - name: 'response', - parts: [{ kind: 'text', text: '' }], - }, - append: true, - lastChunk: true, - }) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: completed.status, - final: true, - }) - await maybeDeliverPush(completed, deps) - } catch (err) { - const releasedTask = await releaseTaskPayment( - authz, - task, - deps, - err instanceof Error ? err.message : String(err), - workObserved || usage !== undefined, - ) - if (finalizationLeaseId) { - const retained = await retainFinalizationForRecovery( - deps.taskStore, - task.id, - finalizationLeaseId, - asError(err), - ) - if (retained) { - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: retained.status, - final: false, - }) - } - return - } - const currentTask = await deps.taskStore.get(task.id) ?? releasedTask - const failed = shouldPreserveTask(currentTask) - ? currentTask - : withStatus(clearTaskSubmission(currentTask), 'failed') - try { - const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) - send({ - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: persisted.status, - final: true, - }) - await maybeDeliverPush(persisted, deps) - } catch (taskError) { - console.error( - `[a2a] failed to persist failed task ${task.id}:`, - taskError instanceof Error ? taskError.message : String(taskError), - ) - } - try { - await deps.state.obs?.onStreamError?.( - { - requestId: authz.requestId, - agentSlug: authz.agent.slug, - startMs: authz.startMs, - }, - { - consumerId: authz.consumerId, - errorMessage: err instanceof Error ? err.message : String(err), - }, - ) - } catch (observerError) { - console.error( - `[a2a] stream observer failed for ${authz.requestId}:`, - observerError instanceof Error ? observerError.message : String(observerError), - ) - } - } finally { - detachRequestAbort() - cancels.clear(task.id) - try { - if (ctrl.desiredSize !== null) ctrl.close() - } catch { - // The response may already be closed by client cancellation. - } - } - })() - }, - cancel() { - controller.abort() - }, - }) - - return new Response(stream, { - headers: { - 'Content-Type': 'text/event-stream', - 'Cache-Control': 'no-cache', - 'X-Request-Id': authz.requestId, - 'X-Agent-Slug': authz.agent.slug, - 'X-Task-Id': task.id, - ...(authz.mppChargeOperation - ? { 'Payment-Receipt': authz.mppChargeOperation.receipt } - : {}), - ...(authz.paymentRecoveryId - ? { 'X-Payment-Operation-Id': authz.paymentRecoveryId } - : {}), + return executeMessageStream( + c, + req, + { + taskStore: deps.taskStore, + config: deps.config, + lifecycle: createTaskLifecycle(deps), + cancels, + deliverPush: (streamTask) => maybeDeliverPush(streamTask, deps), + reportStreamError: async (streamAuthz, error) => { + await deps.state.obs?.onStreamError?.( + { + requestId: streamAuthz.requestId, + agentSlug: streamAuthz.agent.slug, + startMs: streamAuthz.startMs, + }, + { + consumerId: streamAuthz.consumerId, + errorMessage: error instanceof Error ? error.message : String(error), + }, + ) + }, }, - }) -} - -// ── tasks/get + tasks/cancel ────────────────────────────────────────────── - -async function handleTasksGet( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, -): Promise { - const params = req.params as TaskIdParams | undefined - if (!params || typeof params.id !== 'string') { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) - } - const storedTask = await deps.taskStore.get(params.id) - if (!storedTask) { - return c.json( - fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), - ) - } - const accessError = await authorizeTaskAccess(c, req, storedTask, deps) - if (accessError) return accessError - const task = await recoverTaskIfNeeded( - storedTask, - deps, - c.req.param('slug') ?? '', - ) - return c.json(ok(req.id, task)) -} - -async function handleTasksCancel( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, - cancels: CancelRegistry, -): Promise { - const params = req.params as TaskIdParams | undefined - if (!params || typeof params.id !== 'string') { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) - } - const storedTask = await deps.taskStore.get(params.id) - if (!storedTask) { - return c.json( - fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), - ) - } - const accessError = await authorizeTaskAccess(c, req, storedTask, deps) - if (accessError) return accessError - const task = await recoverTaskIfNeeded( - storedTask, - deps, - c.req.param('slug') ?? '', - ) - if (isTerminal(task.status.state)) { - return c.json( - fail( - req.id, - A2A_ERROR_CODES.TASK_NOT_CANCELABLE, - `task '${params.id}' is in terminal state '${task.status.state}'`, - ), - ) - } - - if ( - isTaskFinalizing(task) || - cancels.isFinalizing(task.id) || - (hasActiveTaskExecution(task) && !cancels.has(task.id)) - ) { - return c.json( - fail( - req.id, - A2A_ERROR_CODES.TASK_NOT_CANCELABLE, - hasActiveTaskExecution(task) - ? `task '${task.id}' has an active execution fence` - : `task '${task.id}' is being finalized`, - ), - ) - } - let stillActive = false - let candidate = task - for (let attempt = 0; attempt < 8; attempt += 1) { - if (isTerminal(candidate.status.state)) { - return c.json( - fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' changed before cancellation`), - ) - } - if (isTaskFinalizing(candidate)) { - return c.json( - fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' is being finalized`), - ) - } - if (hasActiveTaskExecution(candidate) && !cancels.has(candidate.id)) { - return c.json( - fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' has an active execution fence`), - ) - } - const canceled = withStatus(candidate, 'canceled') - if (await compareAndSetTask(deps.taskStore, candidate, canceled)) { - stillActive = cancels.cancel(task.id) - // If a stream was active, it observes the abort and emits its own final - // status update and push delivery. Otherwise this handler owns delivery. - if (!stillActive) await maybeDeliverPush(canceled, deps) - return c.json(ok(req.id, canceled)) - } - const current = await deps.taskStore.get(task.id) - if (!current) { - return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${task.id}' not found`)) - } - candidate = current - } - return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'task changed before cancellation')) -} - -// ── tasks/resubscribe ───────────────────────────────────────────────────── - -/** - * Re-attach to a known task via SSE. The minimum-viable shape (and the one - * the spec actually requires): emit the task's current status as one - * status-update event with the right `final` flag, then close. Callers that - * lost their original stream connection can re-subscribe to find out where - * the task ended up; in-flight tasks return their last-known state and the - * client polls (or re-subscribes) for further updates. - * - * Out of scope: live-rebroadcasting deltas from an in-flight stream to a new - * subscriber. That requires per-task pub/sub which we haven't needed yet — - * the typical recovery path is "task already finished, fetch the result." - */ -async function handleTasksResubscribe( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, -): Promise { - const params = req.params as TaskIdParams | undefined - if (!params || typeof params.id !== 'string') { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) - } - const storedTask = await deps.taskStore.get(params.id) - if (!storedTask) { - return c.json( - fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`), - ) - } - const accessError = await authorizeTaskAccess(c, req, storedTask, deps) - if (accessError) return accessError - const task = await recoverTaskIfNeeded( - storedTask, - deps, - c.req.param('slug') ?? '', + authz, + task, ) - const final = isTerminal(task.status.state) || task.status.state === 'input-required' - const event: TaskStatusUpdateEvent = { - kind: 'status-update', - taskId: task.id, - contextId: task.contextId, - status: task.status, - final, - } - const encoder = new TextEncoder() - const stream = new ReadableStream({ - start(ctrl) { - ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) - ctrl.close() - }, - }) - return new Response(stream, { - headers: { - 'Content-Type': 'text/event-stream', - 'Cache-Control': 'no-cache', - 'X-Task-Id': task.id, - }, - }) -} - -// ── tasks/pushNotificationConfig/* ──────────────────────────────────────── - -async function handlePushSet( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, -): Promise { - if (!deps.pushStore) { - return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) - } - const params = req.params as TaskPushNotificationConfig | undefined - if (!params || typeof params.taskId !== 'string' || !params.pushNotificationConfig?.id) { - return c.json( - fail( - req.id, - A2A_ERROR_CODES.INVALID_PARAMS, - 'params.taskId and params.pushNotificationConfig.id required', - ), - ) - } - if (typeof params.pushNotificationConfig.url !== 'string') { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url required')) - } - const task = await deps.taskStore.get(params.taskId) - if (!task) { - return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.taskId}' not found`)) - } - const accessError = await authorizeTaskAccess(c, req, task, deps) - if (accessError) return accessError - const pushUrl = validatePushNotificationUrl(params.pushNotificationConfig.url) - if (!pushUrl) { - return c.json( - fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url is not a safe HTTPS destination'), - ) - } - const urlValidator = deps.config.a2a?.pushUrlValidator - if (!deps.config.x402.demoMode && !urlValidator) { - return c.json( - fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'production push URL validation is not configured'), - ) - } - let allowedByHostPolicy = true - try { - if (urlValidator) allowedByHostPolicy = await urlValidator(pushUrl) - } catch (error) { - allowedByHostPolicy = false - console.error( - `[a2a] push URL policy failed for task ${task.id}:`, - error instanceof Error ? error.message : String(error), - ) - } - if (!allowedByHostPolicy) { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url was rejected')) - } - await deps.pushStore.set(params.taskId, params.pushNotificationConfig) - const stored = await deps.pushStore.get(params.taskId, params.pushNotificationConfig.id) - return c.json(ok(req.id, { taskId: params.taskId, pushNotificationConfig: stored })) -} - -async function handlePushGet( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, -): Promise { - if (!deps.pushStore) { - return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) - } - const params = req.params as TaskPushNotificationConfigGetParams | undefined - if (!params || typeof params.id !== 'string' || typeof params.pushNotificationConfigId !== 'string') { - return c.json( - fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id and params.pushNotificationConfigId required'), - ) - } - const task = await deps.taskStore.get(params.id) - if (!task) { - return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) - } - const accessError = await authorizeTaskAccess(c, req, task, deps) - if (accessError) return accessError - const cfg = await deps.pushStore.get(params.id, params.pushNotificationConfigId) - if (!cfg) { - return c.json( - fail( - req.id, - A2A_ERROR_CODES.TASK_NOT_FOUND, - `push config '${params.pushNotificationConfigId}' not found for task '${params.id}'`, - ), - ) - } - return c.json(ok(req.id, { taskId: params.id, pushNotificationConfig: cfg })) -} - -async function handlePushList( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, -): Promise { - if (!deps.pushStore) { - return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) - } - const params = req.params as TaskIdParams | undefined - if (!params || typeof params.id !== 'string') { - return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) - } - const task = await deps.taskStore.get(params.id) - if (!task) { - return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) - } - const accessError = await authorizeTaskAccess(c, req, task, deps) - if (accessError) return accessError - const configs = await deps.pushStore.list(params.id) - return c.json(ok(req.id, configs.map((cfg) => ({ taskId: params.id, pushNotificationConfig: cfg })))) -} - -async function handlePushDelete( - c: Context, - req: JSONRPCRequest, - deps: A2AHandlerDeps, -): Promise { - if (!deps.pushStore) { - return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) - } - const params = req.params as TaskPushNotificationConfigGetParams | undefined - if (!params || typeof params.id !== 'string' || typeof params.pushNotificationConfigId !== 'string') { - return c.json( - fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id and params.pushNotificationConfigId required'), - ) - } - const task = await deps.taskStore.get(params.id) - if (!task) { - return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) - } - const accessError = await authorizeTaskAccess(c, req, task, deps) - if (accessError) return accessError - await deps.pushStore.delete(params.id, params.pushNotificationConfigId) - return c.json(ok(req.id, null)) } // ── Shared message-send setup (auth + task allocation) ──────────────────── @@ -1313,6 +423,7 @@ async function claimTaskPayment( deps: A2AHandlerDeps, paymentFailureTask?: Task, ): Promise { + const lifecycle = createTaskLifecycle(deps) let paymentTask = task try { await claimPayment(authz, deps.config, deps.state, { @@ -1341,13 +452,13 @@ async function claimTaskPayment( const releasedTask = await releaseTaskPayment( authz, recoveryTask, - deps, + lifecycle.payment, 'payment authorization failed', false, ) const cleanedReleasedTask = clearTaskSubmission(releasedTask) - const releaseRecord = cleanedReleasedTask.metadata?.[PAYMENT_RELEASE_METADATA_KEY] - const failed = releaseRecord !== undefined + const hasReleaseRecord = hasPaymentReleaseRecovery(cleanedReleasedTask) + const failed = hasReleaseRecord ? isTerminal(cleanedReleasedTask.status.state) ? cleanedReleasedTask : withStatus(cleanedReleasedTask, 'failed') @@ -1391,7 +502,7 @@ async function claimTaskPayment( const released = await releaseTaskPayment( authz, recoveryTask, - deps, + lifecycle.payment, 'A2A task changed during payment confirmation', false, ) @@ -1402,195 +513,9 @@ async function claimTaskPayment( } } -async function releaseTaskPayment( - authz: AuthorizedRequest, - task: Task, - deps: A2AHandlerDeps, - reason: string, - workObserved: boolean, -): Promise { - // Store the operation before release because the adapter acknowledgement can be ambiguous. - if ( - !workObserved && - !authz.paymentRecoveryId && - authz.paymentOperation && - deps.config.x402.paymentOperations - ) { - let marked: Task - try { - marked = await beginPaymentReleaseRecovery(deps.taskStore, task, authz, reason) ?? task - } catch (error) { - console.error( - '[a2a] failed to persist payment release recovery for ' + authz.requestId + ':', - error instanceof Error ? error.message : String(error), - ) - return await deps.taskStore.get(task.id) ?? task - } - const record = readPaymentReleaseRecord(marked) - if (!record) return marked - try { - await releasePayment(authz, deps.config, reason) - } catch (releaseError) { - const retained = await retainPaymentReleaseForRecovery( - deps.taskStore, - task.id, - record.lease.id, - releaseError instanceof Error ? releaseError : new Error(String(releaseError)), - ) - console.error( - '[a2a] payment release retained for ' + authz.requestId + ':', - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) - return retained ?? marked - } - return clearPaymentReleaseRecovery(deps.taskStore, marked, record.lease.id) - } - - try { - await releasePaymentAfterFailure(authz, deps.config, reason, workObserved) - } catch (releaseError) { - console.error( - `[a2a] payment release failed for ${authz.requestId}:`, - releaseError instanceof Error ? releaseError.message : String(releaseError), - ) - } - const current = await deps.taskStore.get(task.id) ?? task - return workObserved - ? current - : clearReconciledPaymentRecoveryMarker(current, deps) -} - -/** Keep an owned finalization record durable when settlement acknowledgement is lost. */ -async function retainFinalizationForRecovery( - taskStore: TaskStore, - taskId: string, - leaseId: string, - error: Error, -): Promise { - const current = await taskStore.get(taskId) - if (!current) return undefined - const record = readFinalizationRecord(current) - if (!record || record.lease.id !== leaseId) return undefined - const retry: FinalizationRecord = { - ...record, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, - recoveryAttempts: (record.recoveryAttempts ?? 0) + 1, - recoveryError: error.message, - } - const next = withFinalizationRecord(current, retry) - if (await compareAndSetTask(taskStore, current, next)) return next - return await taskStore.get(taskId) -} - -async function completeCanceledTask( - authz: AuthorizedRequest, - task: Task, - responseText: string, - usage: SandboxUsageReceipt | undefined, - workObserved: boolean, - deps: A2AHandlerDeps, -): Promise { - if (usage) { - const current = await deps.taskStore.get(task.id) ?? task - const finalization = buildFinalizationRecord( - authz, - usage, - responseText - ? responseTextToArtifact(responseText, `${task.id}-artifact-0`) - : current.artifacts?.[0] ?? null, - false, - undefined, - 'canceled', - ) - let finalizingTask: Task | undefined - let candidate = current - for (let attempt = 0; attempt < 8; attempt += 1) { - if (isTaskFinalizing(candidate)) return candidate - if (isTerminal(candidate.status.state) && candidate.status.state !== 'canceled') return candidate - // Store the lease before settlement can move the payment to settling. - const next = withFinalizationRecord(candidate, finalization) - if (await compareAndSetTask(deps.taskStore, candidate, next)) { - finalizingTask = next - break - } - const latest = await deps.taskStore.get(task.id) - if (!latest) break - if (isTerminal(latest.status.state) && latest.status.state !== 'canceled') return latest - candidate = latest - } - if (!finalizingTask) { - throw new Error(`A2A task '${task.id}' changed before cancellation settlement`) - } - - let usageRecordedTask = finalizingTask - try { - await settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, { - onUsageRecorded: async () => { - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - }, - }) - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - } catch (settlementError) { - await releasePaymentAfterFailure( - authz, - deps.config, - settlementError instanceof Error ? settlementError.message : String(settlementError), - true, - ) - const retained = await retainFinalizationForRecovery( - deps.taskStore, - task.id, - finalization.lease.id, - asError(settlementError), - ) - const recoveryTask = retained ?? finalizingTask - console.error( - `[a2a] canceled task settlement retained for ${authz.requestId}:`, - settlementError instanceof Error ? settlementError.message : String(settlementError), - ) - await maybeDeliverPush(recoveryTask, deps) - return recoveryTask - } - - const canceled = withStatus( - clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), - 'canceled', - undefined, - responseText - ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] - : finalizingTask.artifacts, - ) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, canceled)) { - return await deps.taskStore.get(task.id) ?? canceled - } - await maybeDeliverPush(canceled, deps) - return canceled - } - await releaseTaskPayment(authz, task, deps, 'a2a task canceled', workObserved) - const currentTask = await deps.taskStore.get(task.id) - const canceledBase = currentTask?.status.state === 'canceled' - ? currentTask - : withStatus(currentTask ?? task, 'canceled') - const canceled: Task = responseText - ? { - ...canceledBase, - artifacts: [responseTextToArtifact(responseText, `${task.id}-artifact-0`)], - } - : canceledBase - const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask ?? task, canceled) - await maybeDeliverPush(persisted, deps) - return persisted -} - // ── Helpers ─────────────────────────────────────────────────────────────── -const FINALIZING_METADATA_KEY = 'gatewayFinalizing' -const PAYMENT_RELEASE_METADATA_KEY = 'gatewayPaymentRelease' -const PAYMENT_RECOVERY_METADATA_KEY = 'gatewayPaymentRecovery' const EXECUTION_RECOVERY_METADATA_KEY = 'gatewayExecutionRecovery' -const PUSH_DELIVERY_METADATA_KEY = 'gatewayPushDelivery' -const FINALIZATION_LEASE_MS = 5 * 60 * 1000 -const PAYMENT_RELEASE_LEASE_MS = 5 * 60 * 1000 function setPaymentResponseHeaders(c: Context, authz: AuthorizedRequest): void { if (authz.mppChargeOperation) { @@ -1646,13 +571,6 @@ async function authorizeTaskAccess( return c.json(fail(req.id, A2A_ERROR_CODES.TASK_ACCESS_DENIED, 'task access denied'), 403) } -function bindRequestAbort(requestSignal: AbortSignal, controller: AbortController): () => void { - const abort = () => controller.abort() - if (requestSignal.aborted) abort() - else requestSignal.addEventListener('abort', abort, { once: true }) - return () => requestSignal.removeEventListener('abort', abort) -} - function taskHistoryAsChatMessages(task: Task): ChatMessage[] { return (task.history ?? []).flatMap((message) => { const extracted = extractTextFromMessage(message) @@ -1664,221 +582,12 @@ function taskHistoryAsChatMessages(task: Task): ChatMessage[] { }) } -function withTaskOrigin( - metadata: Record | undefined, - agent: { id: string; slug: string }, -): Record { - return { - ...(metadata ?? {}), - [TASK_ORIGIN_METADATA_KEY]: { - version: 1, - agentId: agent.id, - agentSlug: agent.slug, - } satisfies TaskOriginBinding, - } -} - -function withTaskSubmission( - metadata: Record | undefined, - authz: AuthorizedRequest, -): Record { - const origin = metadata?.[TASK_ORIGIN_METADATA_KEY] - return { - ...(metadata ?? {}), - ...(origin === undefined - ? { - [TASK_ORIGIN_METADATA_KEY]: { - version: 1, - agentId: authz.agent.id, - agentSlug: authz.agent.slug, - } satisfies TaskOriginBinding, - } - : {}), - [TASK_SUBMISSION_METADATA_KEY]: { - version: 1, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + TASK_SUBMISSION_LEASE_MS }, - agentId: authz.agent.id, - agentSlug: authz.agent.slug, - requestId: authz.requestId, - consumerId: authz.consumerId, - } satisfies TaskSubmissionRecord, - } -} - -function readTaskOrigin(task: Task): TaskOriginBinding | undefined { - const raw = task.metadata?.[TASK_ORIGIN_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined - const origin = raw as Partial - if ( - origin.version !== 1 || - typeof origin.agentId !== 'string' || - origin.agentId.length === 0 || - typeof origin.agentSlug !== 'string' || - origin.agentSlug.length === 0 - ) { - return undefined - } - return origin as TaskOriginBinding -} - -function readTaskSubmission(task: Task): TaskSubmissionRecord | undefined { - const raw = task.metadata?.[TASK_SUBMISSION_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined - const submission = raw as Partial - if ( - submission.version !== 1 || - !submission.lease || - typeof submission.lease.id !== 'string' || - submission.lease.id.length === 0 || - typeof submission.lease.expiresAt !== 'number' || - !Number.isFinite(submission.lease.expiresAt) || - typeof submission.agentId !== 'string' || - submission.agentId.length === 0 || - typeof submission.agentSlug !== 'string' || - submission.agentSlug.length === 0 || - typeof submission.requestId !== 'string' || - submission.requestId.length === 0 || - typeof submission.consumerId !== 'string' - ) { - return undefined - } - return submission as TaskSubmissionRecord -} - -function clearTaskSubmission(task: Task): Task { - if (!task.metadata || !(TASK_SUBMISSION_METADATA_KEY in task.metadata)) return task - const metadata = { ...task.metadata } - delete metadata[TASK_SUBMISSION_METADATA_KEY] - return Object.keys(metadata).length > 0 - ? { ...task, metadata } - : (() => { - const { metadata: _metadata, ...withoutMetadata } = task - return withoutMetadata - })() -} - -async function clearTaskSubmissionMarker( - taskStore: TaskStore, - expected: Task, -): Promise<{ task: Task; applied: boolean }> { - const current = await taskStore.get(expected.id) - if (!current || JSON.stringify(current) !== JSON.stringify(expected)) { - return { task: current ?? expected, applied: false } - } - const cleared = clearTaskSubmission(current) - if (cleared === current) return { task: current, applied: true } - if (await compareAndSetTask(taskStore, current, cleared)) return { task: cleared, applied: true } - return { task: await taskStore.get(expected.id) ?? expected, applied: false } -} - -function shouldPreserveTask(task: Task): boolean { - return isTerminal(task.status.state) || hasPendingPaymentRecovery(task) -} - -async function persistTaskIfCurrent( - taskStore: TaskStore, - expected: Task, - next: Task, -): Promise { - if (expected === next || JSON.stringify(expected) === JSON.stringify(next)) return expected - if (await compareAndSetTask(taskStore, expected, next)) return next - return await taskStore.get(expected.id) ?? expected -} - async function createTask(taskStore: TaskStore, task: Task): Promise { if (!taskStore.createIfAbsent) { throw new Error('A2A task store does not provide createIfAbsent') } return taskStore.createIfAbsent(task) } - -async function compareAndSetTask(taskStore: TaskStore, expected: Task, next: Task): Promise { - if (!taskStore.compareAndSet) { - throw new Error('A2A task store does not provide compareAndSet') - } - return taskStore.compareAndSet(expected, next) -} - -async function attachPaymentRecoveryMarker( - taskStore: TaskStore, - task: Task, - recoveryId: string | undefined, -): Promise { - if (!recoveryId) return task - const existing = readPaymentRecoveryMarker(task) - if (existing) { - if (existing.id !== recoveryId) { - throw new Error('A2A task already has a different payment recovery identity') - } - return task - } - const next = withPaymentRecoveryMarker(task, recoveryId) - if (await compareAndSetTask(taskStore, task, next)) return next - throw new Error('A2A task changed while payment recovery was attached') -} - -/** Attach only as a retention marker. The returned task must never execute. */ -async function retainPaymentRecoveryMarker( - taskStore: TaskStore, - task: Task, - recoveryId: string | undefined, -): Promise { - if (!recoveryId) return task - let current = await taskStore.get(task.id) ?? task - for (let attempt = 0; attempt < 16; attempt += 1) { - const existing = readPaymentRecoveryMarker(current) - if (existing) { - if (existing.id !== recoveryId) { - throw new Error('A2A task already has a different payment recovery identity') - } - return current - } - const next = withPaymentRecoveryMarker(current, recoveryId) - if (await compareAndSetTask(taskStore, current, next)) return next - const latest = await taskStore.get(task.id) - if (!latest) throw new Error('A2A task disappeared while payment recovery was retained') - current = latest - } - throw new Error('A2A task changed too many times while payment recovery was retained') -} - -function withPaymentRecoveryMarker(task: Task, recoveryId: string): Task { - return { - ...task, - metadata: { - ...(task.metadata ?? {}), - [PAYMENT_RECOVERY_METADATA_KEY]: { version: 1, id: recoveryId }, - }, - } -} - -function preservePaymentRecoveryMarker(base: Task, source: Task): Task { - const marker = readPaymentRecoveryMarker(source) - return marker ? withPaymentRecoveryMarker(base, marker.id) : base -} - -function readPaymentRecoveryMarker(task: Task): TaskPaymentRecoveryMarker | undefined { - const raw = task.metadata?.[PAYMENT_RECOVERY_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined - const marker = raw as Partial - if (marker.version !== 1 || typeof marker.id !== 'string' || marker.id.length === 0) { - return undefined - } - return marker as TaskPaymentRecoveryMarker -} - -function clearPaymentRecoveryMarker(task: Task): Task { - if (!task.metadata || !(PAYMENT_RECOVERY_METADATA_KEY in task.metadata)) return task - const metadata = { ...task.metadata } - delete metadata[PAYMENT_RECOVERY_METADATA_KEY] - return Object.keys(metadata).length > 0 - ? { ...task, metadata } - : (() => { - const { metadata: _metadata, ...withoutMetadata } = task - return withoutMetadata - })() -} - function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): TaskStore { const hasCreateIfAbsent = typeof taskStore.createIfAbsent === 'function' const hasCompareAndSet = typeof taskStore.compareAndSet === 'function' @@ -1930,411 +639,23 @@ function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): } } -function buildFinalizationRecord( - authz: AuthorizedRequest, - receipt: SandboxUsageReceipt, - artifact: Artifact | null, - inputRequired: boolean, - inputRequiredPrompt: string | undefined, - finalState: FinalizationState = inputRequired ? 'input-required' : 'completed', -): FinalizationRecord { - const operation = authz.paymentOperation - return { - version: 1, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, - agentSlug: authz.agent.slug, - requestId: authz.requestId, - consumerId: authz.consumerId, - paymentMethod: authz.paymentMethod, - startMs: authz.startMs, - operationId: operation?.operationId ?? null, - paymentOperation: operation ? serializePaymentOperation(operation) : null, - receipt, - artifact, - inputRequired, - ...(inputRequiredPrompt ? { inputRequiredPrompt } : {}), - finalState, - maxOutputTokens: authz.maxOutputTokens, - executionBudget: authz.executionBudget, - usageRecorded: false, - } -} - -function withFinalizationRecord(task: Task, record: FinalizationRecord): Task { - return { - ...task, - metadata: { ...(task.metadata ?? {}), [FINALIZING_METADATA_KEY]: record }, - } -} - -function markUsageRecordedRecord(task: Task): Task { - const record = readFinalizationRecord(task) - if (!record || record.usageRecorded) return task - return withFinalizationRecord(task, { ...record, usageRecorded: true }) -} - -async function markUsageRecorded(taskStore: TaskStore, task: Task): Promise { - const marked = markUsageRecordedRecord(task) - if (marked === task) return task - if (await compareAndSetTask(taskStore, task, marked)) return marked - return await taskStore.get(task.id) ?? marked -} - -function withPaymentReleaseRecord(task: Task, record: PaymentReleaseRecord): Task { - return { - ...task, - metadata: { ...(task.metadata ?? {}), [PAYMENT_RELEASE_METADATA_KEY]: record }, - } -} - -function clearPaymentReleaseRecord(task: Task): Task { - if (!task.metadata || !(PAYMENT_RELEASE_METADATA_KEY in task.metadata)) return task - const metadata = { ...task.metadata } - delete metadata[PAYMENT_RELEASE_METADATA_KEY] - return Object.keys(metadata).length > 0 - ? { ...task, metadata } - : (() => { - const { metadata: _metadata, ...withoutMetadata } = task - return withoutMetadata - })() -} - -function readPaymentReleaseRecord(task: Task): PaymentReleaseRecord | undefined { - const raw = task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined - const record = raw as Partial - if ( - record.version !== 1 || - !record.lease || - typeof record.lease.id !== 'string' || - record.lease.id.length === 0 || - typeof record.lease.expiresAt !== 'number' || - !Number.isFinite(record.lease.expiresAt) || - typeof record.agentSlug !== 'string' || - record.agentSlug.length === 0 || - typeof record.requestId !== 'string' || - record.requestId.length === 0 || - typeof record.operationId !== 'string' || - record.operationId.length === 0 || - !record.paymentOperation || - typeof record.paymentOperation !== 'object' || - typeof record.reason !== 'string' - ) { - return undefined - } - if (record.paymentOperation.operationId !== record.operationId) return undefined - try { - deserializePaymentOperation(record.paymentOperation) - } catch { - return undefined - } - return record as PaymentReleaseRecord -} - -function buildPaymentReleaseRecord( - authz: AuthorizedRequest, - reason: string, -): PaymentReleaseRecord | undefined { - const operation = authz.paymentOperation - if (!operation) return undefined - return { - version: 1, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, - agentSlug: authz.agent.slug, - requestId: authz.requestId, - operationId: operation.operationId, - paymentOperation: serializePaymentOperation({ ...operation, state: 'releasing' }), - reason, - } -} - -async function beginPaymentReleaseRecovery( - taskStore: TaskStore, - task: Task, - authz: AuthorizedRequest, - reason: string, -): Promise { - const record = buildPaymentReleaseRecord(authz, reason) - if (!record) return undefined - for (let attempt = 0; attempt < 8; attempt += 1) { - const current = await taskStore.get(task.id) ?? task - const existing = readPaymentReleaseRecord(current) - if (existing) { - if (existing.operationId !== record.operationId) { - throw new Error('A2A task already has a different payment release recovery') - } - return current - } - const next = withPaymentReleaseRecord(current, record) - if (await compareAndSetTask(taskStore, current, next)) return next - } - throw new Error('A2A task changed before payment release recovery was stored') -} - -async function retainPaymentReleaseForRecovery( - taskStore: TaskStore, - taskId: string, - leaseId: string, - error: Error, -): Promise { - const current = await taskStore.get(taskId) - if (!current) return undefined - const record = readPaymentReleaseRecord(current) - if (!record || record.lease.id !== leaseId) return undefined - const retry: PaymentReleaseRecord = { - ...record, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, - recoveryAttempts: (record.recoveryAttempts ?? 0) + 1, - recoveryError: error.message, - } - const next = withPaymentReleaseRecord(current, retry) - if (await compareAndSetTask(taskStore, current, next)) return next - return await taskStore.get(taskId) -} - -async function clearPaymentReleaseRecovery( - taskStore: TaskStore, - task: Task, - leaseId: string, -): Promise { - const current = await taskStore.get(task.id) ?? task - const record = readPaymentReleaseRecord(current) - if (!record || record.lease.id !== leaseId) return current - const cleared = clearPaymentRecoveryMarker(clearPaymentReleaseRecord(current)) - if (await compareAndSetTask(taskStore, current, cleared)) return cleared - return await taskStore.get(task.id) ?? cleared -} - -function readFinalizationRecord(task: Task): FinalizationRecord | undefined { - const raw = task.metadata?.[FINALIZING_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined - const record = raw as Partial - if ( - record.version !== 1 || - !record.lease || - typeof record.lease.id !== 'string' || - typeof record.lease.expiresAt !== 'number' - ) { - return undefined - } - return record as FinalizationRecord -} - -function isTaskFinalizing(task: Task): boolean { - const marker = task.metadata?.[FINALIZING_METADATA_KEY] - return marker === true || (typeof marker === 'object' && marker !== null) -} - -async function recoverPaymentReleaseIfNeeded( - task: Task, - deps: A2AHandlerDeps, -): Promise { - const raw = task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] - if (raw === undefined) return task - const record = readPaymentReleaseRecord(task) - if (!record) { - return expirePaymentRelease( - task, - deps, - new Error('A2A payment release recovery record is missing'), - ) - } - if (record.lease.expiresAt > Date.now()) return task - - const renewed: PaymentReleaseRecord = { - ...record, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, - } - const leasedTask = withPaymentReleaseRecord(task, renewed) - if (!await compareAndSetTask(deps.taskStore, task, leasedTask)) { - return await deps.taskStore.get(task.id) ?? task - } - - try { - if (!deps.config.x402.paymentOperations) { - throw new Error('A2A payment release recovery is not configured') - } - const operation = deserializePaymentOperation(renewed.paymentOperation) - await deps.config.x402.paymentOperations.releasePayment(operation, renewed.reason) - if (deps.config.paymentRecovery) { - const recovered = await recoverDurablePayment(renewed.operationId, deps.config, { force: true }) - if (recovered && recovered.state !== 'reconciled') { - throw new Error('durable payment release is still pending') - } - } - const recovered = clearPaymentReleaseRecord(leasedTask) - if (!await compareAndSetTask(deps.taskStore, leasedTask, recovered)) { - return await deps.taskStore.get(task.id) ?? recovered - } - return recovered - } catch (error) { - const recoveryError = error instanceof Error ? error : new Error(String(error)) - const retained = await retainPaymentReleaseForRecovery( - deps.taskStore, - task.id, - renewed.lease.id, - recoveryError, - ) - console.error( - '[a2a] payment release recovery failed for ' + task.id + ':', - recoveryError.message, - ) - return retained ?? leasedTask - } -} - -async function recoverFinalizationIfNeeded( - task: Task, - deps: A2AHandlerDeps, - requestedAgentSlug: string, -): Promise { - if (!isTaskFinalizing(task)) return task - const record = readFinalizationRecord(task) - if (!record) { - return expireFinalization( - task, - deps, - null, - new Error('A2A finalization record is missing'), - ) - } - if (record.lease.expiresAt > Date.now()) return task - - const renewed: FinalizationRecord = { - ...record, - lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, - } - const leasedTask = withFinalizationRecord(task, renewed) - if (!await compareAndSetTask(deps.taskStore, task, leasedTask)) { - return await deps.taskStore.get(task.id) ?? task - } - - try { - const agentSlug = renewed.agentSlug || requestedAgentSlug - const agent = await deps.config.resolveAgent(agentSlug) - if (!agent || !agent.enabled) throw new Error('A2A recovery agent is unavailable') - - const paymentRecovery = readPaymentRecoveryMarker(leasedTask) - if (paymentRecovery && deps.config.paymentRecovery) { - const recovery = await recoverDurablePayment(paymentRecovery.id, deps.config, { - force: true, - usage: renewed.receipt, - }) - if (recovery?.state !== 'reconciled') { - throw new Error('durable payment finalization is still pending') - } - const usageRecordedTask = await markUsageRecorded(deps.taskStore, leasedTask) - const recoveredTask = finalizationResultTask(usageRecordedTask, renewed) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, recoveredTask)) { - return await deps.taskStore.get(task.id) ?? recoveredTask - } - await maybeDeliverPush(recoveredTask, deps) - return recoveredTask - } - - let paymentOperation: PaymentOperation | undefined - if (renewed.operationId || renewed.paymentOperation) { - if (!renewed.operationId || !renewed.paymentOperation) { - throw new Error('A2A payment operation recovery record is incomplete') - } - if (renewed.operationId !== renewed.paymentOperation.operationId) { - throw new Error('A2A payment operation recovery id does not match') - } - if (!deps.config.x402.paymentOperations) { - throw new Error('A2A payment operation recovery is not configured') - } - paymentOperation = deserializePaymentOperation(renewed.paymentOperation) - } - - let paymentAlreadySettled = false - if (paymentOperation && deps.config.x402.paymentOperations) { - const currentOperation = await deps.config.x402.paymentOperations.getPaymentOperation( - paymentOperation.operationId, - ) - if (currentOperation.state === 'not-found') { - throw new Error('A2A payment operation disappeared during finalization recovery') - } - if (currentOperation.operationId !== paymentOperation.operationId) { - throw new Error('A2A payment operation recovery returned a different operation') - } - paymentOperation = currentOperation - paymentAlreadySettled = currentOperation.state === 'settled' - } - - const authz: AuthorizedRequest = { - agent, - consumerId: renewed.consumerId, - paymentMethod: renewed.paymentMethod, - keyInfo: null, - userMessage: '[recovered A2A task]', - rateLimitRemaining: undefined, - requestId: renewed.requestId, - startMs: renewed.startMs, - maxOutputTokens: renewed.maxOutputTokens, - executionBudget: renewed.executionBudget, - requiredPaymentAmount: 0n, - paymentPayload: null, - ...(readPaymentRecoveryMarker(leasedTask) - ? { paymentRecoveryId: readPaymentRecoveryMarker(leasedTask)!.id } - : {}), - ...(paymentOperation - ? { paymentOperation, paymentOperationAcquired: true } - : {}), - } - let usageRecordedTask = leasedTask - await settleAndRecord( - agent, - authz, - renewed.receipt, - deps.config, - deps.state.obs, - { - usageAlreadyRecorded: renewed.usageRecorded === true, - paymentAlreadySettled, - onUsageRecorded: async () => { - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - }, - }, - ) - - usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) - const recovered = finalizationResultTask(usageRecordedTask, renewed) - if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, recovered)) { - return await deps.taskStore.get(task.id) ?? recovered - } - await maybeDeliverPush(recovered, deps) - return recovered - } catch (error) { - const recoveryError = error instanceof Error ? error : new Error(String(error)) - console.error( - `[a2a] finalization recovery failed for ${task.id}:`, - recoveryError.message, - ) - if ( - (readPaymentRecoveryMarker(leasedTask) && deps.config.paymentRecovery) || - (renewed.operationId && renewed.paymentOperation) - ) { - const retained = await retainFinalizationForRecovery( - deps.taskStore, - task.id, - renewed.lease.id, - recoveryError, - ) - if (retained) return retained - } - return expireFinalization(leasedTask, deps, renewed, recoveryError) - } -} - async function recoverTaskIfNeeded( task: Task, deps: A2AHandlerDeps, requestedAgentSlug: string, ): Promise { - const released = await recoverPaymentReleaseIfNeeded(task, deps) - const finalized = await recoverFinalizationIfNeeded(released, deps, requestedAgentSlug) - const paymentRecovered = await recoverPaymentMarkerIfNeeded(finalized, deps) - const submissionRecovered = await recoverSubmissionIfNeeded(paymentRecovered, deps) + const lifecycle = createTaskLifecycle(deps) + const paymentReleased = await recoverPaymentReleaseIfNeeded(task, lifecycle.payment) + const finalized = await recoverFinalizationIfNeeded( + paymentReleased, + lifecycle.finalization, + requestedAgentSlug, + ) + const paymentRecovered = await recoverPaymentMarkerIfNeeded(finalized, lifecycle.payment) + const submissionRecovered = await recoverSubmissionIfNeeded(paymentRecovered, { + taskStore: deps.taskStore, + deliverPush: lifecycle.finalization.deliverPush, + }) return recoverExpiredExecutionIfNeeded(submissionRecovered, deps) } @@ -2364,228 +685,6 @@ async function recoverExpiredExecutionIfNeeded(task: Task, deps: A2AHandlerDeps) return await deps.taskStore.get(task.id) ?? task } -async function recoverSubmissionIfNeeded(task: Task, deps: A2AHandlerDeps): Promise { - const raw = task.metadata?.[TASK_SUBMISSION_METADATA_KEY] - if (raw === undefined) return task - const submission = readTaskSubmission(task) - if (submission && submission.lease.expiresAt > Date.now()) return task - if (task.status.state !== 'submitted') { - return (await clearTaskSubmissionMarker(deps.taskStore, task)).task - } - - const cleanTask = clearTaskSubmission(task) - const failed: Task = { - ...withStatus(cleanTask, 'failed'), - metadata: { - ...(cleanTask.metadata ?? {}), - [TASK_SUBMISSION_RECOVERY_METADATA_KEY]: { - error: submission - ? 'A2A task submission lease expired before payment authorization completed' - : 'A2A task submission lease is invalid', - }, - }, - } - if (await compareAndSetTask(deps.taskStore, task, failed)) { - await maybeDeliverPush(failed, deps) - return failed - } - return await deps.taskStore.get(task.id) ?? task -} - -async function recoverPaymentMarkerIfNeeded( - task: Task, - deps: A2AHandlerDeps, -): Promise { - const marker = readPaymentRecoveryMarker(task) - if (!marker || !deps.config.paymentRecovery) return task - try { - const record = await recoverDurablePayment(marker.id, deps.config) - if (record?.state !== 'reconciled') return task - return clearReconciledPaymentRecoveryMarker(task, deps) - } catch (error) { - console.error( - `[a2a] durable payment recovery failed for ${task.id}:`, - error instanceof Error ? error.message : String(error), - ) - return task - } -} - -async function clearReconciledPaymentRecoveryMarker( - task: Task, - deps: A2AHandlerDeps, -): Promise { - const marker = readPaymentRecoveryMarker(task) - if (!marker || !deps.config.paymentRecovery) return task - const record = await deps.config.paymentRecovery.store.get(marker.id) - if (record?.state !== 'reconciled') return task - const cleared = clearPaymentRecoveryMarker(task) - if (cleared.status.state === 'working' || cleared.status.state === 'submitted') { - const failed: Task = { - ...withStatus(cleared, 'failed'), - metadata: { - ...(cleared.metadata ?? {}), - [EXECUTION_RECOVERY_METADATA_KEY]: { - error: 'payment recovery completed without a task result', - }, - }, - } - if (await compareAndSetTask(deps.taskStore, task, failed)) return failed - return await deps.taskStore.get(task.id) ?? failed - } - if (await compareAndSetTask(deps.taskStore, task, cleared)) return cleared - return await deps.taskStore.get(task.id) ?? cleared -} - -function finalizationResultTask(task: Task, record: FinalizationRecord): Task { - const cleanTask = clearPaymentRecoveryMarker(clearFinalizationMarker(task)) - const finalState = record.finalState ?? ( - task.status.state === 'canceled' - ? 'canceled' - : record.inputRequired - ? 'input-required' - : 'completed' - ) - if (finalState === 'canceled') { - return withStatus( - cleanTask, - 'canceled', - undefined, - record.artifact ? [record.artifact] : cleanTask.artifacts, - ) - } - if (finalState === 'input-required') { - return withStatus( - cleanTask, - 'input-required', - record.inputRequiredPrompt ? agentMessage(cleanTask, record.inputRequiredPrompt) : undefined, - record.artifact ? [record.artifact] : cleanTask.artifacts, - ) - } - return withStatus( - cleanTask, - 'completed', - undefined, - record.artifact ? [record.artifact] : cleanTask.artifacts, - ) -} - -async function expireFinalization( - task: Task, - deps: A2AHandlerDeps, - record: FinalizationRecord | null, - error: Error, -): Promise { - const cleanTask = clearFinalizationMarker(task) - const failed: Task = { - ...withStatus(cleanTask, 'failed'), - metadata: { - ...(cleanTask.metadata ?? {}), - gatewayFinalizationRecovery: { - operationId: record?.operationId ?? null, - error: error.message, - }, - }, - } - if (await compareAndSetTask(deps.taskStore, task, failed)) { - await maybeDeliverPush(failed, deps) - return failed - } - return await deps.taskStore.get(task.id) ?? failed -} - -async function expirePaymentRelease( - task: Task, - deps: A2AHandlerDeps, - error: Error, -): Promise { - const cleanTask = clearPaymentReleaseRecord(task) - const failed: Task = { - ...withStatus(cleanTask, 'failed'), - metadata: { - ...(cleanTask.metadata ?? {}), - gatewayPaymentReleaseRecovery: { error: error.message }, - }, - } - if (await compareAndSetTask(deps.taskStore, task, failed)) { - await maybeDeliverPush(failed, deps) - return failed - } - return await deps.taskStore.get(task.id) ?? failed -} - -function clearFinalizationMarker(task: Task): Task { - if (!task.metadata || !(FINALIZING_METADATA_KEY in task.metadata)) return task - const metadata = { ...task.metadata } - delete metadata[FINALIZING_METADATA_KEY] - return Object.keys(metadata).length > 0 - ? { ...task, metadata } - : (() => { - const { metadata: _metadata, ...withoutMetadata } = task - return withoutMetadata - })() -} - -function isTerminal(state: Task['status']['state']): boolean { - return TERMINAL_STATES.has(state) -} - -function asError(error: unknown): Error { - return error instanceof Error ? error : new Error(String(error)) -} - -function nowIso(): string { - return new Date().toISOString() -} - -function cryptoRandomId(): string { - return crypto.randomUUID().replace(/-/g, '') -} - -/** - * Build a new Task with an updated status (and optional artifacts). Centralises - * the timestamp + status structure so terminal transitions are written - * identically across every code path. - */ -function withStatus( - task: Task, - state: Task['status']['state'], - message?: Message, - artifacts?: Task['artifacts'], -): Task { - const next: Task = { - ...task, - status: { state, timestamp: nowIso(), ...(message ? { message } : {}) }, - ...(artifacts !== undefined ? { artifacts } : {}), - } - return isTerminal(state) || state === 'input-required' ? clearTaskExecution(next) : next -} - -/** - * Synthesize an agent-role message attached to a status (e.g. the - * input-required prompt text). messageId is deterministic-by-task so callers - * can dedupe on retry. - */ -function agentMessage(task: Task, text: string): Message { - return { - kind: 'message', - role: 'agent', - parts: [{ kind: 'text', text }], - messageId: `${task.id}-input-required-${stableMessageDigest(text)}`, - taskId: task.id, - contextId: task.contextId, - } -} - -function stableMessageDigest(value: string): string { - let hash = 2166136261 - for (let index = 0; index < value.length; index += 1) { - hash ^= value.charCodeAt(index) - hash = Math.imul(hash, 16777619) - } - return (hash >>> 0).toString(16).padStart(8, '0') -} - async function readJsonBody(request: Request): Promise { if (!request.body) return await request.json() const reader = request.body.getReader() @@ -2623,104 +722,22 @@ async function readJsonBody(request: Request): Promise { * via `tasks/get` to confirm state. */ async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise { - if (!deps.pushStore || !TERMINAL_STATES.has(task.status.state)) return - const webhookSecret = deps.config.a2a?.webhookSecret - const hasWebhookSecret = typeof webhookSecret === 'string' && webhookSecret.trim().length > 0 - if (!deps.config.x402.demoMode && !hasWebhookSecret) { - console.error(`[agent-gateway] production A2A push requires a webhookSecret for task ${task.id}`) - return - } - try { - const deliveryTask = clearPushDeliveryClaims(task) - const deliveryArgs: Omit = { - task: deliveryTask, - store: deps.pushStore, - fetcher: deps.config.a2a?.pushFetcher, - urlValidator: deps.config.a2a?.pushUrlValidator, - requireUrlValidator: !deps.config.x402.demoMode, - claimDelivery: (taskId, configId, terminalState) => claimTaskPushDelivery( - deps.taskStore, - taskId, - configId, - terminalState, - ), - onDelivery: (result: PushDeliveryResult) => { - if (!result.ok) { - deps.state.obs?.onStreamError?.( - { requestId: result.taskId, agentSlug: task.id, startMs: Date.now() }, - { - consumerId: result.configId, - errorMessage: `push delivery failed (${result.status ?? 'no-status'}): ${result.error ?? 'non-2xx'}`, - }, - ) - } - }, - } - if (hasWebhookSecret) { - await deliverPushNotifications({ ...deliveryArgs, webhookSecret }) - } else { - await deliverDemoPushNotifications(deliveryArgs) - } - } catch (err) { - // Catastrophic failure of the push pipeline itself (e.g. the store threw). - // Logged but never escalated — a busted webhook MUST NOT fail the agent. - console.error( - `[agent-gateway] push delivery threw for task ${task.id}: ${err instanceof Error ? err.message : String(err)}`, - ) - } -} - -async function claimTaskPushDelivery( - taskStore: TaskStore, - taskId: string, - configId: string, - terminalState: Task['status']['state'], -): Promise { - if (!TERMINAL_STATES.has(terminalState)) return false - for (let attempt = 0; attempt < 16; attempt += 1) { - const current = await taskStore.get(taskId) - if (!current || current.status.state !== terminalState) return false - const existing = readPushDeliveryClaims(current) - if (existing?.claims[configId] === terminalState) return false - const next: Task = { - ...current, - metadata: { - ...(current.metadata ?? {}), - [PUSH_DELIVERY_METADATA_KEY]: { - version: 1, - claims: { - ...(existing?.claims ?? {}), - [configId]: terminalState, - }, - } satisfies TaskPushDeliveryClaims, - }, - } - if (await compareAndSetTask(taskStore, current, next)) return true + const pushDeps: PushDeliveryDependencies = { + taskStore: deps.taskStore, + pushStore: deps.pushStore, + demoMode: deps.config.x402.demoMode === true, + webhookSecret: deps.config.a2a?.webhookSecret, + fetcher: deps.config.a2a?.pushFetcher, + urlValidator: deps.config.a2a?.pushUrlValidator, + onDeliveryFailure: (failedTask, result) => { + void deps.state.obs?.onStreamError?.( + { requestId: result.taskId, agentSlug: failedTask.id, startMs: Date.now() }, + { + consumerId: result.configId, + errorMessage: `push delivery failed (${result.status ?? 'no-status'}): ${result.error ?? 'non-2xx'}`, + }, + ) + }, } - throw new Error(`A2A push delivery claim changed too many times for task '${taskId}'`) -} - -function readPushDeliveryClaims(task: Task): TaskPushDeliveryClaims | undefined { - const raw = task.metadata?.[PUSH_DELIVERY_METADATA_KEY] - if (!raw || typeof raw !== 'object') return undefined - const record = raw as Partial - if (!record.claims || typeof record.claims !== 'object') return undefined - const claims = Object.fromEntries( - Object.entries(record.claims).filter(([, state]) => - typeof state === 'string' && TERMINAL_STATES.has(state as Task['status']['state']), - ), - ) as Record - return record.version === 1 ? { version: 1, claims } : undefined -} - -function clearPushDeliveryClaims(task: Task): Task { - if (!task.metadata || !(PUSH_DELIVERY_METADATA_KEY in task.metadata)) return task - const metadata = { ...task.metadata } - delete metadata[PUSH_DELIVERY_METADATA_KEY] - return Object.keys(metadata).length > 0 - ? { ...task, metadata } - : (() => { - const { metadata: _metadata, ...withoutMetadata } = task - return withoutMetadata - })() + await deliverTaskPush(task, pushDeps) } diff --git a/src/a2a/message-send-execution.ts b/src/a2a/message-send-execution.ts new file mode 100644 index 0000000..1d5d3be --- /dev/null +++ b/src/a2a/message-send-execution.ts @@ -0,0 +1,241 @@ +import type { Context } from 'hono' +import { + type AuthorizedRequest, + dispatchSandboxStreamRich, + beginPaymentExecution, + markPaymentExecutionStarted, + renewPaymentExecution, +} from '../dispatch' +import { claimTaskExecution, renewTaskExecution } from './execution-fence' +import type { GatewayConfig, SandboxUsageReceipt } from '../types' +import type { TaskStore } from './task-store' +import { A2A_ERROR_CODES, type JSONRPCRequest, type Task } from './types' +import { clearTaskSubmission } from './task-submission-recovery' +import { + buildFinalizationRecord, + clearFinalizationMarker, + completeCanceledTask, + markUsageRecorded, + retainFinalizationForRecovery, + withFinalizationRecord, +} from './task-finalization' +import { clearPaymentRecoveryMarker, releaseTaskPayment } from './payment-recovery' +import { + agentMessage, + asError, + compareAndSetTask, + isTerminal, + nowIso, + persistTaskIfCurrent, + shouldPreserveTask, + withStatus, +} from './task-state' +import { fail, ok } from './jsonrpc' +import { responseTextToArtifact } from './translate' +import type { TaskFinalizationDependencies } from './task-finalization' +import type { PaymentRecoveryDependencies } from './payment-recovery' + +export interface MessageExecutionDependencies { + taskStore: TaskStore + config: GatewayConfig + payment: PaymentRecoveryDependencies + finalization: TaskFinalizationDependencies +} + +export async function executeMessageSend( + c: Context, + req: JSONRPCRequest, + deps: MessageExecutionDependencies, + authz: AuthorizedRequest, + task: Task, + signal: AbortSignal, +): Promise { + if (isTerminal(task.status.state)) return c.json(ok(req.id, task)) + const taskWithoutSubmission = clearTaskSubmission(task) + let workingTask: Task = task.status.state === 'working' + ? taskWithoutSubmission + : { ...taskWithoutSubmission, status: { state: 'working', timestamp: nowIso() } } + if ( + JSON.stringify(task) !== JSON.stringify(workingTask) && + !await compareAndSetTask(deps.taskStore, task, workingTask) + ) { + await releaseTaskPayment(authz, task, deps.payment, 'A2A task changed before execution started', false) + if (signal.aborted) { + const canceled = await deps.taskStore.get(task.id) + if (canceled?.status.state === 'canceled') { + await deps.finalization.deliverPush(canceled) + return c.json(ok(req.id, canceled)) + } + } + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) + } + + let responseText = '' + let usage: SandboxUsageReceipt | undefined + let workObserved = false + let inputRequiredPrompt: string | undefined + let inputRequiredSeen = false + let finalizationLeaseId: string | undefined + try { + for await (const event of dispatchSandboxStreamRich( + authz.agent, + authz.userMessage, + authz.consumerId, + deps.config, + signal, + task.id, + authz.maxOutputTokens, + async () => { + workingTask = await claimTaskExecution(deps.taskStore, workingTask, authz.requestId) + await beginPaymentExecution(authz, deps.config) + }, + authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, + async () => { + workObserved = true + await markPaymentExecutionStarted(authz, deps.config) + }, + authz.executionBudget.maxInputTokens, + async () => { + workingTask = await renewTaskExecution(deps.taskStore, task.id, authz.requestId) + await renewPaymentExecution(authz, deps.config) + }, + )) { + if (event.kind === 'text') { + responseText += event.delta + workObserved = true + } else if (event.kind === 'activity') { + workObserved = true + } else if (event.kind === 'usage') { + usage = event.usage + } else { + inputRequiredSeen = true + inputRequiredPrompt = event.prompt + workObserved = true + } + } + } catch (err) { + const releasedTask = await releaseTaskPayment( + authz, + workingTask, + deps.payment, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, + ) + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = shouldPreserveTask(currentTask) + ? currentTask + : withStatus(clearTaskSubmission(currentTask), 'failed') + try { + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) + await deps.finalization.deliverPush(persisted) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, err instanceof Error ? err.message : String(err))) + } + + if (signal.aborted) { + const canceled = await completeCanceledTask( + authz, + workingTask, + responseText, + usage, + workObserved, + deps.finalization, + ) + return c.json(ok(req.id, canceled)) + } + + try { + if (!usage) throw new Error('sandbox did not provide a usage receipt') + const finalizationArtifact = responseText + ? responseTextToArtifact(responseText, `${task.id}-artifact-0`) + : task.artifacts?.[0] ?? null + const finalization = buildFinalizationRecord( + authz, + usage, + finalizationArtifact, + inputRequiredSeen, + inputRequiredPrompt, + ) + const finalizingTask = withFinalizationRecord(workingTask, finalization) + if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask?.status.state === 'canceled') { + const canceled = await completeCanceledTask( + authz, + currentTask, + responseText, + usage, + workObserved, + deps.finalization, + ) + return c.json(ok(req.id, canceled)) + } + throw new Error('A2A task changed before payment settlement') + } + finalizationLeaseId = finalization.lease.id + let usageRecordedTask = finalizingTask + await deps.finalization.settle(authz, usage, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + const settledBase = clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)) + const result = inputRequiredSeen + ? withStatus( + settledBase, + 'input-required', + inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, + responseText ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] : task.artifacts, + ) + : withStatus(settledBase, 'completed', undefined, [ + responseTextToArtifact(responseText, `${task.id}-artifact-0`), + ]) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, result)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask && (isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required')) { + return c.json(ok(req.id, currentTask)) + } + throw new Error('A2A task changed after payment settlement') + } + if (inputRequiredSeen) return c.json(ok(req.id, result)) + await deps.finalization.deliverPush(result) + return c.json(ok(req.id, result)) + } catch (err) { + const releasedTask = await releaseTaskPayment( + authz, + workingTask, + deps.payment, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, + ) + if (finalizationLeaseId) { + await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalizationLeaseId, + asError(err), + ) + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) + } + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = shouldPreserveTask(currentTask) + ? currentTask + : withStatus(clearTaskSubmission(currentTask), 'failed') + try { + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) + await deps.finalization.deliverPush(persisted) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'Payment settlement failed')) + } +} diff --git a/src/a2a/message-stream-execution.ts b/src/a2a/message-stream-execution.ts new file mode 100644 index 0000000..ce455c8 --- /dev/null +++ b/src/a2a/message-stream-execution.ts @@ -0,0 +1,392 @@ +import type { Context } from 'hono' +import { + type AuthorizedRequest, + beginPaymentExecution, + dispatchSandboxStreamRich, + markPaymentExecutionStarted, + renewPaymentExecution, +} from '../dispatch' +import type { GatewayConfig, SandboxUsageReceipt } from '../types' +import { claimTaskExecution, renewTaskExecution } from './execution-fence' +import { fail, ok } from './jsonrpc' +import { + clearPaymentRecoveryMarker, + releaseTaskPayment, +} from './payment-recovery' +import { + buildFinalizationRecord, + clearFinalizationMarker, + completeCanceledTask, + markUsageRecorded, + retainFinalizationForRecovery, + withFinalizationRecord, +} from './task-finalization' +import { clearTaskSubmission } from './task-submission-recovery' +import { bindRequestAbort, type TaskCancellationRegistry } from './task-cancellation' +import type { TaskLifecycle } from './task-lifecycle' +import { + agentMessage, + asError, + compareAndSetTask, + isTerminal, + nowIso, + persistTaskIfCurrent, + shouldPreserveTask, + withStatus, +} from './task-state' +import type { TaskStateStore } from './task-state' +import { + A2A_ERROR_CODES, + type JSONRPCRequest, + type StreamingEvent, + type Task, + type TaskArtifactUpdateEvent, + type TaskStatusUpdateEvent, +} from './types' +import { responseTextToArtifact } from './translate' + +export interface MessageStreamExecutionDependencies { + taskStore: TaskStateStore + config: GatewayConfig + lifecycle: TaskLifecycle + cancels: TaskCancellationRegistry + deliverPush: (task: Task) => Promise + reportStreamError: (authz: AuthorizedRequest, error: unknown) => Promise +} + +export async function executeMessageStream( + c: Context, + req: JSONRPCRequest, + deps: MessageStreamExecutionDependencies, + authz: AuthorizedRequest, + task: Task, +): Promise { + const controller = deps.cancels.register(task.id) + const lifecycle = deps.lifecycle + const detachRequestAbort = bindRequestAbort(c.req.raw.signal, controller) + const workingStatus: TaskStatusUpdateEvent = { + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: { state: 'working', timestamp: nowIso() }, + final: false, + } + if (isTerminal(task.status.state)) { + detachRequestAbort() + deps.cancels.clear(task.id) + return c.json(ok(req.id, task)) + } + let workingTask: Task = task.status.state === 'working' + ? task + : { ...task, status: workingStatus.status } + if (task.status.state !== 'working' && !await compareAndSetTask(deps.taskStore, task, workingTask)) { + detachRequestAbort() + deps.cancels.clear(task.id) + const current = await releaseTaskPayment( + authz, + await deps.taskStore.get(task.id) ?? task, + lifecycle.payment, + 'A2A task changed before execution started', + false, + ) + if (current.status.state === 'canceled') return c.json(ok(req.id, current)) + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' changed before execution`)) + } + let responseText = '' + let usage: SandboxUsageReceipt | undefined + let workObserved = false + + const stream = new ReadableStream({ + start(ctrl) { + void (async () => { + const encoder = new TextEncoder() + const send = (event: StreamingEvent) => { + if (ctrl.desiredSize === null) return + try { + ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) + } catch { + // The client can cancel between the desiredSize check and enqueue. + } + } + + let inputRequiredPrompt: string | undefined + let inputRequiredSeen = false + let finalizationLeaseId: string | undefined + try { + send(workingStatus) + + for await (const event of dispatchSandboxStreamRich( + authz.agent, + authz.userMessage, + authz.consumerId, + deps.config, + controller.signal, + task.id, + authz.maxOutputTokens, + async () => { + workingTask = await claimTaskExecution(deps.taskStore, workingTask, authz.requestId) + await beginPaymentExecution(authz, deps.config) + }, + authz.paymentOperation !== undefined || authz.mppChargeOperation !== undefined, + async () => { + workObserved = true + await markPaymentExecutionStarted(authz, deps.config) + }, + authz.executionBudget.maxInputTokens, + async () => { + workingTask = await renewTaskExecution(deps.taskStore, task.id, authz.requestId) + await renewPaymentExecution(authz, deps.config) + }, + )) { + if (event.kind === 'text') { + responseText += event.delta + workObserved = true + const artifactEvent: TaskArtifactUpdateEvent = { + kind: 'artifact-update', + taskId: task.id, + contextId: task.contextId, + artifact: { + artifactId: `${task.id}-artifact-0`, + name: 'response', + parts: [{ kind: 'text', text: event.delta }], + }, + append: true, + } + send(artifactEvent) + } else if (event.kind === 'activity') { + workObserved = true + } else if (event.kind === 'usage') { + usage = event.usage + } else { + inputRequiredSeen = true + inputRequiredPrompt = event.prompt + workObserved = true + } + } + + if (controller.signal.aborted) { + const canceled = await completeCanceledTask( + authz, + workingTask, + responseText, + usage, + workObserved, + lifecycle.finalization, + ) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: canceled.status, + final: true, + }) + return + } + + if (!usage) throw new Error('sandbox did not provide a usage receipt') + const finalizationArtifact = responseTextToArtifact(responseText, `${task.id}-artifact-0`) + const finalization = buildFinalizationRecord( + authz, + usage, + finalizationArtifact, + inputRequiredSeen, + inputRequiredPrompt, + ) + // Let the durable task-store CAS decide the cancellation race. + const finalizingTask = withFinalizationRecord(workingTask, finalization) + if (!await compareAndSetTask(deps.taskStore, workingTask, finalizingTask)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask?.status.state === 'canceled') { + const canceled = await completeCanceledTask( + authz, + currentTask, + responseText, + usage, + workObserved, + lifecycle.finalization, + ) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: canceled.status, + final: true, + }) + return + } + await releaseTaskPayment( + authz, + task, + lifecycle.payment, + 'A2A task changed before payment settlement', + workObserved || usage !== undefined, + ) + return + } + finalizationLeaseId = finalization.lease.id + deps.cancels.beginFinalization(task.id) + let usageRecordedTask = finalizingTask + await lifecycle.finalization.settle(authz, usage, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + + if (inputRequiredSeen) { + const paused = withStatus( + clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), + 'input-required', + inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined, + responseText + ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] + : task.artifacts, + ) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, paused)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: currentTask.status, + final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', + }) + } + return + } + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: paused.status, + final: true, + }) + return + } + + const completed = withStatus( + clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), + 'completed', + undefined, + [responseTextToArtifact(responseText, `${task.id}-artifact-0`)], + ) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, completed)) { + const currentTask = await deps.taskStore.get(task.id) + if (currentTask) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: currentTask.status, + final: isTerminal(currentTask.status.state) || currentTask.status.state === 'input-required', + }) + } + return + } + send({ + kind: 'artifact-update', + taskId: task.id, + contextId: task.contextId, + artifact: { + artifactId: `${task.id}-artifact-0`, + name: 'response', + parts: [{ kind: 'text', text: '' }], + }, + append: true, + lastChunk: true, + }) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: completed.status, + final: true, + }) + await lifecycle.finalization.deliverPush(completed) + } catch (err) { + const releasedTask = await releaseTaskPayment( + authz, + task, + lifecycle.payment, + err instanceof Error ? err.message : String(err), + workObserved || usage !== undefined, + ) + if (finalizationLeaseId) { + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalizationLeaseId, + asError(err), + ) + if (retained) { + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: retained.status, + final: false, + }) + } + return + } + const currentTask = await deps.taskStore.get(task.id) ?? releasedTask + const failed = shouldPreserveTask(currentTask) + ? currentTask + : withStatus(clearTaskSubmission(currentTask), 'failed') + try { + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask, failed) + send({ + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: persisted.status, + final: true, + }) + await deps.deliverPush(persisted) + } catch (taskError) { + console.error( + `[a2a] failed to persist failed task ${task.id}:`, + taskError instanceof Error ? taskError.message : String(taskError), + ) + } + try { + await deps.reportStreamError(authz, err) + } catch (observerError) { + console.error( + `[a2a] stream observer failed for ${authz.requestId}:`, + observerError instanceof Error ? observerError.message : String(observerError), + ) + } + } finally { + detachRequestAbort() + deps.cancels.clear(task.id) + try { + if (ctrl.desiredSize !== null) ctrl.close() + } catch { + // The response may already be closed by client cancellation. + } + } + })() + }, + cancel() { + controller.abort() + }, + }) + + return new Response(stream, { + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + 'X-Request-Id': authz.requestId, + 'X-Agent-Slug': authz.agent.slug, + 'X-Task-Id': task.id, + ...(authz.mppChargeOperation + ? { 'Payment-Receipt': authz.mppChargeOperation.receipt } + : {}), + ...(authz.paymentRecoveryId + ? { 'X-Payment-Operation-Id': authz.paymentRecoveryId } + : {}), + }, + }) +} diff --git a/src/a2a/payment-recovery.ts b/src/a2a/payment-recovery.ts new file mode 100644 index 0000000..de1b56a --- /dev/null +++ b/src/a2a/payment-recovery.ts @@ -0,0 +1,431 @@ +import type { AuthorizedRequest } from '../dispatch' +import type { PaymentOperations, PaymentOperation } from '../payment-operations' +import { + deserializePaymentOperation, + serializePaymentOperation, + type PaymentRecoveryConfig, + type PaymentRecoveryRecord, + type SerializedPaymentOperation, +} from '../payment-recovery' +import type { SandboxUsageReceipt } from '../types' +import type { Task } from './types' +import { + asError, + compareAndSetTask, + clearTaskMetadata, + cryptoRandomId, + withStatus, + type TaskStateStore, +} from './task-state' + +const PAYMENT_RELEASE_METADATA_KEY = 'gatewayPaymentRelease' +const PAYMENT_RECOVERY_METADATA_KEY = 'gatewayPaymentRecovery' +const PAYMENT_RELEASE_LEASE_MS = 5 * 60 * 1000 + +interface TaskPaymentRecoveryMarker { + version: 1 + id: string +} + +interface PaymentReleaseRecord { + version: 1 + lease: { id: string; expiresAt: number } + agentSlug: string + requestId: string + operationId: string + paymentOperation: SerializedPaymentOperation + reason: string + recoveryAttempts?: number + recoveryError?: string +} + +export interface PaymentRecoveryDependencies { + taskStore: TaskStateStore + paymentOperations?: PaymentOperations + paymentRecovery?: PaymentRecoveryConfig + releasePayment: (authz: AuthorizedRequest, reason: string) => Promise + releasePaymentAfterFailure: ( + authz: AuthorizedRequest, + reason: string, + workObserved: boolean, + ) => Promise + recoverDurablePayment: ( + recoveryId: string, + options?: { force?: boolean; usage?: SandboxUsageReceipt }, + ) => Promise + deliverPush: (task: Task) => Promise +} + +export async function attachPaymentRecoveryMarker( + taskStore: TaskStateStore, + task: Task, + recoveryId: string | undefined, +): Promise { + if (!recoveryId) return task + const existing = readPaymentRecoveryMarker(task) + if (existing) { + if (existing.id !== recoveryId) { + throw new Error('A2A task already has a different payment recovery identity') + } + return task + } + const next = withPaymentRecoveryMarker(task, recoveryId) + if (await compareAndSetTask(taskStore, task, next)) return next + throw new Error('A2A task changed while payment recovery was attached') +} + +/** Attach only as a retention marker. The returned task must never execute. */ +export async function retainPaymentRecoveryMarker( + taskStore: TaskStateStore, + task: Task, + recoveryId: string | undefined, +): Promise { + if (!recoveryId) return task + let current = await taskStore.get(task.id) ?? task + for (let attempt = 0; attempt < 16; attempt += 1) { + const existing = readPaymentRecoveryMarker(current) + if (existing) { + if (existing.id !== recoveryId) { + throw new Error('A2A task already has a different payment recovery identity') + } + return current + } + const next = withPaymentRecoveryMarker(current, recoveryId) + if (await compareAndSetTask(taskStore, current, next)) return next + const latest = await taskStore.get(task.id) + if (!latest) throw new Error('A2A task disappeared while payment recovery was retained') + current = latest + } + throw new Error('A2A task changed too many times while payment recovery was retained') +} + +export function preservePaymentRecoveryMarker(base: Task, source: Task): Task { + const marker = readPaymentRecoveryMarker(source) + return marker ? withPaymentRecoveryMarker(base, marker.id) : base +} + +export function readPaymentRecoveryMarker(task: Task): TaskPaymentRecoveryMarker | undefined { + const raw = task.metadata?.[PAYMENT_RECOVERY_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const marker = raw as Partial + if (marker.version !== 1 || typeof marker.id !== 'string' || marker.id.length === 0) { + return undefined + } + return marker as TaskPaymentRecoveryMarker +} + +export function clearPaymentRecoveryMarker(task: Task): Task { + return clearTaskMetadata(task, PAYMENT_RECOVERY_METADATA_KEY) +} + +export function hasPaymentReleaseRecovery(task: Task): boolean { + return task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] !== undefined +} + +export async function releaseTaskPayment( + authz: AuthorizedRequest, + task: Task, + deps: PaymentRecoveryDependencies, + reason: string, + workObserved: boolean, +): Promise { + if ( + !workObserved && + !authz.paymentRecoveryId && + authz.paymentOperation && + deps.paymentOperations + ) { + let marked: Task + try { + marked = await beginPaymentReleaseRecovery(deps.taskStore, task, authz, reason) ?? task + } catch (error) { + console.error( + '[a2a] failed to persist payment release recovery for ' + authz.requestId + ':', + error instanceof Error ? error.message : String(error), + ) + return await deps.taskStore.get(task.id) ?? task + } + const record = readPaymentReleaseRecord(marked) + if (!record) return marked + try { + await deps.releasePayment(authz, reason) + } catch (releaseError) { + const retained = await retainPaymentReleaseForRecovery( + deps.taskStore, + task.id, + record.lease.id, + asError(releaseError), + ) + console.error( + '[a2a] payment release retained for ' + authz.requestId + ':', + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + return retained ?? marked + } + return clearPaymentReleaseRecovery(deps.taskStore, marked, record.lease.id) + } + + try { + await deps.releasePaymentAfterFailure(authz, reason, workObserved) + } catch (releaseError) { + console.error( + `[a2a] payment release failed for ${authz.requestId}:`, + releaseError instanceof Error ? releaseError.message : String(releaseError), + ) + } + const current = await deps.taskStore.get(task.id) ?? task + return workObserved + ? current + : clearReconciledPaymentRecoveryMarker(current, deps) +} + +async function beginPaymentReleaseRecovery( + taskStore: TaskStateStore, + task: Task, + authz: AuthorizedRequest, + reason: string, +): Promise { + const record = buildPaymentReleaseRecord(authz, reason) + if (!record) return undefined + for (let attempt = 0; attempt < 8; attempt += 1) { + const current = await taskStore.get(task.id) ?? task + const existing = readPaymentReleaseRecord(current) + if (existing) { + if (existing.operationId !== record.operationId) { + throw new Error('A2A task already has a different payment release recovery') + } + return current + } + const next = withPaymentReleaseRecord(current, record) + if (await compareAndSetTask(taskStore, current, next)) return next + } + throw new Error('A2A task changed before payment release recovery was stored') +} + +async function retainPaymentReleaseForRecovery( + taskStore: TaskStateStore, + taskId: string, + leaseId: string, + error: Error, +): Promise { + const current = await taskStore.get(taskId) + if (!current) return undefined + const record = readPaymentReleaseRecord(current) + if (!record || record.lease.id !== leaseId) return undefined + const retry: PaymentReleaseRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, + recoveryAttempts: (record.recoveryAttempts ?? 0) + 1, + recoveryError: error.message, + } + const next = withPaymentReleaseRecord(current, retry) + if (await compareAndSetTask(taskStore, current, next)) return next + return await taskStore.get(taskId) +} + +async function clearPaymentReleaseRecovery( + taskStore: TaskStateStore, + task: Task, + leaseId: string, +): Promise { + const current = await taskStore.get(task.id) ?? task + const record = readPaymentReleaseRecord(current) + if (!record || record.lease.id !== leaseId) return current + const cleared = clearPaymentRecoveryMarker(clearPaymentReleaseRecord(current)) + if (await compareAndSetTask(taskStore, current, cleared)) return cleared + return await taskStore.get(task.id) ?? cleared +} + +export async function recoverPaymentReleaseIfNeeded( + task: Task, + deps: PaymentRecoveryDependencies, +): Promise { + const raw = task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] + if (raw === undefined) return task + const record = readPaymentReleaseRecord(task) + if (!record) { + return expirePaymentRelease( + task, + deps, + new Error('A2A payment release recovery record is missing'), + ) + } + if (record.lease.expiresAt > Date.now()) return task + + const renewed: PaymentReleaseRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, + } + const leasedTask = withPaymentReleaseRecord(task, renewed) + if (!await compareAndSetTask(deps.taskStore, task, leasedTask)) { + return await deps.taskStore.get(task.id) ?? task + } + + try { + if (!deps.paymentOperations) { + throw new Error('A2A payment release recovery is not configured') + } + const operation = deserializePaymentOperation(renewed.paymentOperation) + await deps.paymentOperations.releasePayment(operation, renewed.reason) + if (deps.paymentRecovery) { + const recovered = await deps.recoverDurablePayment(renewed.operationId, { force: true }) + if (recovered && recovered.state !== 'reconciled') { + throw new Error('durable payment release is still pending') + } + } + const recovered = clearPaymentReleaseRecord(leasedTask) + if (!await compareAndSetTask(deps.taskStore, leasedTask, recovered)) { + return await deps.taskStore.get(task.id) ?? recovered + } + return recovered + } catch (error) { + const recoveryError = asError(error) + const retained = await retainPaymentReleaseForRecovery( + deps.taskStore, + task.id, + renewed.lease.id, + recoveryError, + ) + console.error( + '[a2a] payment release recovery failed for ' + task.id + ':', + recoveryError.message, + ) + return retained ?? leasedTask + } +} + +export async function recoverPaymentMarkerIfNeeded( + task: Task, + deps: PaymentRecoveryDependencies, +): Promise { + const marker = readPaymentRecoveryMarker(task) + if (!marker || !deps.paymentRecovery) return task + try { + const record = await deps.recoverDurablePayment(marker.id) + if (record?.state !== 'reconciled') return task + return clearReconciledPaymentRecoveryMarker(task, deps) + } catch (error) { + console.error( + `[a2a] durable payment recovery failed for ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + return task + } +} + +async function clearReconciledPaymentRecoveryMarker( + task: Task, + deps: PaymentRecoveryDependencies, +): Promise { + const marker = readPaymentRecoveryMarker(task) + if (!marker || !deps.paymentRecovery) return task + const record = await deps.paymentRecovery.store.get(marker.id) + if (record?.state !== 'reconciled') return task + const cleared = clearPaymentRecoveryMarker(task) + if (cleared.status.state === 'working' || cleared.status.state === 'submitted') { + const failed: Task = { + ...withStatus(cleared, 'failed'), + metadata: { + ...(cleared.metadata ?? {}), + gatewayExecutionRecovery: { + error: 'payment recovery completed without a task result', + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) return failed + return await deps.taskStore.get(task.id) ?? failed + } + if (await compareAndSetTask(deps.taskStore, task, cleared)) return cleared + return await deps.taskStore.get(task.id) ?? cleared +} + +function withPaymentRecoveryMarker(task: Task, recoveryId: string): Task { + return { + ...task, + metadata: { + ...(task.metadata ?? {}), + [PAYMENT_RECOVERY_METADATA_KEY]: { version: 1, id: recoveryId }, + }, + } +} + +function withPaymentReleaseRecord(task: Task, record: PaymentReleaseRecord): Task { + return { + ...task, + metadata: { ...(task.metadata ?? {}), [PAYMENT_RELEASE_METADATA_KEY]: record }, + } +} + +function clearPaymentReleaseRecord(task: Task): Task { + return clearTaskMetadata(task, PAYMENT_RELEASE_METADATA_KEY) +} + +function readPaymentReleaseRecord(task: Task): PaymentReleaseRecord | undefined { + const raw = task.metadata?.[PAYMENT_RELEASE_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const record = raw as Partial + if ( + record.version !== 1 || + !record.lease || + typeof record.lease.id !== 'string' || + record.lease.id.length === 0 || + typeof record.lease.expiresAt !== 'number' || + !Number.isFinite(record.lease.expiresAt) || + typeof record.agentSlug !== 'string' || + record.agentSlug.length === 0 || + typeof record.requestId !== 'string' || + record.requestId.length === 0 || + typeof record.operationId !== 'string' || + record.operationId.length === 0 || + !record.paymentOperation || + typeof record.paymentOperation !== 'object' || + typeof record.reason !== 'string' + ) { + return undefined + } + if (record.paymentOperation.operationId !== record.operationId) return undefined + try { + deserializePaymentOperation(record.paymentOperation) + } catch { + return undefined + } + return record as PaymentReleaseRecord +} + +function buildPaymentReleaseRecord( + authz: AuthorizedRequest, + reason: string, +): PaymentReleaseRecord | undefined { + const operation = authz.paymentOperation + if (!operation) return undefined + return { + version: 1, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + PAYMENT_RELEASE_LEASE_MS }, + agentSlug: authz.agent.slug, + requestId: authz.requestId, + operationId: operation.operationId, + paymentOperation: serializePaymentOperation({ ...operation, state: 'releasing' }), + reason, + } +} + +async function expirePaymentRelease( + task: Task, + deps: PaymentRecoveryDependencies, + error: Error, +): Promise { + const cleanTask = clearPaymentReleaseRecord(task) + const failed: Task = { + ...cleanTask, + status: { state: 'failed', timestamp: new Date().toISOString() }, + metadata: { + ...(cleanTask.metadata ?? {}), + gatewayPaymentReleaseRecovery: { error: error.message }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await deps.deliverPush(failed) + return failed + } + return await deps.taskStore.get(task.id) ?? failed +} diff --git a/src/a2a/push-config-methods.ts b/src/a2a/push-config-methods.ts new file mode 100644 index 0000000..032d9be --- /dev/null +++ b/src/a2a/push-config-methods.ts @@ -0,0 +1,158 @@ +import type { Context } from 'hono' +import { + validatePushNotificationUrl, + type PushNotificationStore, + type TaskPushNotificationConfig, +} from './push-notifications' +import type { TaskStateStore } from './task-state' +import { + A2A_ERROR_CODES, + type JSONRPCRequest, + type Task, + type TaskIdParams, + type TaskPushNotificationConfigGetParams, +} from './types' +import { fail, ok } from './jsonrpc' + +export interface PushConfigMethodDependencies { + taskStore: TaskStateStore + pushStore?: PushNotificationStore + demoMode: boolean + urlValidator?: (url: URL) => boolean | Promise + authorizeTaskAccess: ( + c: Context, + req: JSONRPCRequest, + task: Task, + ) => Promise +} + +export async function handlePushSet( + c: Context, + req: JSONRPCRequest, + deps: PushConfigMethodDependencies, +): Promise { + if (!deps.pushStore) { + return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) + } + const params = req.params as TaskPushNotificationConfig | undefined + if (!params || typeof params.taskId !== 'string' || !params.pushNotificationConfig?.id) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.taskId and params.pushNotificationConfig.id required'), + ) + } + if (typeof params.pushNotificationConfig.url !== 'string') { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url required')) + } + const task = await deps.taskStore.get(params.taskId) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.taskId}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, task) + if (accessError) return accessError + const pushUrl = validatePushNotificationUrl(params.pushNotificationConfig.url) + if (!pushUrl) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url is not a safe HTTPS destination'), + ) + } + if (!deps.demoMode && !deps.urlValidator) { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'production push URL validation is not configured'), + ) + } + let allowedByHostPolicy = true + try { + if (deps.urlValidator) allowedByHostPolicy = await deps.urlValidator(pushUrl) + } catch (error) { + allowedByHostPolicy = false + console.error( + `[a2a] push URL policy failed for task ${task.id}:`, + error instanceof Error ? error.message : String(error), + ) + } + if (!allowedByHostPolicy) { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'pushNotificationConfig.url was rejected')) + } + await deps.pushStore.set(params.taskId, params.pushNotificationConfig) + const stored = await deps.pushStore.get(params.taskId, params.pushNotificationConfig.id) + return c.json(ok(req.id, { taskId: params.taskId, pushNotificationConfig: stored })) +} + +export async function handlePushGet( + c: Context, + req: JSONRPCRequest, + deps: PushConfigMethodDependencies, +): Promise { + if (!deps.pushStore) { + return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) + } + const params = req.params as TaskPushNotificationConfigGetParams | undefined + if (!params || typeof params.id !== 'string' || typeof params.pushNotificationConfigId !== 'string') { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id and params.pushNotificationConfigId required'), + ) + } + const task = await deps.taskStore.get(params.id) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, task) + if (accessError) return accessError + const cfg = await deps.pushStore.get(params.id, params.pushNotificationConfigId) + if (!cfg) { + return c.json( + fail( + req.id, + A2A_ERROR_CODES.TASK_NOT_FOUND, + `push config '${params.pushNotificationConfigId}' not found for task '${params.id}'`, + ), + ) + } + return c.json(ok(req.id, { taskId: params.id, pushNotificationConfig: cfg })) +} + +export async function handlePushList( + c: Context, + req: JSONRPCRequest, + deps: PushConfigMethodDependencies, +): Promise { + if (!deps.pushStore) { + return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) + } + const params = req.params as TaskIdParams | undefined + if (!params || typeof params.id !== 'string') { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) + } + const task = await deps.taskStore.get(params.id) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, task) + if (accessError) return accessError + const configs = await deps.pushStore.list(params.id) + return c.json(ok(req.id, configs.map((cfg) => ({ taskId: params.id, pushNotificationConfig: cfg })))) +} + +export async function handlePushDelete( + c: Context, + req: JSONRPCRequest, + deps: PushConfigMethodDependencies, +): Promise { + if (!deps.pushStore) { + return c.json(fail(req.id, A2A_ERROR_CODES.PUSH_NOT_SUPPORTED, 'push notifications not configured')) + } + const params = req.params as TaskPushNotificationConfigGetParams | undefined + if (!params || typeof params.id !== 'string' || typeof params.pushNotificationConfigId !== 'string') { + return c.json( + fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id and params.pushNotificationConfigId required'), + ) + } + const task = await deps.taskStore.get(params.id) + if (!task) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, task) + if (accessError) return accessError + await deps.pushStore.delete(params.id, params.pushNotificationConfigId) + return c.json(ok(req.id, null)) +} diff --git a/src/a2a/task-cancellation.ts b/src/a2a/task-cancellation.ts new file mode 100644 index 0000000..c872062 --- /dev/null +++ b/src/a2a/task-cancellation.ts @@ -0,0 +1,50 @@ +export class TaskCancellationRegistry { + private readonly controllers = new Map() + private readonly finalizing = new Set() + + register(taskId: string): AbortController { + const controller = new AbortController() + this.controllers.set(taskId, controller) + return controller + } + + clear(taskId: string): void { + this.controllers.delete(taskId) + this.finalizing.delete(taskId) + } + + beginFinalization(taskId: string): boolean { + const controller = this.controllers.get(taskId) + if (!controller || controller.signal.aborted || this.finalizing.has(taskId)) return false + this.finalizing.add(taskId) + return true + } + + isFinalizing(taskId: string): boolean { + return this.finalizing.has(taskId) + } + + has(taskId: string): boolean { + const controller = this.controllers.get(taskId) + return controller !== undefined && !controller.signal.aborted + } + + cancel(taskId: string): boolean { + if (this.finalizing.has(taskId)) return false + const controller = this.controllers.get(taskId) + if (!controller) return false + controller.abort() + this.controllers.delete(taskId) + return true + } +} + +export function bindRequestAbort( + requestSignal: AbortSignal, + controller: AbortController, +): () => void { + const abort = () => controller.abort() + if (requestSignal.aborted) abort() + else requestSignal.addEventListener('abort', abort, { once: true }) + return () => requestSignal.removeEventListener('abort', abort) +} diff --git a/src/a2a/task-finalization.ts b/src/a2a/task-finalization.ts new file mode 100644 index 0000000..45cb660 --- /dev/null +++ b/src/a2a/task-finalization.ts @@ -0,0 +1,451 @@ +import type { + AuthorizedRequest, + SettleAndRecordOptions, +} from '../dispatch' +import type { PaymentOperation, PaymentOperations } from '../payment-operations' +import { + deserializePaymentOperation, + serializePaymentOperation, + type PaymentRecoveryConfig, + type PaymentRecoveryRecord, + type SerializedPaymentOperation, +} from '../payment-recovery' +import type { AgentMeta, PaymentMethod, SandboxExecutionBudget, SandboxUsageReceipt } from '../types' +import { + clearPaymentRecoveryMarker, + readPaymentRecoveryMarker, +} from './payment-recovery' +import { clearTaskSubmission } from './task-submission-recovery' +import type { Task } from './types' +import { + agentMessage, + asError, + compareAndSetTask, + clearTaskMetadata, + cryptoRandomId, + isTerminal, + persistTaskIfCurrent, + type TaskStateStore, + withStatus, +} from './task-state' +import { responseTextToArtifact } from './translate' +import type { Artifact } from './types' + +const FINALIZING_METADATA_KEY = 'gatewayFinalizing' +const FINALIZATION_LEASE_MS = 5 * 60 * 1000 + +export type FinalizationState = 'completed' | 'input-required' | 'canceled' + +export interface FinalizationRecord { + version: 1 + lease: { id: string; expiresAt: number } + agentSlug: string + requestId: string + consumerId: string + paymentMethod: PaymentMethod + startMs: number + operationId: string | null + paymentOperation: SerializedPaymentOperation | null + receipt: SandboxUsageReceipt + artifact: Artifact | null + inputRequired: boolean + inputRequiredPrompt?: string + finalState?: FinalizationState + maxOutputTokens: number + executionBudget: SandboxExecutionBudget + usageRecorded: boolean + recoveryAttempts?: number + recoveryError?: string +} + +export interface TaskFinalizationDependencies { + taskStore: TaskStateStore + settle: ( + authz: AuthorizedRequest, + usage: SandboxUsageReceipt, + options?: SettleAndRecordOptions, + ) => Promise + resolveAgent: (slug: string) => Promise + paymentOperations?: PaymentOperations + paymentRecovery?: PaymentRecoveryConfig + recoverDurablePayment: ( + recoveryId: string, + options?: { force?: boolean; usage?: SandboxUsageReceipt }, + ) => Promise + releasePaymentAfterFailure: ( + authz: AuthorizedRequest, + reason: string, + workObserved: boolean, + ) => Promise + releaseTaskPayment: ( + authz: AuthorizedRequest, + task: Task, + reason: string, + workObserved: boolean, + ) => Promise + deliverPush: (task: Task) => Promise +} + +export function buildFinalizationRecord( + authz: AuthorizedRequest, + receipt: SandboxUsageReceipt, + artifact: Artifact | null, + inputRequired: boolean, + inputRequiredPrompt: string | undefined, + finalState: FinalizationState = inputRequired ? 'input-required' : 'completed', +): FinalizationRecord { + const operation = authz.paymentOperation + return { + version: 1, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, + agentSlug: authz.agent.slug, + requestId: authz.requestId, + consumerId: authz.consumerId, + paymentMethod: authz.paymentMethod, + startMs: authz.startMs, + operationId: operation?.operationId ?? null, + paymentOperation: operation ? serializePaymentOperation(operation) : null, + receipt, + artifact, + inputRequired, + ...(inputRequiredPrompt ? { inputRequiredPrompt } : {}), + finalState, + maxOutputTokens: authz.maxOutputTokens, + executionBudget: authz.executionBudget, + usageRecorded: false, + } +} + +export function withFinalizationRecord(task: Task, record: FinalizationRecord): Task { + return { + ...task, + metadata: { ...(task.metadata ?? {}), [FINALIZING_METADATA_KEY]: record }, + } +} + +export function readFinalizationRecord(task: Task): FinalizationRecord | undefined { + const raw = task.metadata?.[FINALIZING_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const record = raw as Partial + if ( + record.version !== 1 || + !record.lease || + typeof record.lease.id !== 'string' || + typeof record.lease.expiresAt !== 'number' + ) { + return undefined + } + return record as FinalizationRecord +} + +export function isTaskFinalizing(task: Task): boolean { + const marker = task.metadata?.[FINALIZING_METADATA_KEY] + return marker === true || (typeof marker === 'object' && marker !== null) +} + +export function clearFinalizationMarker(task: Task): Task { + return clearTaskMetadata(task, FINALIZING_METADATA_KEY) +} + +function markUsageRecordedRecord(task: Task): Task { + const record = readFinalizationRecord(task) + if (!record || record.usageRecorded) return task + return withFinalizationRecord(task, { ...record, usageRecorded: true }) +} + +export async function markUsageRecorded( + taskStore: TaskStateStore, + task: Task, +): Promise { + const marked = markUsageRecordedRecord(task) + if (marked === task) return task + if (await compareAndSetTask(taskStore, task, marked)) return marked + return await taskStore.get(task.id) ?? marked +} + +export async function retainFinalizationForRecovery( + taskStore: TaskStateStore, + taskId: string, + leaseId: string, + error: Error, +): Promise { + const current = await taskStore.get(taskId) + if (!current) return undefined + const record = readFinalizationRecord(current) + if (!record || record.lease.id !== leaseId) return undefined + const retry: FinalizationRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, + recoveryAttempts: (record.recoveryAttempts ?? 0) + 1, + recoveryError: error.message, + } + const next = withFinalizationRecord(current, retry) + if (await compareAndSetTask(taskStore, current, next)) return next + return await taskStore.get(taskId) +} + +export async function completeCanceledTask( + authz: AuthorizedRequest, + task: Task, + responseText: string, + usage: SandboxUsageReceipt | undefined, + workObserved: boolean, + deps: TaskFinalizationDependencies, +): Promise { + if (usage) { + const current = await deps.taskStore.get(task.id) ?? task + const finalization = buildFinalizationRecord( + authz, + usage, + responseText + ? responseTextToArtifact(responseText, `${task.id}-artifact-0`) + : current.artifacts?.[0] ?? null, + false, + undefined, + 'canceled', + ) + let finalizingTask: Task | undefined + let candidate = current + for (let attempt = 0; attempt < 8; attempt += 1) { + if (isTaskFinalizing(candidate)) return candidate + if (isTerminal(candidate.status.state) && candidate.status.state !== 'canceled') return candidate + const next = withFinalizationRecord(candidate, finalization) + if (await compareAndSetTask(deps.taskStore, candidate, next)) { + finalizingTask = next + break + } + const latest = await deps.taskStore.get(task.id) + if (!latest) break + if (isTerminal(latest.status.state) && latest.status.state !== 'canceled') return latest + candidate = latest + } + if (!finalizingTask) { + throw new Error(`A2A task '${task.id}' changed before cancellation settlement`) + } + + let usageRecordedTask = finalizingTask + try { + await deps.settle(authz, usage, { + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + } catch (settlementError) { + await deps.releasePaymentAfterFailure( + authz, + settlementError instanceof Error ? settlementError.message : String(settlementError), + true, + ) + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + finalization.lease.id, + asError(settlementError), + ) + const recoveryTask = retained ?? finalizingTask + console.error( + `[a2a] canceled task settlement retained for ${authz.requestId}:`, + settlementError instanceof Error ? settlementError.message : String(settlementError), + ) + await deps.deliverPush(recoveryTask) + return recoveryTask + } + + const canceled = withStatus( + clearPaymentRecoveryMarker(clearFinalizationMarker(usageRecordedTask)), + 'canceled', + undefined, + responseText + ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)] + : finalizingTask.artifacts, + ) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, canceled)) { + return await deps.taskStore.get(task.id) ?? canceled + } + await deps.deliverPush(canceled) + return canceled + } + await deps.releaseTaskPayment(authz, task, 'a2a task canceled', workObserved) + const currentTask = await deps.taskStore.get(task.id) + const canceledBase = currentTask?.status.state === 'canceled' + ? currentTask + : withStatus(currentTask ?? task, 'canceled') + const canceled: Task = responseText + ? { + ...canceledBase, + artifacts: [responseTextToArtifact(responseText, `${task.id}-artifact-0`)], + } + : canceledBase + const persisted = await persistTaskIfCurrent(deps.taskStore, currentTask ?? task, canceled) + await deps.deliverPush(persisted) + return persisted +} + +export async function recoverFinalizationIfNeeded( + task: Task, + deps: TaskFinalizationDependencies, + requestedAgentSlug: string, +): Promise { + if (!isTaskFinalizing(task)) return task + const record = readFinalizationRecord(task) + if (!record) { + return expireFinalization(task, deps, null, new Error('A2A finalization record is missing')) + } + if (record.lease.expiresAt > Date.now()) return task + + const renewed: FinalizationRecord = { + ...record, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + FINALIZATION_LEASE_MS }, + } + const leasedTask = withFinalizationRecord(task, renewed) + if (!await compareAndSetTask(deps.taskStore, task, leasedTask)) { + return await deps.taskStore.get(task.id) ?? task + } + + try { + const agentSlug = renewed.agentSlug || requestedAgentSlug + const agent = await deps.resolveAgent(agentSlug) + if (!agent || !agent.enabled) throw new Error('A2A recovery agent is unavailable') + + const paymentRecovery = readPaymentRecoveryMarker(leasedTask) + if (paymentRecovery && deps.paymentRecovery) { + const recovery = await deps.recoverDurablePayment(paymentRecovery.id, { + force: true, + usage: renewed.receipt, + }) + if (recovery?.state !== 'reconciled') { + throw new Error('durable payment finalization is still pending') + } + const usageRecordedTask = await markUsageRecorded(deps.taskStore, leasedTask) + const recoveredTask = finalizationResultTask(usageRecordedTask, renewed) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, recoveredTask)) { + return await deps.taskStore.get(task.id) ?? recoveredTask + } + await deps.deliverPush(recoveredTask) + return recoveredTask + } + + let paymentOperation: PaymentOperation | undefined + if (renewed.operationId || renewed.paymentOperation) { + if (!renewed.operationId || !renewed.paymentOperation) { + throw new Error('A2A payment operation recovery record is incomplete') + } + if (renewed.operationId !== renewed.paymentOperation.operationId) { + throw new Error('A2A payment operation recovery id does not match') + } + if (!deps.paymentOperations) { + throw new Error('A2A payment operation recovery is not configured') + } + paymentOperation = deserializePaymentOperation(renewed.paymentOperation) + } + + let paymentAlreadySettled = false + if (paymentOperation && deps.paymentOperations) { + const currentOperation = await deps.paymentOperations.getPaymentOperation(paymentOperation.operationId) + if (currentOperation.state === 'not-found') { + throw new Error('A2A payment operation disappeared during finalization recovery') + } + if (currentOperation.operationId !== paymentOperation.operationId) { + throw new Error('A2A payment operation recovery returned a different operation') + } + paymentOperation = currentOperation + paymentAlreadySettled = currentOperation.state === 'settled' + } + + const authz: AuthorizedRequest = { + agent, + consumerId: renewed.consumerId, + paymentMethod: renewed.paymentMethod, + keyInfo: null, + userMessage: '[recovered A2A task]', + rateLimitRemaining: undefined, + requestId: renewed.requestId, + startMs: renewed.startMs, + maxOutputTokens: renewed.maxOutputTokens, + executionBudget: renewed.executionBudget, + requiredPaymentAmount: 0n, + paymentPayload: null, + ...(paymentRecovery ? { paymentRecoveryId: paymentRecovery.id } : {}), + ...(paymentOperation ? { paymentOperation, paymentOperationAcquired: true } : {}), + } + let usageRecordedTask = leasedTask + await deps.settle(authz, renewed.receipt, { + usageAlreadyRecorded: renewed.usageRecorded === true, + paymentAlreadySettled, + onUsageRecorded: async () => { + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + }, + }) + usageRecordedTask = await markUsageRecorded(deps.taskStore, usageRecordedTask) + const recovered = finalizationResultTask(usageRecordedTask, renewed) + if (!await compareAndSetTask(deps.taskStore, usageRecordedTask, recovered)) { + return await deps.taskStore.get(task.id) ?? recovered + } + await deps.deliverPush(recovered) + return recovered + } catch (error) { + const recoveryError = asError(error) + console.error(`[a2a] finalization recovery failed for ${task.id}:`, recoveryError.message) + if ( + (readPaymentRecoveryMarker(leasedTask) && deps.paymentRecovery) || + (renewed.operationId && renewed.paymentOperation) + ) { + const retained = await retainFinalizationForRecovery( + deps.taskStore, + task.id, + renewed.lease.id, + recoveryError, + ) + if (retained) return retained + } + return expireFinalization(leasedTask, deps, renewed, recoveryError) + } +} + +function finalizationResultTask(task: Task, record: FinalizationRecord): Task { + const cleanTask = clearPaymentRecoveryMarker(clearFinalizationMarker(task)) + const finalState = record.finalState ?? ( + task.status.state === 'canceled' + ? 'canceled' + : record.inputRequired + ? 'input-required' + : 'completed' + ) + if (finalState === 'canceled') { + return withStatus(cleanTask, 'canceled', undefined, record.artifact ? [record.artifact] : cleanTask.artifacts) + } + if (finalState === 'input-required') { + return withStatus( + cleanTask, + 'input-required', + record.inputRequiredPrompt ? agentMessage(cleanTask, record.inputRequiredPrompt) : undefined, + record.artifact ? [record.artifact] : cleanTask.artifacts, + ) + } + return withStatus(cleanTask, 'completed', undefined, record.artifact ? [record.artifact] : cleanTask.artifacts) +} + +async function expireFinalization( + task: Task, + deps: TaskFinalizationDependencies, + record: FinalizationRecord | null, + error: Error, +): Promise { + const cleanTask = clearFinalizationMarker(task) + const failed: Task = { + ...withStatus(cleanTask, 'failed'), + metadata: { + ...(cleanTask.metadata ?? {}), + gatewayFinalizationRecovery: { + operationId: record?.operationId ?? null, + error: error.message, + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await deps.deliverPush(failed) + return failed + } + return await deps.taskStore.get(task.id) ?? failed +} diff --git a/src/a2a/task-lifecycle.ts b/src/a2a/task-lifecycle.ts new file mode 100644 index 0000000..925b3f4 --- /dev/null +++ b/src/a2a/task-lifecycle.ts @@ -0,0 +1,54 @@ +import { + type GatewayState, + releasePayment, + releasePaymentAfterFailure, + settleAndRecord, +} from '../dispatch' +import { recoverPayment as recoverDurablePayment } from '../payment-recovery-worker' +import type { GatewayConfig } from '../types' +import { releaseTaskPayment } from './payment-recovery' +import type { PaymentRecoveryDependencies } from './payment-recovery' +import type { TaskFinalizationDependencies } from './task-finalization' +import type { Task } from './types' +import type { TaskStateStore } from './task-state' + +export interface TaskLifecycleDependencies { + taskStore: TaskStateStore + config: GatewayConfig + state: GatewayState + deliverPush: (task: Task) => Promise +} + +export interface TaskLifecycle { + payment: PaymentRecoveryDependencies + finalization: TaskFinalizationDependencies +} + +export function createTaskLifecycle(deps: TaskLifecycleDependencies): TaskLifecycle { + const payment: PaymentRecoveryDependencies = { + taskStore: deps.taskStore, + paymentOperations: deps.config.x402.paymentOperations, + paymentRecovery: deps.config.paymentRecovery, + releasePayment: (authz, reason) => releasePayment(authz, deps.config, reason), + releasePaymentAfterFailure: (authz, reason, workObserved) => + releasePaymentAfterFailure(authz, deps.config, reason, workObserved), + recoverDurablePayment: (recoveryId, options) => + recoverDurablePayment(recoveryId, deps.config, options), + deliverPush: deps.deliverPush, + } + const finalization: TaskFinalizationDependencies = { + taskStore: deps.taskStore, + settle: (authz, usage, options) => + settleAndRecord(authz.agent, authz, usage, deps.config, deps.state.obs, options), + resolveAgent: deps.config.resolveAgent, + paymentOperations: deps.config.x402.paymentOperations, + paymentRecovery: deps.config.paymentRecovery, + recoverDurablePayment: (recoveryId, options) => + recoverDurablePayment(recoveryId, deps.config, options), + releasePaymentAfterFailure: payment.releasePaymentAfterFailure, + releaseTaskPayment: (authz, task, reason, workObserved) => + releaseTaskPayment(authz, task, payment, reason, workObserved), + deliverPush: deps.deliverPush, + } + return { payment, finalization } +} diff --git a/src/a2a/task-methods.ts b/src/a2a/task-methods.ts new file mode 100644 index 0000000..5a5c0f6 --- /dev/null +++ b/src/a2a/task-methods.ts @@ -0,0 +1,163 @@ +import type { Context } from 'hono' +import { hasActiveTaskExecution } from './execution-fence' +import { fail, ok } from './jsonrpc' +import { releaseTaskPayment } from './payment-recovery' +import type { PaymentRecoveryDependencies } from './payment-recovery' +import { isTaskFinalizing } from './task-finalization' +import type { TaskCancellationRegistry } from './task-cancellation' +import { + compareAndSetTask, + isTerminal, + withStatus, +} from './task-state' +import type { TaskStateStore } from './task-state' +import { + A2A_ERROR_CODES, + type JSONRPCRequest, + type Task, + type TaskIdParams, + type TaskStatusUpdateEvent, +} from './types' + +export interface TaskMethodDependencies { + taskStore: TaskStateStore + payment: PaymentRecoveryDependencies + cancels: TaskCancellationRegistry + authorizeTaskAccess: ( + c: Context, + req: JSONRPCRequest, + task: Task, + ) => Promise + recoverTask: (task: Task, requestedAgentSlug: string) => Promise + deliverPush: (task: Task) => Promise +} + +export async function handleTasksGet( + c: Context, + req: JSONRPCRequest, + deps: TaskMethodDependencies, +): Promise { + const params = req.params as TaskIdParams | undefined + if (!params || typeof params.id !== 'string') { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) + } + const storedTask = await deps.taskStore.get(params.id) + if (!storedTask) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, storedTask) + if (accessError) return accessError + const task = await deps.recoverTask(storedTask, c.req.param('slug') ?? '') + return c.json(ok(req.id, task)) +} + +export async function handleTasksCancel( + c: Context, + req: JSONRPCRequest, + deps: TaskMethodDependencies, +): Promise { + const params = req.params as TaskIdParams | undefined + if (!params || typeof params.id !== 'string') { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) + } + const storedTask = await deps.taskStore.get(params.id) + if (!storedTask) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, storedTask) + if (accessError) return accessError + const task = await deps.recoverTask(storedTask, c.req.param('slug') ?? '') + if (isTerminal(task.status.state)) { + return c.json( + fail( + req.id, + A2A_ERROR_CODES.TASK_NOT_CANCELABLE, + `task '${params.id}' is in terminal state '${task.status.state}'`, + ), + ) + } + if ( + isTaskFinalizing(task) || + deps.cancels.isFinalizing(task.id) || + (hasActiveTaskExecution(task) && !deps.cancels.has(task.id)) + ) { + return c.json( + fail( + req.id, + A2A_ERROR_CODES.TASK_NOT_CANCELABLE, + hasActiveTaskExecution(task) + ? `task '${task.id}' has an active execution fence` + : `task '${task.id}' is being finalized`, + ), + ) + } + let candidate = task + for (let attempt = 0; attempt < 8; attempt += 1) { + if (isTerminal(candidate.status.state)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' changed before cancellation`), + ) + } + if (isTaskFinalizing(candidate)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' is being finalized`), + ) + } + if (hasActiveTaskExecution(candidate) && !deps.cancels.has(candidate.id)) { + return c.json( + fail(req.id, A2A_ERROR_CODES.TASK_NOT_CANCELABLE, `task '${task.id}' has an active execution fence`), + ) + } + const canceled = withStatus(candidate, 'canceled') + if (await compareAndSetTask(deps.taskStore, candidate, canceled)) { + const stillActive = deps.cancels.cancel(task.id) + if (!stillActive) await deps.deliverPush(canceled) + return c.json(ok(req.id, canceled)) + } + const current = await deps.taskStore.get(task.id) + if (!current) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${task.id}' not found`)) + } + candidate = current + } + return c.json(fail(req.id, A2A_ERROR_CODES.INTERNAL_ERROR, 'task changed before cancellation')) +} + +export async function handleTasksResubscribe( + c: Context, + req: JSONRPCRequest, + deps: TaskMethodDependencies, +): Promise { + const params = req.params as TaskIdParams | undefined + if (!params || typeof params.id !== 'string') { + return c.json(fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, 'params.id required')) + } + const storedTask = await deps.taskStore.get(params.id) + if (!storedTask) { + return c.json(fail(req.id, A2A_ERROR_CODES.TASK_NOT_FOUND, `task '${params.id}' not found`)) + } + const accessError = await deps.authorizeTaskAccess(c, req, storedTask) + if (accessError) return accessError + const task = await deps.recoverTask(storedTask, c.req.param('slug') ?? '') + const event: TaskStatusUpdateEvent = { + kind: 'status-update', + taskId: task.id, + contextId: task.contextId, + status: task.status, + final: isTerminal(task.status.state) || task.status.state === 'input-required', + } + const encoder = new TextEncoder() + const stream = new ReadableStream({ + start(ctrl) { + ctrl.enqueue(encoder.encode(`data: ${JSON.stringify(ok(req.id, event))}\n\n`)) + ctrl.close() + }, + }) + return new Response(stream, { + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + 'X-Task-Id': task.id, + }, + }) +} diff --git a/src/a2a/task-push-delivery.ts b/src/a2a/task-push-delivery.ts new file mode 100644 index 0000000..de98b33 --- /dev/null +++ b/src/a2a/task-push-delivery.ts @@ -0,0 +1,119 @@ +import { + deliverDemoPushNotifications, + deliverPushNotifications, + type PushDeliveryResult, + type PushNotificationDeliveryOptions, + type PushNotificationStore, +} from './push-notifications' +import type { Task } from './types' +import { + compareAndSetTask, + clearTaskMetadata, + TERMINAL_STATES, + type TaskStateStore, +} from './task-state' + +const PUSH_DELIVERY_METADATA_KEY = 'gatewayPushDelivery' + +interface TaskPushDeliveryClaims { + version: 1 + claims: Record +} + +export interface PushDeliveryDependencies { + taskStore: TaskStateStore + pushStore?: PushNotificationStore + demoMode: boolean + webhookSecret?: string + fetcher?: PushNotificationDeliveryOptions['fetcher'] + urlValidator?: PushNotificationDeliveryOptions['urlValidator'] + onDeliveryFailure?: (task: Task, result: PushDeliveryResult) => void +} + +export async function deliverTaskPush( + task: Task, + deps: PushDeliveryDependencies, +): Promise { + if (!deps.pushStore || !TERMINAL_STATES.has(task.status.state)) return + const webhookSecret = deps.webhookSecret + const hasWebhookSecret = typeof webhookSecret === 'string' && webhookSecret.trim().length > 0 + if (!deps.demoMode && !hasWebhookSecret) { + console.error(`[agent-gateway] production A2A push requires a webhookSecret for task ${task.id}`) + return + } + try { + const deliveryTask = clearPushDeliveryClaims(task) + const deliveryArgs: Omit = { + task: deliveryTask, + store: deps.pushStore, + fetcher: deps.fetcher, + urlValidator: deps.urlValidator, + requireUrlValidator: !deps.demoMode, + claimDelivery: (taskId, configId, terminalState) => claimTaskPushDelivery( + deps.taskStore, + taskId, + configId, + terminalState, + ), + onDelivery: (result) => { + if (!result.ok) deps.onDeliveryFailure?.(task, result) + }, + } + if (hasWebhookSecret) { + await deliverPushNotifications({ ...deliveryArgs, webhookSecret }) + } else { + await deliverDemoPushNotifications(deliveryArgs) + } + } catch (err) { + console.error( + `[agent-gateway] push delivery threw for task ${task.id}: ${err instanceof Error ? err.message : String(err)}`, + ) + } +} + +async function claimTaskPushDelivery( + taskStore: TaskStateStore, + taskId: string, + configId: string, + terminalState: Task['status']['state'], +): Promise { + if (!TERMINAL_STATES.has(terminalState)) return false + for (let attempt = 0; attempt < 16; attempt += 1) { + const current = await taskStore.get(taskId) + if (!current || current.status.state !== terminalState) return false + const existing = readPushDeliveryClaims(current) + if (existing?.claims[configId] === terminalState) return false + const next: Task = { + ...current, + metadata: { + ...(current.metadata ?? {}), + [PUSH_DELIVERY_METADATA_KEY]: { + version: 1, + claims: { + ...(existing?.claims ?? {}), + [configId]: terminalState, + }, + } satisfies TaskPushDeliveryClaims, + }, + } + if (await compareAndSetTask(taskStore, current, next)) return true + } + throw new Error(`A2A push delivery claim changed too many times for task '${taskId}'`) +} + +function clearPushDeliveryClaims(task: Task): Task { + return clearTaskMetadata(task, PUSH_DELIVERY_METADATA_KEY) +} + +function readPushDeliveryClaims(task: Task): TaskPushDeliveryClaims | undefined { + const raw = task.metadata?.[PUSH_DELIVERY_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const record = raw as Partial + if (!record.claims || typeof record.claims !== 'object') return undefined + const claims = Object.fromEntries( + Object.entries(record.claims).filter(([, state]) => + typeof state === 'string' && TERMINAL_STATES.has(state as Task['status']['state']), + ), + ) as Record + return record.version === 1 ? { version: 1, claims } : undefined +} diff --git a/src/a2a/task-recovery.ts b/src/a2a/task-recovery.ts new file mode 100644 index 0000000..bff3135 --- /dev/null +++ b/src/a2a/task-recovery.ts @@ -0,0 +1,11 @@ +import type { Task } from './types' + +const PAYMENT_RECOVERY_KEYS = [ + 'gatewayFinalizing', + 'gatewayPaymentRelease', + 'gatewayPaymentRecovery', +] as const + +export function hasPendingPaymentRecovery(task: Task): boolean { + return PAYMENT_RECOVERY_KEYS.some((key) => task.metadata?.[key] !== undefined) +} diff --git a/src/a2a/task-state.ts b/src/a2a/task-state.ts new file mode 100644 index 0000000..3126a61 --- /dev/null +++ b/src/a2a/task-state.ts @@ -0,0 +1,99 @@ +import type { Task, TaskStatus, Message } from './types' +import { clearTaskExecution } from './execution-fence' +import { hasPendingPaymentRecovery } from './task-recovery' + +export interface TaskStateStore { + get(id: string): Promise + compareAndSet?(expected: Task, next: Task): Promise +} + +export const TERMINAL_STATES: ReadonlySet = new Set([ + 'completed', + 'canceled', + 'failed', + 'rejected', +]) + +export function isTerminal(state: Task['status']['state']): boolean { + return TERMINAL_STATES.has(state) +} + +export function shouldPreserveTask(task: Task): boolean { + return isTerminal(task.status.state) || hasPendingPaymentRecovery(task) +} + +export async function compareAndSetTask( + taskStore: TaskStateStore, + expected: Task, + next: Task, +): Promise { + if (!taskStore.compareAndSet) { + throw new Error('A2A task store does not provide compareAndSet') + } + return taskStore.compareAndSet(expected, next) +} + +export async function persistTaskIfCurrent( + taskStore: TaskStateStore, + expected: Task, + next: Task, +): Promise { + if (expected === next || JSON.stringify(expected) === JSON.stringify(next)) return expected + if (await compareAndSetTask(taskStore, expected, next)) return next + return await taskStore.get(expected.id) ?? expected +} + +export function clearTaskMetadata(task: Task, key: string): Task { + if (!task.metadata || !(key in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[key] + if (Object.keys(metadata).length > 0) return { ...task, metadata } + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata +} + +export function withStatus( + task: Task, + state: TaskStatus['state'], + message?: Message, + artifacts?: Task['artifacts'], +): Task { + const next: Task = { + ...task, + status: { state, timestamp: nowIso(), ...(message ? { message } : {}) }, + ...(artifacts !== undefined ? { artifacts } : {}), + } + return isTerminal(state) || state === 'input-required' ? clearTaskExecution(next) : next +} + +export function agentMessage(task: Task, text: string): Message { + return { + kind: 'message', + role: 'agent', + parts: [{ kind: 'text', text }], + messageId: `${task.id}-input-required-${stableMessageDigest(text)}`, + taskId: task.id, + contextId: task.contextId, + } +} + +export function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)) +} + +export function nowIso(): string { + return new Date().toISOString() +} + +export function cryptoRandomId(): string { + return crypto.randomUUID().replace(/-/g, '') +} + +function stableMessageDigest(value: string): string { + let hash = 2166136261 + for (let index = 0; index < value.length; index += 1) { + hash ^= value.charCodeAt(index) + hash = Math.imul(hash, 16777619) + } + return (hash >>> 0).toString(16).padStart(8, '0') +} diff --git a/src/a2a/task-store.ts b/src/a2a/task-store.ts index 7e6d219..949d9b8 100644 --- a/src/a2a/task-store.ts +++ b/src/a2a/task-store.ts @@ -7,6 +7,9 @@ import type { Task } from './types' import { inspectTaskExecution } from './execution-fence' +import { hasPendingPaymentRecovery } from './task-recovery' + +export { hasPendingPaymentRecovery } from './task-recovery' export interface TaskStore { get(id: string): Promise @@ -27,17 +30,6 @@ export interface TaskStore { const DEFAULT_TTL_MS = 60 * 60 * 1000 -const PAYMENT_RECOVERY_KEYS = [ - 'gatewayFinalizing', - 'gatewayPaymentRelease', - 'gatewayPaymentRecovery', -] as const - -/** Payment recovery tasks must remain readable until reconciliation clears the marker. */ -export function hasPendingPaymentRecovery(task: Task): boolean { - return PAYMENT_RECOVERY_KEYS.some((key) => task.metadata?.[key] !== undefined) -} - export class InMemoryTaskStore implements TaskStore { private readonly entries = new Map() diff --git a/src/a2a/task-submission-recovery.ts b/src/a2a/task-submission-recovery.ts new file mode 100644 index 0000000..3fadf32 --- /dev/null +++ b/src/a2a/task-submission-recovery.ts @@ -0,0 +1,178 @@ +import type { Task } from './types' +import { + compareAndSetTask, + cryptoRandomId, + type TaskStateStore, + withStatus, +} from './task-state' + +const TASK_ORIGIN_METADATA_KEY = 'gatewayOrigin' +const TASK_SUBMISSION_METADATA_KEY = 'gatewaySubmission' +const TASK_SUBMISSION_RECOVERY_METADATA_KEY = 'gatewaySubmissionRecovery' +const TASK_SUBMISSION_LEASE_MS = 5 * 60 * 1000 + +export interface TaskOriginAgent { + id: string + slug: string +} + +export interface TaskSubmissionIdentity { + agent: TaskOriginAgent + requestId: string + consumerId: string +} + +interface TaskOriginBinding { + version: 1 + agentId: string + agentSlug: string +} + +export interface TaskSubmissionRecord { + version: 1 + lease: { id: string; expiresAt: number } + agentId: string + agentSlug: string + requestId: string + consumerId: string +} + +export interface SubmissionRecoveryDependencies { + taskStore: TaskStateStore + deliverPush: (task: Task) => Promise +} + +export function withTaskOrigin( + metadata: Record | undefined, + agent: TaskOriginAgent, +): Record { + return { + ...(metadata ?? {}), + [TASK_ORIGIN_METADATA_KEY]: { + version: 1, + agentId: agent.id, + agentSlug: agent.slug, + } satisfies TaskOriginBinding, + } +} + +export function withTaskSubmission( + metadata: Record | undefined, + identity: TaskSubmissionIdentity, +): Record { + const origin = metadata?.[TASK_ORIGIN_METADATA_KEY] + return { + ...(metadata ?? {}), + ...(origin === undefined + ? { + [TASK_ORIGIN_METADATA_KEY]: { + version: 1, + agentId: identity.agent.id, + agentSlug: identity.agent.slug, + } satisfies TaskOriginBinding, + } + : {}), + [TASK_SUBMISSION_METADATA_KEY]: { + version: 1, + lease: { id: cryptoRandomId(), expiresAt: Date.now() + TASK_SUBMISSION_LEASE_MS }, + agentId: identity.agent.id, + agentSlug: identity.agent.slug, + requestId: identity.requestId, + consumerId: identity.consumerId, + } satisfies TaskSubmissionRecord, + } +} + +export function readTaskOrigin(task: Task): TaskOriginBinding | undefined { + const raw = task.metadata?.[TASK_ORIGIN_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const origin = raw as Partial + if ( + origin.version !== 1 || + typeof origin.agentId !== 'string' || + origin.agentId.length === 0 || + typeof origin.agentSlug !== 'string' || + origin.agentSlug.length === 0 + ) { + return undefined + } + return origin as TaskOriginBinding +} + +export function readTaskSubmission(task: Task): TaskSubmissionRecord | undefined { + const raw = task.metadata?.[TASK_SUBMISSION_METADATA_KEY] + if (!raw || typeof raw !== 'object') return undefined + const submission = raw as Partial + if ( + submission.version !== 1 || + !submission.lease || + typeof submission.lease.id !== 'string' || + submission.lease.id.length === 0 || + typeof submission.lease.expiresAt !== 'number' || + !Number.isFinite(submission.lease.expiresAt) || + typeof submission.agentId !== 'string' || + submission.agentId.length === 0 || + typeof submission.agentSlug !== 'string' || + submission.agentSlug.length === 0 || + typeof submission.requestId !== 'string' || + submission.requestId.length === 0 || + typeof submission.consumerId !== 'string' + ) { + return undefined + } + return submission as TaskSubmissionRecord +} + +export function clearTaskSubmission(task: Task): Task { + if (!task.metadata || !(TASK_SUBMISSION_METADATA_KEY in task.metadata)) return task + const metadata = { ...task.metadata } + delete metadata[TASK_SUBMISSION_METADATA_KEY] + if (Object.keys(metadata).length > 0) return { ...task, metadata } + const { metadata: _metadata, ...withoutMetadata } = task + return withoutMetadata +} + +export async function recoverSubmissionIfNeeded( + task: Task, + deps: SubmissionRecoveryDependencies, +): Promise { + const raw = task.metadata?.[TASK_SUBMISSION_METADATA_KEY] + if (raw === undefined) return task + const submission = readTaskSubmission(task) + if (submission && submission.lease.expiresAt > Date.now()) return task + if (task.status.state !== 'submitted') { + return (await clearTaskSubmissionMarker(deps.taskStore, task)).task + } + + const cleanTask = clearTaskSubmission(task) + const failed: Task = { + ...withStatus(cleanTask, 'failed'), + metadata: { + ...(cleanTask.metadata ?? {}), + [TASK_SUBMISSION_RECOVERY_METADATA_KEY]: { + error: submission + ? 'A2A task submission lease expired before payment authorization completed' + : 'A2A task submission lease is invalid', + }, + }, + } + if (await compareAndSetTask(deps.taskStore, task, failed)) { + await deps.deliverPush(failed) + return failed + } + return await deps.taskStore.get(task.id) ?? task +} + +async function clearTaskSubmissionMarker( + taskStore: TaskStateStore, + expected: Task, +): Promise<{ task: Task; applied: boolean }> { + const current = await taskStore.get(expected.id) + if (!current || JSON.stringify(current) !== JSON.stringify(expected)) { + return { task: current ?? expected, applied: false } + } + const cleared = clearTaskSubmission(current) + if (cleared === current) return { task: current, applied: true } + if (await compareAndSetTask(taskStore, current, cleared)) return { task: cleared, applied: true } + return { task: await taskStore.get(expected.id) ?? expected, applied: false } +} diff --git a/src/observer-types.ts b/src/observer-types.ts new file mode 100644 index 0000000..1865c59 --- /dev/null +++ b/src/observer-types.ts @@ -0,0 +1,63 @@ +import type { + GatewayUsageEvent, + PaymentMethod, +} from './payment-types' + +export interface RequestContext { + requestId: string + agentSlug: string + startMs: number +} + +export interface AuthFailureReason { + method: PaymentMethod + code: string + httpStatus: number +} + +export interface GatewayObserver { + /** Called at the start of every chat completions POST. */ + onRequestStart?: (ctx: RequestContext) => void | Promise + + /** Called when a payment method has been successfully verified. */ + onPaymentVerified?: (ctx: RequestContext, info: { + method: PaymentMethod + consumerId: string + keyId?: string + }) => void | Promise + + /** Called when auth fails — every branch. */ + onAuthFailure?: (ctx: RequestContext, reason: AuthFailureReason) => void | Promise + + /** Called when a consumer hits the rate limit. */ + onRateLimited?: (ctx: RequestContext, info: { + consumerId: string + retryAfterSeconds: number + }) => void | Promise + + /** Called when the request body exceeds the 64KB limit. */ + onBodyTooLarge?: (ctx: RequestContext, contentLength: number) => void | Promise + + /** Called when prompt-injection patterns are detected. */ + onInjectionDetected?: (ctx: RequestContext, info: { + consumerId: string + patterns: string[] + blocked: boolean + }) => void | Promise + + /** Called after a successful stream completes and recordUsage has fired. */ + onRequestComplete?: (ctx: RequestContext, usage: GatewayUsageEvent) => void | Promise + + /** Called when the sandbox throws. The error message is pre-scrubbed. */ + onStreamError?: (ctx: RequestContext, info: { + consumerId: string + errorMessage: string + }) => void | Promise + + /** Called when settlement fails. Payment already occurred; this is async bookkeeping. */ + onSettlementError?: (ctx: RequestContext, info: { + consumerId: string + method: PaymentMethod + errorMessage: string + }) => void | Promise +} diff --git a/src/observer.ts b/src/observer.ts index 358c7e1..ad2e4a9 100644 --- a/src/observer.ts +++ b/src/observer.ts @@ -10,70 +10,10 @@ * When no observer is configured, the gateway stays silent. */ -import type { PaymentMethod, GatewayUsageEvent } from './types' +import type { GatewayUsageEvent, PaymentMethod } from './payment-types' +import type { AuthFailureReason, GatewayObserver, RequestContext } from './observer-types' -export interface RequestContext { - requestId: string - agentSlug: string - startMs: number -} - -export interface AuthFailureReason { - method: 'x402' | 'mpp' | 'apikey' | 'none' - code: string - httpStatus: number -} - -export interface GatewayObserver { - /** Called at the start of every chat completions POST. */ - onRequestStart?: (ctx: RequestContext) => void | Promise - - /** Called when a payment method has been successfully verified. */ - onPaymentVerified?: (ctx: RequestContext, info: { - method: PaymentMethod - consumerId: string - keyId?: string - }) => void | Promise - - /** Called when auth fails — every branch. */ - onAuthFailure?: (ctx: RequestContext, reason: AuthFailureReason) => void | Promise - - /** Called when a consumer hits the rate limit. */ - onRateLimited?: (ctx: RequestContext, info: { - consumerId: string - retryAfterSeconds: number - }) => void | Promise - - /** Called when the request body exceeds the 64KB limit. */ - onBodyTooLarge?: (ctx: RequestContext, contentLength: number) => void | Promise - - /** - * Called when prompt-injection patterns are detected. - * `blocked` is true when blockInjection config is on and the request was - * rejected; false when the patterns were logged but the request proceeded. - */ - onInjectionDetected?: (ctx: RequestContext, info: { - consumerId: string - patterns: string[] - blocked: boolean - }) => void | Promise - - /** Called after a successful stream completes and recordUsage has fired. */ - onRequestComplete?: (ctx: RequestContext, usage: GatewayUsageEvent) => void | Promise - - /** Called when the sandbox throws. The error message is pre-scrubbed. */ - onStreamError?: (ctx: RequestContext, info: { - consumerId: string - errorMessage: string - }) => void | Promise - - /** Called when settlement fails. Payment already occurred; this is async bookkeeping. */ - onSettlementError?: (ctx: RequestContext, info: { - consumerId: string - method: PaymentMethod - errorMessage: string - }) => void | Promise -} +export type { AuthFailureReason, GatewayObserver, RequestContext } from './observer-types' // --------------------------------------------------------------------------- // Convenience implementations diff --git a/src/payment-operations.ts b/src/payment-operations.ts index bc94b88..a30f70f 100644 --- a/src/payment-operations.ts +++ b/src/payment-operations.ts @@ -1,5 +1,7 @@ -import type { SandboxUsageReceipt } from './types' -import type { PaymentSettlementBasis } from './payment-recovery' +import type { + PaymentSettlementBasis, + SandboxUsageReceipt, +} from './payment-types' /** Version negotiated by gateways that use durable payment operations. */ export const PAYMENT_PROTOCOL_VERSION = 2 as const diff --git a/src/payment-recovery.ts b/src/payment-recovery.ts index 1f54b93..8948700 100644 --- a/src/payment-recovery.ts +++ b/src/payment-recovery.ts @@ -2,9 +2,12 @@ import type { MppChargeOperation } from './mpp-payment' import type { PaymentOperation } from './payment-operations' import type { PaymentMethod, + PaymentSettlementBasis, SandboxExecutionBudget, SandboxUsageReceipt, -} from './types' +} from './payment-types' + +export type { PaymentSettlementBasis } from './payment-types' export const PAYMENT_RECOVERY_VERSION = 1 as const @@ -17,8 +20,6 @@ export type PaymentRecoveryState = | 'releasing' | 'reconciled' -export type PaymentSettlementBasis = 'usage-receipt' | 'quoted-ceiling' - export interface SerializedPaymentOperation { protocolVersion: 2 operationId: string diff --git a/src/payment-types.ts b/src/payment-types.ts new file mode 100644 index 0000000..9ecb021 --- /dev/null +++ b/src/payment-types.ts @@ -0,0 +1,48 @@ +export type PaymentMethod = 'x402' | 'mpp' | 'apikey' | 'none' + +export interface SandboxExecutionBudget { + maxInputTokens: number + maxOutputTokens: number + maxReasoningTokens: number + maxToolTokens: number + maxToolCalls: number + maxProviderCostUsd: number +} + +export interface SandboxUsageReceipt { + inputTokens: number + outputTokens: number + reasoningTokens: number + toolTokens: number + toolCallCount: number + providerCostUsd: number + /** True only when the provider/adapter enforced every supplied budget. */ + budgetEnforced: boolean +} + +export type PaymentSettlementBasis = 'usage-receipt' | 'quoted-ceiling' + +export interface GatewayUsageEvent { + /** Correlates usage, settlement, and observer records for one request. */ + requestId: string + agentId: string + agentSlug: string + consumerId: string + paymentMethod: PaymentMethod + inputTokens: number + outputTokens: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + reasoningTokens?: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + toolTokens?: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + toolCallCount?: number + /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ + providerCostUsd?: number + totalCostUsd: number + ownerEarnedUsd: number + platformFeeUsd: number + durationMs: number + /** Exact receipt in normal operation; quoted ceiling only after receipt timeout. */ + settlementBasis?: PaymentSettlementBasis +} diff --git a/src/types.ts b/src/types.ts index b619ac0..8916f55 100644 --- a/src/types.ts +++ b/src/types.ts @@ -4,7 +4,21 @@ import type { PaymentOperations, } from './payment-operations' import type { MppAuthenticatedCredential, MppChargeLifecycle } from './mpp-payment' -import type { PaymentRecoveryConfig, PaymentSettlementBasis } from './payment-recovery' +import type { PaymentRecoveryConfig } from './payment-recovery' +import type { GatewayObserver } from './observer-types' +import type { + GatewayUsageEvent, + PaymentMethod, + SandboxExecutionBudget, + SandboxUsageReceipt, +} from './payment-types' + +export type { + GatewayUsageEvent, + PaymentMethod, + SandboxExecutionBudget, + SandboxUsageReceipt, +} from './payment-types' // --- Agent resolution --- @@ -79,17 +93,6 @@ export interface AgentMeta { // --- Payment --- -export type PaymentMethod = 'x402' | 'mpp' | 'apikey' | 'none' - -export interface SandboxExecutionBudget { - maxInputTokens: number - maxOutputTokens: number - maxReasoningTokens: number - maxToolTokens: number - maxToolCalls: number - maxProviderCostUsd: number -} - export interface X402Config { /** Ethereum operator address for SpendAuth verification */ operatorAddress: string @@ -184,40 +187,6 @@ export interface ApiKeyInfo { dailyLimit?: number } -// --- Usage tracking --- - -export interface GatewayUsageEvent { - /** - * Per-request id (matches `RequestContext.requestId`). Lets - * `recordUsage` correlate the usage row to the same request that - * `settlePayment` settles, observability hooks observe, and - * `onRequestComplete` reports — without re-deriving from a - * synthetic key. Required field as of 0.4.0; the gateway always has - * it in scope at the recordUsage call site. - */ - requestId: string - agentId: string - agentSlug: string - consumerId: string - paymentMethod: PaymentMethod - inputTokens: number - outputTokens: number - /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ - reasoningTokens?: number - /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ - toolTokens?: number - /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ - toolCallCount?: number - /** Optional in 0.7.2 so 0.7.1 event constructors remain source-compatible. */ - providerCostUsd?: number - totalCostUsd: number - ownerEarnedUsd: number - platformFeeUsd: number - durationMs: number - /** Exact receipt in normal operation; quoted ceiling only after receipt timeout. */ - settlementBasis?: PaymentSettlementBasis -} - // --- Sandbox interface --- export interface SandboxStreamEvent { @@ -243,17 +212,6 @@ export interface SandboxStreamEvent { } } -export interface SandboxUsageReceipt { - inputTokens: number - outputTokens: number - reasoningTokens: number - toolTokens: number - toolCallCount: number - providerCostUsd: number - /** True only when the provider/adapter enforced every supplied budget. */ - budgetEnforced: boolean -} - export interface SandboxBox { streamPrompt( message: string, @@ -372,7 +330,7 @@ export interface GatewayConfig { * and settlement failures. See ./observer.ts for the interface and * ConsoleObserver / CompositeObserver implementations. */ - observer?: import('./observer').GatewayObserver + observer?: GatewayObserver /** * A2A protocol configuration. The gateway exposes A2A with an in-memory diff --git a/tests/a2a-state-machines.test.ts b/tests/a2a-state-machines.test.ts new file mode 100644 index 0000000..c659708 --- /dev/null +++ b/tests/a2a-state-machines.test.ts @@ -0,0 +1,201 @@ +import { describe, expect, it, vi } from 'vitest' + +import type { AuthorizedRequest } from '../src/dispatch' +import { + attachPaymentRecoveryMarker, + clearPaymentRecoveryMarker, + readPaymentRecoveryMarker, + retainPaymentRecoveryMarker, +} from '../src/a2a/payment-recovery' +import { + buildFinalizationRecord, + clearFinalizationMarker, + markUsageRecorded, + readFinalizationRecord, + withFinalizationRecord, +} from '../src/a2a/task-finalization' +import { + InMemoryPushNotificationStore, + type PushNotificationConfig, +} from '../src/a2a/push-notifications' +import { deliverTaskPush } from '../src/a2a/task-push-delivery' +import { + clearTaskSubmission, + readTaskOrigin, + readTaskSubmission, + recoverSubmissionIfNeeded, + withTaskOrigin, + withTaskSubmission, +} from '../src/a2a/task-submission-recovery' +import { InMemoryTaskStore } from '../src/a2a/task-store' +import type { Task } from '../src/a2a/types' + +function makeTask(id: string, state: Task['status']['state'] = 'submitted'): Task { + return { + kind: 'task', + id, + contextId: `ctx-${id}`, + status: { state, timestamp: new Date().toISOString() }, + history: [], + } +} + +function makeAuthz(): AuthorizedRequest { + const agent = { + id: 'agent-1', + ownerId: 'owner-1', + slug: 'agent-1', + pricePerTokenUsd: 0.00002, + platformFeePercent: 0.2, + sandboxEndpoint: null, + remoteSandboxId: null, + remoteBearerToken: null, + enabled: true, + } + return { + agent, + consumerId: 'consumer-1', + paymentMethod: 'none', + keyInfo: null, + userMessage: 'hello', + rateLimitRemaining: undefined, + requestId: 'request-1', + startMs: Date.now(), + maxOutputTokens: 16, + executionBudget: { + maxInputTokens: 16, + maxOutputTokens: 16, + maxReasoningTokens: 16, + maxToolTokens: 16, + maxToolCalls: 2, + maxProviderCostUsd: 1, + }, + requiredPaymentAmount: 0n, + paymentPayload: null, + } +} + +const receipt = { + inputTokens: 1, + outputTokens: 2, + reasoningTokens: 0, + toolTokens: 0, + toolCallCount: 0, + providerCostUsd: 0.00006, + budgetEnforced: true, +} + +describe('A2A extracted state machines', () => { + it('preserves task origin while clearing only the submission lease', () => { + const task = makeTask('submission') + const withOrigin = withTaskOrigin(task.metadata, { id: 'agent-1', slug: 'agent-1' }) + const submitted = { + ...task, + metadata: withTaskSubmission(withOrigin, { + agent: { id: 'agent-1', slug: 'agent-1' }, + requestId: 'request-1', + consumerId: 'consumer-1', + }), + } + + expect(readTaskOrigin(submitted)).toMatchObject({ agentId: 'agent-1', agentSlug: 'agent-1' }) + expect(readTaskSubmission(submitted)).toMatchObject({ requestId: 'request-1' }) + expect(readTaskSubmission(clearTaskSubmission(submitted))).toBeUndefined() + expect(readTaskOrigin(clearTaskSubmission(submitted))).toMatchObject({ agentId: 'agent-1' }) + }) + + it('fails an expired submission and delivers its terminal transition once', async () => { + const store = new InMemoryTaskStore() + const task = makeTask('expired-submission') + const metadata = withTaskSubmission(task.metadata, { + agent: { id: 'agent-1', slug: 'agent-1' }, + requestId: 'request-1', + consumerId: 'consumer-1', + }) + const expired = { + ...task, + metadata: { + ...metadata, + gatewaySubmission: { + ...(metadata.gatewaySubmission as Record), + lease: { id: 'expired', expiresAt: 0 }, + }, + }, + } + await store.put(expired) + const delivered: Task[] = [] + + const recovered = await recoverSubmissionIfNeeded(expired, { + taskStore: store, + deliverPush: async (deliveredTask) => { + delivered.push(deliveredTask) + }, + }) + + expect(recovered.status.state).toBe('failed') + expect(recovered.metadata?.gatewaySubmission).toBeUndefined() + expect(recovered.metadata?.gatewaySubmissionRecovery).toBeDefined() + expect(delivered).toHaveLength(1) + }) + + it('claims a successful push delivery once across repeated terminal reads', async () => { + const taskStore = new InMemoryTaskStore() + const pushStore = new InMemoryPushNotificationStore() + const fetcher = vi.fn(async () => new Response(null, { status: 200 })) + const config: PushNotificationConfig = { id: 'webhook-1', url: 'https://example.test/done' } + const task = makeTask('push-once', 'completed') + await taskStore.put(task) + await pushStore.set(task.id, config) + + const deps = { taskStore, pushStore, demoMode: true, fetcher } + await deliverTaskPush((await taskStore.get(task.id))!, deps) + await deliverTaskPush((await taskStore.get(task.id))!, deps) + + expect(fetcher).toHaveBeenCalledTimes(1) + expect((await taskStore.get(task.id))?.metadata?.gatewayPushDelivery).toMatchObject({ + claims: { 'webhook-1': 'completed' }, + }) + }) + + it('does not claim a push rejected by URL policy, then retries it', async () => { + const taskStore = new InMemoryTaskStore() + const pushStore = new InMemoryPushNotificationStore() + const fetcher = vi.fn(async () => new Response(null, { status: 200 })) + const urlValidator = vi.fn().mockResolvedValueOnce(false).mockResolvedValueOnce(true) + const task = makeTask('push-retry', 'failed') + await taskStore.put(task) + await pushStore.set(task.id, { id: 'webhook-1', url: 'https://example.test/retry' }) + + const deps = { taskStore, pushStore, demoMode: true, fetcher, urlValidator } + await deliverTaskPush((await taskStore.get(task.id))!, deps) + expect((await taskStore.get(task.id))?.metadata?.gatewayPushDelivery).toBeUndefined() + await deliverTaskPush((await taskStore.get(task.id))!, deps) + + expect(urlValidator).toHaveBeenCalledTimes(2) + expect(fetcher).toHaveBeenCalledTimes(1) + }) + + it('records and clears finalization usage without changing the task contract', async () => { + const store = new InMemoryTaskStore() + const authz = makeAuthz() + const task = makeTask('finalization', 'working') + const record = buildFinalizationRecord(authz, receipt, null, false, undefined) + const marked = withFinalizationRecord(task, record) + await store.put(marked) + + const recorded = await markUsageRecorded(store, marked) + expect(readFinalizationRecord(recorded)?.usageRecorded).toBe(true) + expect(clearFinalizationMarker(recorded).metadata?.gatewayFinalizing).toBeUndefined() + }) + + it('attaches, retains, and clears a payment recovery identity atomically', async () => { + const store = new InMemoryTaskStore() + const task = makeTask('payment-marker') + await store.put(task) + + const attached = await attachPaymentRecoveryMarker(store, task, 'recovery-1') + const retained = await retainPaymentRecoveryMarker(store, attached, 'recovery-1') + expect(readPaymentRecoveryMarker(retained)).toEqual({ version: 1, id: 'recovery-1' }) + expect(readPaymentRecoveryMarker(clearPaymentRecoveryMarker(retained))).toBeUndefined() + }) +}) diff --git a/tests/module-size.test.ts b/tests/module-size.test.ts new file mode 100644 index 0000000..53ee6be --- /dev/null +++ b/tests/module-size.test.ts @@ -0,0 +1,26 @@ +import { readdirSync, readFileSync } from 'node:fs' +import { dirname, join, relative } from 'node:path' +import { fileURLToPath } from 'node:url' +import { describe, expect, it } from 'vitest' + +const sourceRoot = join(dirname(fileURLToPath(import.meta.url)), '../src/a2a') +const handlerPath = join(sourceRoot, 'handler.ts') +const extractedModulePaths = readdirSync(sourceRoot) + .filter((name) => name.endsWith('.ts') && name !== 'handler.ts') + .sort() + .map((name) => join(sourceRoot, name)) + +function countLines(path: string): number { + const source = readFileSync(path, 'utf8') + return source.length === 0 ? 0 : source.split(/\r?\n/).length - (source.endsWith('\n') ? 1 : 0) +} + +describe('A2A module size boundaries', () => { + it('keeps the handler below 1,000 lines and every A2A module below 500', () => { + expect(countLines(handlerPath), relative(process.cwd(), handlerPath)).toBeLessThan(1_000) + expect(extractedModulePaths.length).toBeGreaterThan(0) + for (const path of extractedModulePaths) { + expect(countLines(path), relative(process.cwd(), path)).toBeLessThan(500) + } + }) +}) From de86465736dd64e94cfbd3c51daad3c1c37603f1 Mon Sep 17 00:00:00 2001 From: Drew Stone Date: Sat, 15 Aug 2026 19:00:44 -0600 Subject: [PATCH 32/32] fix(a2a): clear execution state on release failure --- src/a2a/payment-recovery.ts | 6 +++--- tests/pr11-regressions.test.ts | 38 ++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/src/a2a/payment-recovery.ts b/src/a2a/payment-recovery.ts index de1b56a..45e1bad 100644 --- a/src/a2a/payment-recovery.ts +++ b/src/a2a/payment-recovery.ts @@ -415,11 +415,11 @@ async function expirePaymentRelease( error: Error, ): Promise { const cleanTask = clearPaymentReleaseRecord(task) + const terminalTask = withStatus(cleanTask, 'failed') const failed: Task = { - ...cleanTask, - status: { state: 'failed', timestamp: new Date().toISOString() }, + ...terminalTask, metadata: { - ...(cleanTask.metadata ?? {}), + ...(terminalTask.metadata ?? {}), gatewayPaymentReleaseRecovery: { error: error.message }, }, } diff --git a/tests/pr11-regressions.test.ts b/tests/pr11-regressions.test.ts index 2cee600..238a8c7 100644 --- a/tests/pr11-regressions.test.ts +++ b/tests/pr11-regressions.test.ts @@ -4,6 +4,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { InMemoryTaskStore, type TaskStore } from '../src/a2a/task-store' import { SqlTaskStore, type SqlAdapter } from '../src/a2a/task-store-sql' import { InMemoryPushNotificationStore } from '../src/a2a/push-notifications' +import { recoverPaymentReleaseIfNeeded } from '../src/a2a/payment-recovery' import { createAgentGateway } from '../src/middleware' import { dispatchSandboxStreamRich, requiredX402Amount } from '../src/dispatch' import { MemoryNonceStore, claimStoredNonce, type NonceStore } from '../src/nonce-store' @@ -1512,6 +1513,43 @@ describe('PR #11 production regressions', () => { .toEqual({ version: 1, id: 'payment-recovery-resubscribe' }) }) + it('clears the execution marker when malformed payment release metadata fails a task', async () => { + const taskStore = new InMemoryTaskStore() + const task: Task = { + kind: 'task', + id: 'malformed-payment-release', + contextId: 'malformed-payment-release-context', + status: { state: 'working', timestamp: new Date().toISOString() }, + metadata: { + gatewayExecution: { + version: 1, + requestId: 'worker-a', + lease: { id: 'worker-a', expiresAt: Date.now() + 60_000 }, + }, + gatewayPaymentRelease: { version: 1 }, + }, + } + await taskStore.put(task) + const deliverPush = vi.fn(async () => undefined) + + const recovered = await recoverPaymentReleaseIfNeeded(task, { + taskStore, + releasePayment: async () => undefined, + releasePaymentAfterFailure: async () => undefined, + recoverDurablePayment: async () => undefined, + deliverPush, + }) + + expect(recovered.status.state).toBe('failed') + expect(recovered.metadata?.gatewayExecution).toBeUndefined() + expect(recovered.metadata?.gatewayPaymentRelease).toBeUndefined() + expect(recovered.metadata?.gatewayPaymentReleaseRecovery).toMatchObject({ + error: expect.stringContaining('record is missing'), + }) + expect((await taskStore.get(task.id))?.metadata?.gatewayExecution).toBeUndefined() + expect(deliverPush).toHaveBeenCalledOnce() + }) + it('rejects direct private push destinations before fetch', async () => { const { deliverPushNotifications } = await import('../src/a2a/push-notifications') const fetcher = vi.fn(async () => new Response('unexpected', { status: 200 }))