@tangle-network/agent-gateway 0.7.0 → 0.8.0

package/src/a2a/handler.ts

@@ -13,36 +13,87 @@
 import type { Context } from 'hono'
 
 import {
-  type A2ADispatchEvent,
   type AuthorizedRequest,
   type GatewayState,
   authenticateAndGuard,
-  dispatchSandboxStreamRich,
-  estimateTokens,
-  settleAndRecord,
+  claimPayment,
 } from '../dispatch'
-import type { GatewayConfig } from '../types'
+import type {
+  ChatMessage,
+  GatewayConfig,
+} from '../types'
 import { buildAgentCard } from './agent-card'
+import {
+  clearTaskExecution,
+  hasExpiredTaskExecution,
+  hasMalformedTaskExecution,
+  inspectTaskExecution,
+} from './execution-fence'
+import { executeMessageSend } from './message-send-execution'
+import { executeMessageStream } from './message-stream-execution'
 import { fail, ok, parseEnvelope } from './jsonrpc'
 import {
-  deliverPushNotifications,
   type PushNotificationStore,
-  type TaskPushNotificationConfig,
 } from './push-notifications'
-import 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 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
@@ -51,51 +102,60 @@ export interface A2AHandlerDeps {
   pushStore?: PushNotificationStore
 }
 
-/** Terminal task states — fire-once push delivery occurs on these transitions. */
-const TERMINAL_STATES: ReadonlySet<Task['status']['state']> = new Set([
-  'completed',
-  'canceled',
-  'failed',
-  'rejected',
-])
+const MAX_A2A_BODY_BYTES = 64 * 1024
 
-/**
- * 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<string, AbortController>()
+class RequestBodyTooLargeError extends Error {}
 
-  register(taskId: string): AbortController {
-    const c = new AbortController()
-    this.controllers.set(taskId, c)
-    return c
-  }
+function createTaskLifecycle(deps: A2AHandlerDeps): TaskLifecycle {
+  return buildTaskLifecycle({
+    taskStore: deps.taskStore,
+    config: deps.config,
+    state: deps.state,
+    deliverPush: (task) => maybeDeliverPush(task, deps),
+  })
+}
 
-  clear(taskId: string): void {
-    this.controllers.delete(taskId)
+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 {
-    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),
   }
 }
 
 export function createA2AHandlers(deps: A2AHandlerDeps) {
-  const cancels = new CancelRegistry()
+  const runtimeDeps: A2AHandlerDeps = {
+    ...deps,
+    taskStore: normalizeTaskStore(deps.taskStore, deps.config.x402.demoMode === true),
+  }
+  const cancels = new TaskCancellationRegistry()
 
   // GET /:slug/.well-known/agent.json
   const handleAgentCard = async (c: Context): Promise<Response> => {
     const slug = c.req.param('slug')
     if (!slug) return c.json({ error: 'slug required' }, 400)
     const agent = await deps.config.resolveAgent(slug)
-    if (!agent) {
+    if (!agent || !agent.enabled) {
       return c.json({ error: 'Agent not found or not published' }, 404)
     }
     const url = new URL(c.req.url)
@@ -112,14 +172,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)
@@ -129,23 +195,23 @@ export function createA2AHandlers(deps: A2AHandlerDeps) {
 
     switch (parsed.method) {
       case 'message/send':
-        return handleMessageSend(c, slug, parsed, deps)
+        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, buildTaskMethodDependencies(runtimeDeps, cancels))
       case 'tasks/cancel':
-        return handleTasksCancel(c, parsed, deps, cancels)
+        return handleTasksCancel(c, parsed, buildTaskMethodDependencies(runtimeDeps, cancels))
       case 'tasks/resubscribe':
-        return handleTasksResubscribe(c, parsed, deps)
+        return handleTasksResubscribe(c, parsed, buildTaskMethodDependencies(runtimeDeps, cancels))
       case 'tasks/pushNotificationConfig/set':
-        return handlePushSet(c, parsed, deps)
+        return handlePushSet(c, parsed, buildPushConfigMethodDependencies(runtimeDeps))
       case 'tasks/pushNotificationConfig/get':
-        return handlePushGet(c, parsed, deps)
+        return handlePushGet(c, parsed, buildPushConfigMethodDependencies(runtimeDeps))
       case 'tasks/pushNotificationConfig/list':
-        return handlePushList(c, parsed, deps)
+        return handlePushList(c, parsed, buildPushConfigMethodDependencies(runtimeDeps))
       case 'tasks/pushNotificationConfig/delete':
-        return handlePushDelete(c, parsed, deps)
+        return handlePushDelete(c, parsed, buildPushConfigMethodDependencies(runtimeDeps))
       default:
         return c.json(
           fail(parsed.id, A2A_ERROR_CODES.METHOD_NOT_FOUND, `unknown method '${parsed.method}'`),
@@ -163,77 +229,28 @@ async function handleMessageSend(
   slug: string,
   req: JSONRPCRequest,
   deps: A2AHandlerDeps,
+  cancels: TaskCancellationRegistry,
 ): Promise<Response> {
   const guard = await guardMessageRequest(c, slug, req, deps)
   if (guard instanceof Response) return guard
   const { authz, task } = guard
-
-  let responseText = ''
-  let outputTokens = 0
-  let inputRequiredPrompt: string | undefined
-  let inputRequiredSeen = false
+  setPaymentResponseHeaders(c, authz)
+  const controller = cancels.register(task.id)
+  const detachRequestAbort = bindRequestAbort(c.req.raw.signal, controller)
+  const lifecycle = createTaskLifecycle(deps)
   try {
-    for await (const event of dispatchSandboxStreamRich(
-      authz.agent,
-      authz.userMessage,
-      authz.consumerId,
-      deps.config,
-      undefined,
-      task.id,
-    )) {
-      if (event.kind === 'text') {
-        responseText += event.delta
-        outputTokens += estimateTokens(event.delta)
-      } else {
-        inputRequiredSeen = true
-        inputRequiredPrompt = event.prompt
-      }
-    }
-  } catch (err) {
-    const failed = withStatus(task, 'failed')
-    await deps.taskStore.put(failed)
-    await maybeDeliverPush(failed, deps)
-    return c.json(
-      fail(
-        req.id,
-        A2A_ERROR_CODES.INTERNAL_ERROR,
-        err instanceof Error ? err.message : String(err),
-      ),
-    )
-  }
-
-  // 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,
-    estimateTokens(authz.userMessage),
-    outputTokens,
-    deps.config,
-    deps.state.obs,
-  )
-
-  if (inputRequiredSeen) {
-    const paused = withStatus(
+    return await executeMessageSend(
+      c,
+      req,
+      { taskStore: deps.taskStore, config: deps.config, ...lifecycle },
+      authz,
       task,
-      'input-required',
-      inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined,
-      responseText
-        ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)]
-        : task.artifacts,
+      controller.signal,
     )
-    await deps.taskStore.put(paused)
-    // input-required is non-terminal — do NOT deliver push notifications.
-    return c.json(ok(req.id, paused))
+  } finally {
+    detachRequestAbort()
+    cancels.clear(task.id)
   }
-
-  const completed = withStatus(task, '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) ──────────────────────────────────────────────────
@@ -243,386 +260,38 @@ async function handleMessageStream(
   slug: string,
   req: JSONRPCRequest,
   deps: A2AHandlerDeps,
-  cancels: CancelRegistry,
+  cancels: TaskCancellationRegistry,
 ): Promise<Response> {
   const guard = await guardMessageRequest(c, slug, req, deps)
   if (guard instanceof Response) return guard
   const { authz, task } = guard
-
-  const controller = cancels.register(task.id)
-  const inputTokens = estimateTokens(authz.userMessage)
-  let outputTokens = 0
-  let responseText = ''
-
-  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`))
-      }
-
-      // Status: working
-      const workingStatus: TaskStatusUpdateEvent = {
-        kind: 'status-update',
-        taskId: task.id,
-        contextId: task.contextId,
-        status: { state: 'working', timestamp: nowIso() },
-        final: false,
-      }
-      await deps.taskStore.put({ ...task, status: workingStatus.status })
-      send(workingStatus)
-
-      let inputRequiredPrompt: string | undefined
-      let inputRequiredSeen = false
-      try {
-        for await (const event of dispatchSandboxStreamRich(
-          authz.agent,
-          authz.userMessage,
-          authz.consumerId,
-          deps.config,
-          controller.signal,
-          task.id,
-        )) {
-          if (event.kind === 'text') {
-            responseText += event.delta
-            outputTokens += estimateTokens(event.delta)
-            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 {
-            inputRequiredSeen = true
-            inputRequiredPrompt = event.prompt
-          }
-        }
-
-        // Caller aborted via tasks/cancel — emit canceled, do not settle.
-        if (controller.signal.aborted) {
-          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)
-          return
-        }
-
-        // Settle once for whatever the sandbox produced (full or partial).
-        await settleAndRecord(
-          authz.agent,
-          authz,
-          inputTokens,
-          outputTokens,
-          deps.config,
-          deps.state.obs,
-        )
-
-        if (inputRequiredSeen) {
-          const paused = withStatus(
-            task,
-            'input-required',
-            inputRequiredPrompt ? agentMessage(task, inputRequiredPrompt) : undefined,
-            responseText
-              ? [responseTextToArtifact(responseText, `${task.id}-artifact-0`)]
-              : task.artifacts,
-          )
-          await deps.taskStore.put(paused)
-          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: artifact lastChunk + completed status.
-        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,
-        })
-        const completed = withStatus(task, 'completed', undefined, [
-          responseTextToArtifact(responseText, `${task.id}-artifact-0`),
-        ])
-        await deps.taskStore.put(completed)
-        send({
-          kind: 'status-update',
-          taskId: task.id,
-          contextId: task.contextId,
-          status: completed.status,
-          final: true,
-        })
-        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)
+  setPaymentResponseHeaders(c, authz)
+  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: authz.requestId,
-            agentSlug: authz.agent.slug,
-            startMs: authz.startMs,
+            requestId: streamAuthz.requestId,
+            agentSlug: streamAuthz.agent.slug,
+            startMs: streamAuthz.startMs,
           },
           {
-            consumerId: authz.consumerId,
-            errorMessage: err instanceof Error ? err.message : String(err),
+            consumerId: streamAuthz.consumerId,
+            errorMessage: error instanceof Error ? error.message : String(error),
           },
         )
-      } finally {
-        cancels.clear(task.id)
-        ctrl.close()
-      }
-    },
-  })
-
-  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,
-    },
-  })
-}
-
-// ── tasks/get + tasks/cancel ──────────────────────────────────────────────
-
-async function handleTasksGet(
-  c: Context,
-  req: JSONRPCRequest,
-  deps: A2AHandlerDeps,
-): Promise<Response> {
-  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`),
-    )
-  }
-  return c.json(ok(req.id, task))
-}
-
-async function handleTasksCancel(
-  c: Context,
-  req: JSONRPCRequest,
-  deps: A2AHandlerDeps,
-  cancels: CancelRegistry,
-): Promise<Response> {
-  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`),
-    )
-  }
-  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}'`,
-      ),
-    )
-  }
-
-  const stillActive = cancels.cancel(task.id)
-  const canceled: Task = {
-    ...task,
-    status: { state: 'canceled', timestamp: nowIso() },
-  }
-  await deps.taskStore.put(canceled)
-
-  // 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)
-  }
-  return c.json(ok(req.id, canceled))
-}
-
-// ── 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<Response> {
-  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 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<Response> {
-  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`))
-  }
-  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<Response> {
-  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 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<Response> {
-  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 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<Response> {
-  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'),
-    )
-  }
-  await deps.pushStore.delete(params.id, params.pushNotificationConfigId)
-  return c.json(ok(req.id, null))
+    authz,
+    task,
+  )
 }
 
 // ── Shared message-send setup (auth + task allocation) ────────────────────
@@ -647,23 +316,47 @@ 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 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 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 recoverTaskIfNeeded(storedExisting, deps, slug)
       if (existing.status.state !== 'input-required') {
         return c.json(
           fail(
@@ -680,11 +373,18 @@ async function guardMessageRequest(
       }
       const continued: Task = {
         ...existing,
-        status: { state: 'working', timestamp: nowIso() },
+        status: { state: 'submitted', timestamp: nowIso() },
         history: [...(existing.history ?? []), appendedMessage],
+        metadata: withTaskSubmission(existing.metadata, authz),
       }
-      await deps.taskStore.put(continued)
-      return { authz, task: 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 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.
@@ -703,62 +403,316 @@ 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(
+      fail(req.id, A2A_ERROR_CODES.INVALID_PARAMS, `task '${task.id}' already exists`),
+    )
+  }
+  const claimedTask = await claimTaskPayment(c, req, task, authz, deps)
+  if (claimedTask instanceof Response) return claimedTask
+  return { authz, task: claimedTask }
+}
+
+async function claimTaskPayment(
+  c: Context,
+  req: JSONRPCRequest,
+  task: Task,
+  authz: AuthorizedRequest,
+  deps: A2AHandlerDeps,
+  paymentFailureTask?: Task,
+): Promise<Response | Task> {
+  const lifecycle = createTaskLifecycle(deps)
+  let paymentTask = task
+  try {
+    await claimPayment(authz, deps.config, deps.state, {
+      onRecoveryPrepared: async (recoveryId) => {
+        paymentTask = await attachPaymentRecoveryMarker(
+          deps.taskStore,
+          paymentTask,
+          recoveryId,
+        )
+      },
+    })
+  } catch {
+    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,
+      lifecycle.payment,
+      'payment authorization failed',
+      false,
+    )
+    const cleanedReleasedTask = clearTaskSubmission(releasedTask)
+    const hasReleaseRecord = hasPaymentReleaseRecovery(cleanedReleasedTask)
+    const failed = hasReleaseRecord
+      ? isTerminal(cleanedReleasedTask.status.state)
+        ? cleanedReleasedTask
+        : withStatus(cleanedReleasedTask, 'failed')
+      : paymentFailureTask
+        ? 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)
+      }
+    } 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'))
+  }
+
+  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,
+      lifecycle.payment,
+      '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`),
+    )
   }
-  await deps.taskStore.put(task)
-  return { authz, task }
 }
 
 // ── Helpers ───────────────────────────────────────────────────────────────
 
-function isTerminal(state: Task['status']['state']): boolean {
-  return (
-    state === 'completed' ||
-    state === 'canceled' ||
-    state === 'failed' ||
-    state === 'rejected'
-  )
+const EXECUTION_RECOVERY_METADATA_KEY = 'gatewayExecutionRecovery'
+
+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)
+  }
 }
 
-function nowIso(): string {
-  return new Date().toISOString()
+async function authorizeTaskAccess(
+  c: Context,
+  req: JSONRPCRequest,
+  task: Task,
+  deps: A2AHandlerDeps,
+): Promise<Response | undefined> {
+  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) {
+    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: requestedAgentSlug,
+      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 cryptoRandomId(): string {
-  return crypto.randomUUID().replace(/-/g, '')
+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,
+    }]
+  })
 }
 
-/**
- * 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 {
+async function createTask(taskStore: TaskStore, task: Task): Promise<boolean> {
+  if (!taskStore.createIfAbsent) {
+    throw new Error('A2A task store does not provide createIfAbsent')
+  }
+  return taskStore.createIfAbsent(task)
+}
+function normalizeTaskStore(taskStore: TaskStore, allowUnsafeFallback: boolean): TaskStore {
+  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, compareAndSet, and compareAndSetExecution',
+    )
+  }
   return {
-    ...task,
-    status: { state, timestamp: nowIso(), ...(message ? { message } : {}) },
-    ...(artifacts !== undefined ? { artifacts } : {}),
+    get: (id) => taskStore.get(id),
+    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
+    },
   }
 }
 
-/**
- * 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}-status-${task.status.state}-${nowIso()}`,
-    taskId: task.id,
-    contextId: task.contextId,
+async function recoverTaskIfNeeded(
+  task: Task,
+  deps: A2AHandlerDeps,
+  requestedAgentSlug: string,
+): Promise<Task> {
+  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)
+}
+
+async function recoverExpiredExecutionIfNeeded(task: Task, deps: A2AHandlerDeps): Promise<Task> {
+  const malformed = hasMalformedTaskExecution(task)
+  if (
+    task.status.state !== 'working' ||
+    (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: inspection.state === 'malformed'
+          ? `A2A execution marker was malformed: ${inspection.reason}`
+          : '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 readJsonBody(request: Request): Promise<unknown> {
+  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
 }
 
 /**
@@ -768,31 +722,22 @@ function agentMessage(task: Task, text: string): Message {
  * via `tasks/get` to confirm state.
  */
 async function maybeDeliverPush(task: Task, deps: A2AHandlerDeps): Promise<void> {
-  if (!deps.pushStore || !TERMINAL_STATES.has(task.status.state)) return
-  try {
-    await deliverPushNotifications({
-      task,
-      store: deps.pushStore,
-      webhookSecret: deps.config.a2a?.webhookSecret,
-      fetcher: deps.config.a2a?.pushFetcher,
-      onDelivery: (result) => {
-        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'}`,
-            },
-          )
-        }
-      },
-    })
-  } 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)}`,
-    )
+  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'}`,
+        },
+      )
+    },
   }
+  await deliverTaskPush(task, pushDeps)
 }
-