@tangle-network/agent-gateway 0.7.1 → 0.8.1

package/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<string, Task['status']['state']>
+}
+
+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<void> {
+  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<PushNotificationDeliveryOptions, 'webhookSecret'> = {
+      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<boolean> {
+  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<TaskPushDeliveryClaims>
+  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<string, Task['status']['state']>
+  return record.version === 1 ? { version: 1, claims } : undefined
+}

package/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)
+}

package/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<Task | undefined>
+  compareAndSet?(expected: Task, next: Task): Promise<boolean>
+}
+
+export const TERMINAL_STATES: ReadonlySet<Task['status']['state']> = 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<boolean> {
+  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<Task> {
+  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')
+}

package/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
@@ -41,7 +40,8 @@
  *   await store.migrate()
  */
 
-import type { TaskStore } from './task-store'
+import { inspectTaskExecution } from './execution-fence'
+import { hasPendingPaymentRecovery, type TaskStore } from './task-store'
 import type { Task } from './types'
 
 /**
@@ -98,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) => `
@@ -124,50 +126,236 @@ export class SqlTaskStore implements TaskStore {
     return this.opts.table ?? 'a2a_tasks'
   }
 
+  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
+      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,
+          executionRequestId: row.execution_request_id ?? null,
+          executionLeaseExpiresAt: row.execution_lease_expires_at ?? null,
+        }
+      : 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<number> {
+    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<void> {
     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))
   }
 
   async get(id: string): Promise<Task | 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]
+    const row = await this.readRow(id)
     if (!row) return undefined
-    if (Date.now() - row.updated_at > this.ttlMs) {
-      // 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])
+    const task = JSON.parse(row.payload) as 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 and updated_at both fence the row.
+      await this.deleteObservedRow(id, row.payload, row.updatedAt)
       return undefined
     }
-    return JSON.parse(row.payload) as Task
+    return task
+  }
+
+  private async insert(task: Task): Promise<number> {
+    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, execution_request_id, execution_lease_expires_at)
+       VALUES (?, ?, ?, ?, ?, ?, ?)`,
+      [
+        task.id,
+        task.contextId,
+        task.status.state,
+        payload,
+        Date.now(),
+        executionRequestId,
+        executionLeaseExpiresAt,
+      ],
+    )
+    return result.rowsAffected
   }
 
   async put(task: Task): Promise<void> {
     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,
+        ],
       )
     }
   }
 
+  async createIfAbsent(task: Task): Promise<boolean> {
+    try {
+      return (await this.insert(task)) === 1
+    } catch (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
+      }
+    }
+  }
+
+  async compareAndSet(expected: Task, next: Task): Promise<boolean> {
+    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<boolean> {
+    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 = ?,
+           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
+  }
+
   async delete(id: string): Promise<void> {
-    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)],
+    )
   }
 
   /**
@@ -183,7 +371,17 @@ 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)
   }
 }
+
+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]
+}

package/src/a2a/task-store.ts

@@ -6,10 +6,25 @@
  */
 
 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<Task | undefined>
   put(task: Task): Promise<void>
+  /** Insert only when the task id is absent. Required outside explicit demo mode. */
+  createIfAbsent?(task: Task): Promise<boolean>
+  /** Replace only when the stored task still equals `expected`. Required for production races. */
+  compareAndSet?(expected: Task, next: Task): Promise<boolean>
+  /** Replace an execution marker only while its owner lease is still live. */
+  compareAndSetExecution?(
+    expected: Task,
+    next: Task,
+    requestId: string,
+    now: number,
+  ): Promise<boolean>
   delete(id: string): Promise<void>
 }
 
@@ -32,7 +47,47 @@ export class InMemoryTaskStore implements TaskStore {
     this.entries.set(task.id, { task: clone(task), expiresAt: Date.now() + this.ttlMs })
   }
 
+  async createIfAbsent(task: Task): Promise<boolean> {
+    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<boolean> {
+    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 compareAndSetExecution(
+    expected: Task,
+    next: Task,
+    requestId: string,
+    now: number,
+  ): Promise<boolean> {
+    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<void> {
+    const entry = this.entries.get(id)
+    if (entry && hasPendingPaymentRecovery(entry.task)) return
     this.entries.delete(id)
   }
 
@@ -43,7 +98,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)
+      }
     }
   }
 }