@falai/agent 2.2.4 → 2.4.1

package/src/adapters/PostgreSQLAdapter.ts

@@ -25,7 +25,9 @@ import type {
   MessageData,
   SessionStatus,
   PersistenceAdapter,
+  SessionUpdateOptions,
 } from "../types";
+import { SessionConflictError } from "../types/errors";
 import { createSessionId } from "../utils";
 
 /**
@@ -133,11 +135,18 @@ export class PostgreSQLAdapter<TData = Record<string, unknown>> implements Persi
         message_count INTEGER DEFAULT 0,
         last_message_at TIMESTAMP,
         completed_at TIMESTAMP,
+        version INTEGER,
         created_at TIMESTAMP DEFAULT NOW(),
         updated_at TIMESTAMP DEFAULT NOW()
       )
     `);
 
+    // Migration: add the optimistic-concurrency version column to tables
+    // created by pre-2.4 versions
+    await this.client.query(`
+      ALTER TABLE ${sessionTable} ADD COLUMN IF NOT EXISTS version INTEGER
+    `);
+
     await this.client.query(`
       CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON ${sessionTable}(user_id);
       CREATE INDEX IF NOT EXISTS idx_sessions_status ON ${sessionTable}(status);
@@ -187,9 +196,9 @@ class PostgreSQLSessionRepository<TData = Record<string, unknown>>
     const now = new Date();
 
     const result = await this.client.query<SessionData<TData>>(
-      `INSERT INTO ${this.tableName} 
-       (id, user_id, agent_name, status, collected_data, message_count, created_at, updated_at)
-       VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+      `INSERT INTO ${this.tableName}
+       (id, user_id, agent_name, status, collected_data, message_count, version, created_at, updated_at)
+       VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
        RETURNING *`,
       [
         id,
@@ -198,6 +207,7 @@ class PostgreSQLSessionRepository<TData = Record<string, unknown>>
         data.status || "active",
         JSON.stringify(data.collectedData || {}),
         data.messageCount || 0,
+        data.version ?? 1,
         now,
         now,
       ]
@@ -244,7 +254,8 @@ class PostgreSQLSessionRepository<TData = Record<string, unknown>>
 
   async update(
     id: string,
-    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>,
+    options?: SessionUpdateOptions
   ): Promise<SessionData<TData> | null> {
     const fields: string[] = [];
     const values: unknown[] = [];
@@ -282,10 +293,39 @@ class PostgreSQLSessionRepository<TData = Record<string, unknown>>
     fields.push(`updated_at = $${paramIndex++}`);
     values.push(new Date());
 
+    if (options?.expectedVersion !== undefined) {
+      const expectedVersion = options.expectedVersion;
+
+      // Compare-and-swap in a single statement: rows without a stored version
+      // (pre-2.4) are accepted and adopt expectedVersion as their base
+      fields.push(`version = COALESCE(version, $${paramIndex++}) + 1`);
+      values.push(expectedVersion);
+
+      const idParam = paramIndex++;
+      values.push(id);
+      const versionParam = paramIndex++;
+      values.push(expectedVersion);
+
+      const result = await this.client.query<SessionData<TData>>(
+        `UPDATE ${this.tableName}
+         SET ${fields.join(", ")}
+         WHERE id = $${idParam} AND (version IS NULL OR version = $${versionParam})
+         RETURNING *`,
+        values
+      );
+
+      if (result.rows[0]) return result.rows[0];
+
+      const existing = await this.findById(id);
+      if (!existing) return null;
+      throw new SessionConflictError(id, expectedVersion, existing.version);
+    }
+
+    fields.push(`version = COALESCE(version, 0) + 1`);
     values.push(id);
 
     const result = await this.client.query<SessionData<TData>>(
-      `UPDATE ${this.tableName} 
+      `UPDATE ${this.tableName}
        SET ${fields.join(", ")}
        WHERE id = $${paramIndex}
        RETURNING *`,

package/src/adapters/PrismaAdapter.ts

@@ -14,7 +14,9 @@ import type {
   MessageRole,
   Event,
   CreateSessionData,
+  SessionUpdateOptions,
 } from "../types";
+import { SessionConflictError } from "../types/errors";
 import { logger } from '../utils'
 import { createSessionId } from '../utils';
 
@@ -202,6 +204,13 @@ class PrismaSessionRepository<TData = Record<string, unknown>>
   private prisma: PrismaClient;
   private tableName: string;
   private fieldMap: Partial<Record<keyof SessionData<TData>, string>>;
+  /**
+   * Whether the user's Prisma model has a `version` column. Unknown until the
+   * first write; when Prisma rejects the field, writes fall back to omitting
+   * it so models without the column keep working (optimistic locking is then
+   * inactive — add `version Int?` to the model to enable it).
+   */
+  private versionSupported?: boolean;
 
   constructor(
     prismaClient: PrismaClient,
@@ -213,6 +222,34 @@ class PrismaSessionRepository<TData = Record<string, unknown>>
     this.fieldMap = fieldMappings || {};
   }
 
+  /**
+   * Run a write that includes the `version` field, retrying without it once
+   * if the model evidently lacks the column. The outcome is remembered.
+   */
+  private async writeWithVersionFallback<T>(
+    write: (includeVersion: boolean) => Promise<T>
+  ): Promise<T> {
+    if (this.versionSupported === false) {
+      return write(false);
+    }
+    try {
+      const result = await write(true);
+      this.versionSupported = true;
+      return result;
+    } catch (error) {
+      const versionField = this.fieldMap.version || "version";
+      if (
+        this.versionSupported === undefined &&
+        error instanceof Error &&
+        error.message.includes(versionField)
+      ) {
+        this.versionSupported = false;
+        return write(false);
+      }
+      throw error;
+    }
+  }
+
   /**
    * Map our standard field names to custom schema field names
    */
@@ -254,6 +291,7 @@ class PrismaSessionRepository<TData = Record<string, unknown>>
       messageCount: unmapped.messageCount as number | undefined,
       lastMessageAt: unmapped.lastMessageAt as Date | undefined,
       completedAt: unmapped.completedAt as Date | undefined,
+      version: (unmapped.version as number | null) ?? undefined,
       createdAt: unmapped.createdAt as Date,
       updatedAt: unmapped.updatedAt as Date,
     };
@@ -264,14 +302,14 @@ class PrismaSessionRepository<TData = Record<string, unknown>>
   }
 
   async create(data: CreateSessionData<TData>): Promise<SessionData<TData>> {
-    const mapped = this.mapFields({
-      ...data,
-      id:
-        data.id ||
-        createSessionId(),
-    });
-    const result = await this.getModel().create({
-      data: mapped,
+    const id = data.id || createSessionId();
+    const result = await this.writeWithVersionFallback((includeVersion) => {
+      const mapped = this.mapFields(
+        includeVersion
+          ? { ...data, id, version: data.version ?? 1 }
+          : { ...data, id }
+      );
+      return this.getModel().create({ data: mapped });
     });
     return this.unmapFields(result);
   }
@@ -314,15 +352,43 @@ class PrismaSessionRepository<TData = Record<string, unknown>>
 
   async update(
     id: string,
-    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>,
+    options?: SessionUpdateOptions
   ): Promise<SessionData<TData> | null> {
-    const mapped = this.mapFields(data as Record<string, unknown>);
-    const result = await this.getModel().update({
-      where: { [this.fieldMap.id || "id"]: id },
-      data: {
-        ...mapped,
-        [this.fieldMap.updatedAt || "updatedAt"]: new Date(),
-      },
+    // Read-check-write: the Prisma model is user-defined, so we cannot emit a
+    // conditional UPDATE generically. Add `version Int?` to your Prisma model
+    // for optimistic locking to be effective.
+    const existing = await this.findById(id);
+    if (!existing) return null;
+
+    if (
+      options?.expectedVersion !== undefined &&
+      existing.version !== undefined &&
+      existing.version !== options.expectedVersion
+    ) {
+      throw new SessionConflictError(
+        id,
+        options.expectedVersion,
+        existing.version
+      );
+    }
+
+    const result = await this.writeWithVersionFallback((includeVersion) => {
+      const mapped = this.mapFields(
+        includeVersion
+          ? {
+            ...(data as Record<string, unknown>),
+            version: (existing.version ?? options?.expectedVersion ?? 0) + 1,
+          }
+          : { ...(data as Record<string, unknown>) }
+      );
+      return this.getModel().update({
+        where: { [this.fieldMap.id || "id"]: id },
+        data: {
+          ...mapped,
+          [this.fieldMap.updatedAt || "updatedAt"]: new Date(),
+        },
+      });
     });
     return this.unmapFields(result);
   }

package/src/adapters/RedisAdapter.ts

@@ -11,7 +11,9 @@ import type {
   SessionRepository,
   SessionStatus,
   CollectedStateData,
+  SessionUpdateOptions,
 } from "../types";
+import { SessionConflictError } from "../types/errors";
 import { createSessionId, logger } from "../utils";
 
 /**
@@ -139,6 +141,7 @@ class RedisSessionRepository<TData = Record<string, unknown>>
       updatedAt: now,
       status: data.status || "active",
       messageCount: data.messageCount || 0,
+      version: data.version ?? 1,
     };
 
     await this.redis.setex(this.getKey(id), this.ttl, JSON.stringify(session));
@@ -201,14 +204,29 @@ class RedisSessionRepository<TData = Record<string, unknown>>
 
   async update(
     id: string,
-    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>,
+    options?: SessionUpdateOptions
   ): Promise<SessionData<TData> | null> {
     const existing = await this.findById(id);
     if (!existing) return null;
 
+    // Check-then-set on the stored JSON — not fully atomic (no WATCH/MULTI)
+    if (
+      options?.expectedVersion !== undefined &&
+      existing.version !== undefined &&
+      existing.version !== options.expectedVersion
+    ) {
+      throw new SessionConflictError(
+        id,
+        options.expectedVersion,
+        existing.version
+      );
+    }
+
     const updated: SessionData<TData> = {
       ...existing,
       ...data,
+      version: (existing.version ?? options?.expectedVersion ?? 0) + 1,
       updatedAt: new Date(),
     };
 

package/src/adapters/SQLiteAdapter.ts

@@ -29,7 +29,9 @@ import type {
   PersistenceAdapter,
   CollectedStateData,
   CreateSessionData,
+  SessionUpdateOptions,
 } from "../types";
+import { SessionConflictError } from "../types/errors";
 import { createSessionId } from "../utils";
 
 /**
@@ -125,11 +127,21 @@ export class SQLiteAdapter<TData = Record<string, unknown>> implements Persisten
         message_count INTEGER DEFAULT 0,
         last_message_at TEXT,
         completed_at TEXT,
+        version INTEGER,
         created_at TEXT NOT NULL,
         updated_at TEXT NOT NULL
       )
     `);
 
+    // Migration: add the optimistic-concurrency version column to tables
+    // created by pre-2.4 versions
+    const sessionColumns = this.db
+      .prepare(`PRAGMA table_info(${sessionTable})`)
+      .all();
+    if (!sessionColumns.some((column) => column.name === "version")) {
+      this.db.exec(`ALTER TABLE ${sessionTable} ADD COLUMN version INTEGER`);
+    }
+
     // Create indexes for sessions
     this.db.exec(`
       CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON ${sessionTable}(user_id);
@@ -183,14 +195,15 @@ class SQLiteSessionRepository<TData = Record<string, unknown>>
       id,
       status: data.status || "active",
       messageCount: data.messageCount || 0,
+      version: data.version ?? 1,
       createdAt: now,
       updatedAt: now,
     };
 
     const stmt = this.db.prepare(`
-      INSERT INTO ${this.tableName} 
-      (id, user_id, agent_name, status, collected_data, message_count, created_at, updated_at)
-      VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+      INSERT INTO ${this.tableName}
+      (id, user_id, agent_name, status, collected_data, message_count, version, created_at, updated_at)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
     `);
 
     stmt.run(
@@ -200,6 +213,7 @@ class SQLiteSessionRepository<TData = Record<string, unknown>>
       session.status,
       JSON.stringify(data.collectedData || {}),
       session.messageCount,
+      session.version,
       now.toISOString(),
       now.toISOString()
     );
@@ -239,7 +253,8 @@ class SQLiteSessionRepository<TData = Record<string, unknown>>
 
   async update(
     id: string,
-    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>
+    data: Partial<Omit<SessionData<TData>, "id" | "createdAt">>,
+    options?: SessionUpdateOptions
   ): Promise<SessionData<TData> | null> {
     const fields: string[] = [];
     const values: unknown[] = [];
@@ -276,10 +291,36 @@ class SQLiteSessionRepository<TData = Record<string, unknown>>
     fields.push("updated_at = ?");
     values.push(new Date().toISOString());
 
+    if (options?.expectedVersion !== undefined) {
+      const expectedVersion = options.expectedVersion;
+
+      // Compare-and-swap in a single statement: rows without a stored version
+      // (pre-2.4) are accepted and adopt expectedVersion as their base
+      fields.push("version = COALESCE(version, ?) + 1");
+      values.push(expectedVersion);
+      values.push(id, expectedVersion);
+
+      const stmt = this.db.prepare(
+        `UPDATE ${this.tableName}
+         SET ${fields.join(", ")}
+         WHERE id = ? AND (version IS NULL OR version = ?)`
+      );
+
+      const result = stmt.run(...values);
+      if (result.changes === 0) {
+        const existing = await this.findById(id);
+        if (!existing) return null;
+        throw new SessionConflictError(id, expectedVersion, existing.version);
+      }
+
+      return await this.findById(id);
+    }
+
+    fields.push("version = COALESCE(version, 0) + 1");
     values.push(id);
 
     const stmt = this.db.prepare(
-      `UPDATE ${this.tableName} 
+      `UPDATE ${this.tableName}
        SET ${fields.join(", ")}
        WHERE id = ?`
     );
@@ -358,6 +399,7 @@ class SQLiteSessionRepository<TData = Record<string, unknown>>
       completedAt: row.completed_at
         ? new Date(row.completed_at as string)
         : undefined,
+      version: (row.version as number | null) ?? undefined,
       createdAt: new Date(row.created_at as string),
       updatedAt: new Date(row.updated_at as string),
     };