diff --git a/README.md b/README.md index f64830ba..511a4596 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ You can install the A2A SDK using `npm`. npm install @a2a-js/sdk ``` -### For Server Usage +### A2A Server If you plan to use the A2A server functionality (`A2AExpressApp`), you'll also need to install Express as it's a peer dependency: @@ -29,6 +29,24 @@ If you plan to use the A2A server functionality (`A2AExpressApp`), you'll also n npm install express ``` +### Persistent Task Storage + +For production deployments requiring persistent task storage, install Drizzle ORM, a database driver, and drizzle-kit for migrations: + +```bash +# SQLite +npm install drizzle-orm better-sqlite3 +npm install -D drizzle-kit + +# PostgreSQL +npm install drizzle-orm pg +npm install -D drizzle-kit + +# MySQL +npm install drizzle-orm mysql2 +npm install -D drizzle-kit +``` + You can also find JavaScript samples [here](https://github.com/google-a2a/a2a-samples/tree/main/samples/js). --- @@ -669,6 +687,87 @@ app.post('/webhook/task-updates', (req, res) => { }); ``` +--- + +## Persistent Task Storage + +By default, the SDK provides an `InMemoryTaskStore` which stores tasks in memory. For production deployments, use the `DatabaseTaskStore` with Drizzle ORM to persist tasks to SQLite, PostgreSQL, or MySQL. + +### Quick Setup + +**1. Create a schema file** (e.g., `src/schema.ts`): + +```typescript +// For SQLite +export { sqliteTasks as tasks } from '@a2a-js/sdk/server/drizzle'; +// For PostgreSQL: export { pgTasks as tasks } from '@a2a-js/sdk/server/drizzle'; +// For MySQL: export { mysqlTasks as tasks } from '@a2a-js/sdk/server/drizzle'; +``` + +**2. Create `drizzle.config.ts`** in your project root: + +```typescript +import { defineConfig } from 'drizzle-kit'; + +export default defineConfig({ + schema: './src/schema.ts', + out: './drizzle', + dialect: 'sqlite', // or 'postgresql' or 'mysql' + dbCredentials: { + url: 'tasks.db', // or your connection string + }, +}); +``` + +**3. Apply the schema:** + +```bash +# For development (applies changes directly): +npx drizzle-kit push + +# For production (version-controlled migrations): +npx drizzle-kit generate +npx drizzle-kit migrate +``` + +### Usage Examples + +**SQLite:** + +```typescript +import { drizzle } from 'drizzle-orm/better-sqlite3'; +import Database from 'better-sqlite3'; +import { DatabaseTaskStore, sqliteTasks } from '@a2a-js/sdk/server/drizzle'; + +const sqlite = new Database('tasks.db'); +const db = drizzle(sqlite); +const taskStore = new DatabaseTaskStore({ db, table: sqliteTasks, dialect: 'sqlite' }); +``` + +**PostgreSQL** (with connection pooling): + +```typescript +import { drizzle } from 'drizzle-orm/node-postgres'; +import { Pool } from 'pg'; +import { DatabaseTaskStore, pgTasks } from '@a2a-js/sdk/server/drizzle'; + +const pool = new Pool({ connectionString: process.env.DATABASE_URL }); +const db = drizzle(pool); +const taskStore = new DatabaseTaskStore({ db, table: pgTasks, dialect: 'postgresql' }); +``` + +**MySQL** (with connection pooling): + +```typescript +import { drizzle } from 'drizzle-orm/mysql2'; +import mysql from 'mysql2/promise'; +import { DatabaseTaskStore, mysqlTasks } from '@a2a-js/sdk/server/drizzle'; + +const pool = mysql.createPool(process.env.DATABASE_URL); +const db = drizzle(pool); +const taskStore = new DatabaseTaskStore({ db, table: mysqlTasks, dialect: 'mysql' }); +``` + ## License This project is licensed under the terms of the [Apache 2.0 License](https://raw.githubusercontent.com/google-a2a/a2a-python/refs/heads/main/LICENSE). diff --git a/package-lock.json b/package-lock.json index 268ba240..54a9cbe3 100644 --- a/package-lock.json +++ b/package-lock.json @@ -22,6 +22,7 @@ "@types/supertest": "^6.0.3", "c8": "^10.1.3", "chai": "^5.2.0", + "drizzle-orm": "^0.44.7", "esbuild": "^0.27.0", "eslint": "^9.39.1", "eslint-config-prettier": "^10.1.8", @@ -43,9 +44,13 @@ "node": ">=18" }, "peerDependencies": { + "drizzle-orm": ">=0.30.0", "express": "^4.21.2 || ^5.1.0" }, "peerDependenciesMeta": { + "drizzle-orm": { + "optional": true + }, "express": { "optional": true } @@ -5980,6 +5985,132 @@ "yaml": "^2.8.0" } }, + "node_modules/drizzle-orm": { + "version": "0.44.7", + "resolved": "https://registry.npmjs.org/drizzle-orm/-/drizzle-orm-0.44.7.tgz", + "integrity": "sha512-quIpnYznjU9lHshEOAYLoZ9s3jweleHlZIAWR/jX9gAWNg/JhQ1wj0KGRf7/Zm+obRrYd9GjPVJg790QY9N5AQ==", + "dev": true, + "license": "Apache-2.0", + "peerDependencies": { + "@aws-sdk/client-rds-data": ">=3", + "@cloudflare/workers-types": ">=4", + "@electric-sql/pglite": ">=0.2.0", + "@libsql/client": ">=0.10.0", + "@libsql/client-wasm": ">=0.10.0", + "@neondatabase/serverless": ">=0.10.0", + "@op-engineering/op-sqlite": ">=2", + "@opentelemetry/api": "^1.4.1", + "@planetscale/database": ">=1.13", + "@prisma/client": "*", + "@tidbcloud/serverless": "*", + "@types/better-sqlite3": "*", + "@types/pg": "*", + "@types/sql.js": "*", + "@upstash/redis": ">=1.34.7", + "@vercel/postgres": ">=0.8.0", + "@xata.io/client": "*", + "better-sqlite3": ">=7", + "bun-types": "*", + "expo-sqlite": ">=14.0.0", + "gel": ">=2", + "knex": "*", + "kysely": "*", + "mysql2": ">=2", + "pg": ">=8", + "postgres": ">=3", + "sql.js": ">=1", + "sqlite3": ">=5" + }, + "peerDependenciesMeta": { + "@aws-sdk/client-rds-data": { + "optional": true + }, + "@cloudflare/workers-types": { + "optional": true + }, + "@electric-sql/pglite": { + "optional": true + }, + "@libsql/client": { + "optional": true + }, + "@libsql/client-wasm": { + "optional": true + }, + "@neondatabase/serverless": { + "optional": true + }, + "@op-engineering/op-sqlite": { + "optional": true + }, + "@opentelemetry/api": { + "optional": true + }, + "@planetscale/database": { + "optional": true + }, + "@prisma/client": { + "optional": true + }, + "@tidbcloud/serverless": { + "optional": true + }, + "@types/better-sqlite3": { + "optional": true + }, + "@types/pg": { + "optional": true + }, + "@types/sql.js": { + "optional": true + }, + "@upstash/redis": { + "optional": true + }, + "@vercel/postgres": { + "optional": true + }, + "@xata.io/client": { + "optional": true + }, + "better-sqlite3": { + "optional": true + }, + "bun-types": { + "optional": true + }, + "expo-sqlite": { + "optional": true + }, + "gel": { + "optional": true + }, + "knex": { + "optional": true + }, + "kysely": { + "optional": true + }, + "mysql2": { + "optional": true + }, + "pg": { + "optional": true + }, + "postgres": { + "optional": true + }, + "prisma": { + "optional": true + }, + "sql.js": { + "optional": true + }, + "sqlite3": { + "optional": true + } + } + }, "node_modules/dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", diff --git a/package.json b/package.json index d3c1b3f9..72189823 100644 --- a/package.json +++ b/package.json @@ -28,6 +28,11 @@ "import": "./dist/server/express/index.js", "require": "./dist/server/express/index.cjs" }, + "./server/drizzle": { + "types": "./dist/server/drizzle/index.d.ts", + "import": "./dist/server/drizzle/index.js", + "require": "./dist/server/drizzle/index.cjs" + }, "./client": { "types": "./dist/client/index.d.ts", "import": "./dist/client/index.js", @@ -51,6 +56,7 @@ "@types/supertest": "^6.0.3", "c8": "^10.1.3", "chai": "^5.2.0", + "drizzle-orm": "^0.44.7", "esbuild": "^0.27.0", "eslint": "^9.39.1", "eslint-config-prettier": "^10.1.8", @@ -88,11 +94,15 @@ "uuid": "^11.1.0" }, "peerDependencies": { + "drizzle-orm": ">=0.30.0", "express": "^4.21.2 || ^5.1.0" }, "peerDependenciesMeta": { "express": { "optional": true + }, + "drizzle-orm": { + "optional": true } }, "mocha": { diff --git a/src/server/drizzle/database_task_store.ts b/src/server/drizzle/database_task_store.ts new file mode 100644 index 00000000..1e7ecea8 --- /dev/null +++ b/src/server/drizzle/database_task_store.ts @@ -0,0 +1,218 @@ +/** + * DatabaseTaskStore - A persistent task store implementation using Drizzle ORM. + * + * This module provides a database-backed implementation of the TaskStore interface, + * supporting PostgreSQL, MySQL, and SQLite databases through Drizzle ORM. + * + * @example + * ```typescript + * import { drizzle } from 'drizzle-orm/better-sqlite3'; + * import Database from 'better-sqlite3'; + * import { DatabaseTaskStore, sqliteTasks } from '@a2a-js/sdk/server/drizzle'; + * + * const sqlite = new Database('tasks.db'); + * const db = drizzle(sqlite); + * + * const taskStore = new DatabaseTaskStore(db, sqliteTasks); + * ``` + */ +import { eq, type InferSelectModel } from 'drizzle-orm'; +import type { Task } from '../../types.js'; +import type { TaskStore } from '../store.js'; +import type { SqliteTasksTable, PgTasksTable, MysqlTasksTable } from './schema.js'; + +/** + * Union type for all supported Drizzle database instances. + * Supports any Drizzle database that has select, insert, and delete methods. + * + * This type is intentionally broad to accommodate different Drizzle dialects + * (SQLite, PostgreSQL, MySQL) which have slightly different method signatures. + */ +export type DrizzleDatabase = { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + select: () => SelectQueryBuilder; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + insert: (table: TTable) => InsertQueryBuilder; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + delete: (table: TTable) => DeleteQueryBuilder; +}; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +type SelectQueryBuilder = { + from: (table: TTable) => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + where: (condition: TCondition) => Promise; + }; +}; + +type InsertQueryBuilder = { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + values: (values: TValues) => InsertValuesQueryBuilder; +}; + +type InsertValuesQueryBuilder = { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + onConflictDoUpdate?: (config: TConfig) => Promise; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + onDuplicateKeyUpdate?: (config: TConfig) => Promise; +}; + +type DeleteQueryBuilder = { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + where: (condition: TCondition) => Promise; +}; + +/** + * Configuration options for DatabaseTaskStore. + */ +export interface DatabaseTaskStoreOptions< + T extends SqliteTasksTable | PgTasksTable | MysqlTasksTable, +> { + /** + * The Drizzle database instance to use for storage. + */ + db: DrizzleDatabase; + + /** + * The tasks table schema to use. + * Use sqliteTasks, pgTasks, or mysqlTasks from the schema module. + */ + table: T; + + /** + * The database dialect. Used to determine the correct upsert strategy. + * Required to ensure correct SQL generation for your database. + * + * @example 'sqlite' for SQLite databases + * @example 'postgresql' for PostgreSQL databases + * @example 'mysql' for MySQL/MariaDB databases + */ + dialect: T extends SqliteTasksTable + ? 'sqlite' + : T extends PgTasksTable + ? 'postgresql' + : T extends MysqlTasksTable + ? 'mysql' + : 'sqlite' | 'postgresql' | 'mysql'; +} + +type TaskRow = + | InferSelectModel + | InferSelectModel + | InferSelectModel; + +/** + * A persistent task store implementation using Drizzle ORM. + * + * Supports PostgreSQL, MySQL, and SQLite databases through Drizzle ORM. + * Tasks are stored in a single table with JSON columns for complex nested data. + */ +export class DatabaseTaskStore + implements TaskStore +{ + private db: DrizzleDatabase; + private table: T; + private dialect: 'sqlite' | 'postgresql' | 'mysql'; + + constructor(options: DatabaseTaskStoreOptions) { + this.db = options.db; + this.table = options.table; + this.dialect = options.dialect; + } + + /** + * Saves a task to the database. + * If a task with the same ID exists, it will be updated. + * The updatedAt timestamp is automatically set by the database schema. + * + * @param task The task to save. + */ + async save(task: Task): Promise { + const row = this.taskToRow(task); + const insertQuery = this.db.insert(this.table).values(row); + + const updateFields = { + contextId: row.contextId, + kind: row.kind, + status: row.status, + artifacts: row.artifacts, + history: row.history, + metadata: row.metadata, + }; + + if (this.dialect === 'mysql') { + // MySQL uses ON DUPLICATE KEY UPDATE + if (!insertQuery.onDuplicateKeyUpdate) { + throw new Error('Database does not support onDuplicateKeyUpdate'); + } + await insertQuery.onDuplicateKeyUpdate({ set: updateFields }); + } else { + // SQLite and PostgreSQL use ON CONFLICT DO UPDATE + if (!insertQuery.onConflictDoUpdate) { + throw new Error('Database does not support onConflictDoUpdate'); + } + await insertQuery.onConflictDoUpdate({ + target: this.table.id, + set: updateFields, + }); + } + } + + /** + * Loads a task from the database by ID. + * + * @param taskId The ID of the task to load. + * @returns The task if found, or undefined if not found. + */ + async load(taskId: string): Promise { + const results = (await this.db + .select() + .from(this.table) + .where(eq(this.table.id, taskId))) as TaskRow[]; + + if (results.length === 0) { + return undefined; + } + + return this.rowToTask(results[0]); + } + + /** + * Deletes a task from the database by ID. + * + * @param taskId The ID of the task to delete. + */ + async delete(taskId: string): Promise { + await this.db.delete(this.table).where(eq(this.table.id, taskId)); + } + + /** + * Converts a Task object to a database row. + */ + private taskToRow(task: Task) { + return { + id: task.id, + contextId: task.contextId, + kind: task.kind, + status: task.status, + artifacts: task.artifacts ?? null, + history: task.history ?? null, + metadata: task.metadata ?? null, + }; + } + + /** + * Converts a database row to a Task object. + */ + private rowToTask(row: TaskRow): Task { + return { + id: row.id, + contextId: row.contextId, + kind: row.kind as 'task', + status: row.status, + ...(row.artifacts && { artifacts: row.artifacts }), + ...(row.history && { history: row.history }), + ...(row.metadata && { metadata: row.metadata }), + }; + } +} diff --git a/src/server/drizzle/index.ts b/src/server/drizzle/index.ts new file mode 100644 index 00000000..52731cc6 --- /dev/null +++ b/src/server/drizzle/index.ts @@ -0,0 +1,46 @@ +/** + * Drizzle ORM integration for persistent task storage. + * + * This module provides a database-backed implementation of the TaskStore interface + * using Drizzle ORM, supporting PostgreSQL, MySQL, and SQLite databases. + * + * @example + * ```typescript + * // SQLite example + * import { drizzle } from 'drizzle-orm/better-sqlite3'; + * import Database from 'better-sqlite3'; + * import { DatabaseTaskStore, sqliteTasks } from '@a2a-js/sdk/server/drizzle'; + * + * const sqlite = new Database('tasks.db'); + * const db = drizzle(sqlite); + * + * const taskStore = new DatabaseTaskStore({ + * db, + * table: sqliteTasks, + * dialect: 'sqlite', + * }); + * ``` + * + * @example + * ```typescript + * // PostgreSQL example + * import { drizzle } from 'drizzle-orm/node-postgres'; + * import { Pool } from 'pg'; + * import { DatabaseTaskStore, pgTasks } from '@a2a-js/sdk/server/drizzle'; + * + * const pool = new Pool({ connectionString: process.env.DATABASE_URL }); + * const db = drizzle(pool); + * + * const taskStore = new DatabaseTaskStore({ + * db, + * table: pgTasks, + * dialect: 'postgresql', + * }); + * ``` + */ + +export { DatabaseTaskStore } from './database_task_store.js'; +export type { DrizzleDatabase, DatabaseTaskStoreOptions } from './database_task_store.js'; + +export { sqliteTasks, pgTasks, mysqlTasks } from './schema.js'; +export type { SqliteTasksTable, PgTasksTable, MysqlTasksTable, TasksTable } from './schema.js'; diff --git a/src/server/drizzle/schema.ts b/src/server/drizzle/schema.ts new file mode 100644 index 00000000..2a74746d --- /dev/null +++ b/src/server/drizzle/schema.ts @@ -0,0 +1,72 @@ +/** + * Drizzle ORM schema for persistent task storage. + * This schema defines the tasks table used by DatabaseTaskStore. + */ +import { sqliteTable, text, integer } from 'drizzle-orm/sqlite-core'; +import { pgTable, text as pgText, jsonb, timestamp } from 'drizzle-orm/pg-core'; +import { mysqlTable, text as mysqlText, json, timestamp as mysqlTimestamp } from 'drizzle-orm/mysql-core'; +import type { TaskStatus, Artifact, Message1 } from '../../types.js'; + +/** + * SQLite tasks table schema. + * Stores task data with JSON columns for complex nested objects. + */ +export const sqliteTasks = sqliteTable('tasks', { + id: text('id').primaryKey(), + contextId: text('context_id').notNull(), + kind: text('kind').notNull().$type<'task'>(), + status: text('status', { mode: 'json' }).notNull().$type(), + artifacts: text('artifacts', { mode: 'json' }).$type(), + history: text('history', { mode: 'json' }).$type(), + metadata: text('metadata', { mode: 'json' }).$type | null>(), + createdAt: integer('created_at', { mode: 'timestamp_ms' }) + .notNull() + .$defaultFn(() => new Date()), + updatedAt: integer('updated_at', { mode: 'timestamp_ms' }) + .notNull() + .$defaultFn(() => new Date()) + .$onUpdate(() => new Date()), +}); + +/** + * PostgreSQL tasks table schema. + * Uses JSONB for efficient JSON storage and querying. + */ +export const pgTasks = pgTable('tasks', { + id: pgText('id').primaryKey(), + contextId: pgText('context_id').notNull(), + kind: pgText('kind').notNull().$type<'task'>(), + status: jsonb('status').notNull().$type(), + artifacts: jsonb('artifacts').$type(), + history: jsonb('history').$type(), + metadata: jsonb('metadata').$type | null>(), + createdAt: timestamp('created_at', { mode: 'date' }).notNull().defaultNow(), + updatedAt: timestamp('updated_at', { mode: 'date' }) + .notNull() + .defaultNow() + .$onUpdate(() => new Date()), +}); + +/** + * MySQL tasks table schema. + * Uses JSON columns for complex nested objects. + */ +export const mysqlTasks = mysqlTable('tasks', { + id: mysqlText('id').primaryKey(), + contextId: mysqlText('context_id').notNull(), + kind: mysqlText('kind').notNull().$type<'task'>(), + status: json('status').notNull().$type(), + artifacts: json('artifacts').$type(), + history: json('history').$type(), + metadata: json('metadata').$type | null>(), + createdAt: mysqlTimestamp('created_at', { mode: 'date' }).notNull().defaultNow(), + updatedAt: mysqlTimestamp('updated_at', { mode: 'date' }) + .notNull() + .defaultNow() + .$onUpdate(() => new Date()), +}); + +export type SqliteTasksTable = typeof sqliteTasks; +export type PgTasksTable = typeof pgTasks; +export type MysqlTasksTable = typeof mysqlTasks; +export type TasksTable = SqliteTasksTable | PgTasksTable | MysqlTasksTable; diff --git a/src/server/store.ts b/src/server/store.ts index 48f52c9e..fd9b1d9d 100644 --- a/src/server/store.ts +++ b/src/server/store.ts @@ -19,6 +19,13 @@ export interface TaskStore { * @returns A promise resolving to an object containing the Task, or undefined if not found. */ load(taskId: string): Promise; + + /** + * Deletes a task by task ID. + * @param taskId The ID of the task to delete. + * @returns A promise resolving when the delete operation is complete. + */ + delete(taskId: string): Promise; } // ======================== @@ -39,4 +46,8 @@ export class InMemoryTaskStore implements TaskStore { // Store copies to prevent internal mutation if caller reuses objects this.store.set(task.id, { ...task }); } + + async delete(taskId: string): Promise { + this.store.delete(taskId); + } } diff --git a/tsup.config.ts b/tsup.config.ts index 24030b81..4abaa5fc 100644 --- a/tsup.config.ts +++ b/tsup.config.ts @@ -5,6 +5,7 @@ export default defineConfig({ 'src/index.ts', 'src/server/index.ts', 'src/server/express/index.ts', + 'src/server/drizzle/index.ts', 'src/client/index.ts', ], format: ['esm', 'cjs'],