@agent-native/core 0.114.12 → 0.114.13

package/src/application-state/store.ts

@@ -1,9 +1,11 @@
 import {
   getDbExec,
+  getDialect,
   isLocalDatabase,
   isConnectionError,
   isPostgres,
   intType,
+  type DbExec,
 } from "../db/client.js";
 import { ensureIndexExists, ensureTableExists } from "../db/ddl-guard.js";
 import { widenIntColumnsToBigInt } from "../db/widen-columns.js";
@@ -194,39 +196,238 @@ export async function appStateDelete(
 export async function appStateCompareAndSet(
   sessionId: string,
   key: string,
-  expectedValue: Record<string, unknown>,
+  expectedValue: Record<string, unknown> | null,
   nextValue: Record<string, unknown> | null,
   options?: StoreWriteOptions,
 ): Promise<boolean> {
   await ensureTable();
   const client = getDbExec();
+  const changed = await executeAppStateCompareAndSet(
+    client,
+    sessionId,
+    key,
+    expectedValue,
+    nextValue,
+  );
+  if (changed) {
+    if (nextValue === null) {
+      emitAppStateDelete(key, options?.requestSource, sessionId);
+    } else {
+      emitAppStateChange(key, options?.requestSource, sessionId);
+    }
+  }
+  return changed;
+}
+
+export interface AppStateCompareAndSetOperation {
+  key: string;
+  expectedValue: Record<string, unknown> | null;
+  nextValue: Record<string, unknown> | null;
+}
+
+const APP_STATE_CAS_MISMATCH = Symbol("app-state-cas-mismatch");
+
+async function executeAppStateCompareAndSet(
+  client: DbExec,
+  sessionId: string,
+  key: string,
+  expectedValue: Record<string, unknown> | null,
+  nextValue: Record<string, unknown> | null,
+): Promise<boolean> {
+  const statement = buildAppStateCompareAndSetStatement(
+    sessionId,
+    key,
+    expectedValue,
+    nextValue,
+  );
+  const result = await client.execute(statement);
+  return result.rowsAffected > 0;
+}
+
+function buildAppStateCompareAndSetStatement(
+  sessionId: string,
+  key: string,
+  expectedValue: Record<string, unknown> | null,
+  nextValue: Record<string, unknown> | null,
+): { sql: string; args: unknown[] } {
+  if (expectedValue === null) {
+    if (nextValue === null) {
+      throw new Error(
+        "Application state CAS cannot replace absence with absence.",
+      );
+    }
+    const next = serializeAppStateValue(key, nextValue);
+    return {
+      sql: isPostgres()
+        ? `INSERT INTO application_state (session_id, key, value, updated_at) SELECT ?, ?, ?, ? WHERE NOT EXISTS (SELECT 1 FROM application_state WHERE session_id = ? AND key = ?) ON CONFLICT (session_id, key) DO NOTHING`
+        : `INSERT OR IGNORE INTO application_state (session_id, key, value, updated_at) SELECT ?, ?, ?, ? WHERE NOT EXISTS (SELECT 1 FROM application_state WHERE session_id = ? AND key = ?)`,
+      args: [sessionId, key, next, Date.now(), sessionId, key],
+    };
+  }
+
   const expected = JSON.stringify(expectedValue);
   if (nextValue === null) {
-    const result = await client.execute({
+    return {
       sql: `DELETE FROM application_state WHERE session_id = ? AND key = ? AND value = ?`,
       args: [sessionId, key, expected],
-    });
-    const changed = result.rowsAffected > 0;
-    if (changed) emitAppStateDelete(key, options?.requestSource, sessionId);
-    return changed;
+    };
   }
 
-  const next = JSON.stringify(nextValue);
+  const next = serializeAppStateValue(key, nextValue);
+  return {
+    sql: `UPDATE application_state SET value = ?, updated_at = ? WHERE session_id = ? AND key = ? AND value = ?`,
+    args: [next, Date.now(), sessionId, key, expected],
+  };
+}
+
+function buildD1CompareAndSetGuard(
+  sessionId: string,
+  guardKey: string,
+  operation: AppStateCompareAndSetOperation,
+): { sql: string; args: unknown[] } {
+  const expected = operation.expectedValue;
+  const condition =
+    expected === null
+      ? "NOT EXISTS (SELECT 1 FROM application_state WHERE session_id = ? AND key = ?)"
+      : "EXISTS (SELECT 1 FROM application_state WHERE session_id = ? AND key = ? AND value = ?)";
+  return {
+    sql: `INSERT INTO application_state (session_id, key, value, updated_at) SELECT ?, ?, '{}', ? WHERE NOT (${condition})`,
+    args:
+      expected === null
+        ? [sessionId, guardKey, Date.now(), sessionId, operation.key]
+        : [
+            sessionId,
+            guardKey,
+            Date.now(),
+            sessionId,
+            operation.key,
+            JSON.stringify(expected),
+          ],
+  };
+}
+
+function isD1CompareAndSetMismatch(error: unknown): boolean {
+  const message = error instanceof Error ? error.message : String(error);
+  return /unique constraint failed:\s*application_state\.session_id,\s*application_state\.key/i.test(
+    message,
+  );
+}
+
+function serializeAppStateValue(
+  key: string,
+  value: Record<string, unknown>,
+): string {
+  const serialized = JSON.stringify(value);
   if (
     !isLocalDatabase() &&
-    Buffer.byteLength(next, "utf8") > MAX_HOSTED_APP_STATE_VALUE_BYTES
+    Buffer.byteLength(serialized, "utf8") > MAX_HOSTED_APP_STATE_VALUE_BYTES
   ) {
     throw new Error(
       `application_state value "${key}" is too large for hosted SQL storage. Store large files, base64, or blobs in file storage and write only a URL or handle.`,
     );
   }
-  const result = await client.execute({
-    sql: `UPDATE application_state SET value = ?, updated_at = ? WHERE session_id = ? AND key = ? AND value = ?`,
-    args: [next, Date.now(), sessionId, key, expected],
-  });
-  const changed = result.rowsAffected > 0;
-  if (changed) emitAppStateChange(key, options?.requestSource, sessionId);
-  return changed;
+  return serialized;
+}
+
+export async function appStateCompareAndSetMany(
+  sessionId: string,
+  operations: readonly AppStateCompareAndSetOperation[],
+  options?: StoreWriteOptions,
+): Promise<boolean> {
+  if (operations.length === 0) return true;
+  const keys = new Set(operations.map(({ key }) => key));
+  if (keys.size !== operations.length) {
+    throw new Error("Application state multi-key CAS requires unique keys.");
+  }
+  for (const { key, nextValue } of operations) {
+    if (nextValue) serializeAppStateValue(key, nextValue);
+  }
+  const orderedOperations = [...operations].sort((a, b) =>
+    a.key.localeCompare(b.key),
+  );
+
+  await ensureTable();
+  const client = getDbExec();
+  if (getDialect() === "d1") {
+    if (!client.atomicBatch) {
+      throw new Error(
+        "D1 application-state CAS requires atomic batch support.",
+      );
+    }
+    const guardKey = `__agent_native_cas_guard__:${Date.now()}:${Math.random().toString(36).slice(2)}`;
+    const guardInsert = {
+      sql: `INSERT INTO application_state (session_id, key, value, updated_at) VALUES (?, ?, '{}', ?)`,
+      args: [sessionId, guardKey, Date.now()],
+    };
+    const guards = orderedOperations.map((operation) =>
+      buildD1CompareAndSetGuard(sessionId, guardKey, operation),
+    );
+    const mutations = orderedOperations.map((operation) =>
+      buildAppStateCompareAndSetStatement(
+        sessionId,
+        operation.key,
+        operation.expectedValue,
+        operation.nextValue,
+      ),
+    );
+    let results;
+    try {
+      results = await client.atomicBatch([
+        guardInsert,
+        ...guards,
+        ...mutations,
+        {
+          sql: `DELETE FROM application_state WHERE session_id = ? AND key = ?`,
+          args: [sessionId, guardKey],
+        },
+      ]);
+    } catch (error) {
+      if (isD1CompareAndSetMismatch(error)) return false;
+      throw error;
+    }
+    const mutationResults = results.slice(
+      1 + guards.length,
+      1 + guards.length + mutations.length,
+    );
+    if (
+      mutationResults.length !== mutations.length ||
+      mutationResults.some((result) => result.rowsAffected !== 1)
+    ) {
+      throw new Error(
+        "D1 application-state CAS completed without applying every mutation.",
+      );
+    }
+  } else {
+    if (!client.transaction) {
+      throw new Error("Application state multi-key CAS requires transactions.");
+    }
+    try {
+      await client.transaction(async (tx) => {
+        for (const operation of orderedOperations) {
+          const changed = await executeAppStateCompareAndSet(
+            tx,
+            sessionId,
+            operation.key,
+            operation.expectedValue,
+            operation.nextValue,
+          );
+          if (!changed) throw APP_STATE_CAS_MISMATCH;
+        }
+      });
+    } catch (error) {
+      if (error === APP_STATE_CAS_MISMATCH) return false;
+      throw error;
+    }
+  }
+
+  for (const { key, nextValue } of operations) {
+    if (nextValue === null) {
+      emitAppStateDelete(key, options?.requestSource, sessionId);
+    } else {
+      emitAppStateChange(key, options?.requestSource, sessionId);
+    }
+  }
+  return true;
 }
 
 export async function appStateList(

package/src/db/client.ts

@@ -29,6 +29,9 @@ export interface DbExec {
     sql: string | { sql: string; args?: unknown[] },
   ): Promise<{ rows: any[]; rowsAffected: number }>;
   transaction?<T>(fn: (tx: DbExec) => Promise<T>): Promise<T>;
+  atomicBatch?(
+    statements: readonly (string | { sql: string; args?: unknown[] })[],
+  ): Promise<Array<{ rows: any[]; rowsAffected: number }>>;
   /**
    * Release the underlying connection/pool held by this exec.
    * Only non-singleton execs created via `createDbExec()` (e.g. the migration
@@ -1100,7 +1103,17 @@ async function createDbExecInternal(
     };
     return {
       execute,
-      transaction: explicitTransaction(execute),
+      async atomicBatch(statements) {
+        const prepared = statements.map((statement) => {
+          if (typeof statement === "string") return d1.prepare(statement);
+          return d1.prepare(statement.sql).bind(...(statement.args ?? []));
+        });
+        const results = await d1.batch(prepared);
+        return results.map((result: any) => ({
+          rows: result.results || [],
+          rowsAffected: result.meta?.changes ?? 0,
+        }));
+      },
     };
   }
 
@@ -1558,6 +1571,10 @@ export function getDbExec(): DbExec {
       // After init, swap to a sanitizing wrapper around the real client
       const wrapper: DbExec = {
         execute: (s) => _exec!.execute(sanitize(s)),
+        atomicBatch: _exec!.atomicBatch
+          ? (statements) =>
+              _exec!.atomicBatch!(statements.map((s) => sanitize(s)))
+          : undefined,
         transaction: _exec!.transaction
           ? (fn) =>
               _exec!.transaction!((tx) =>
@@ -1582,6 +1599,10 @@ export function getDbExec(): DbExec {
       }
       const wrapper: DbExec = {
         execute: (s) => _exec!.execute(sanitize(s)),
+        atomicBatch: _exec!.atomicBatch
+          ? (statements) =>
+              _exec!.atomicBatch!(statements.map((s) => sanitize(s)))
+          : undefined,
         transaction: _exec!.transaction
           ? (innerFn) =>
               _exec!.transaction!((tx) =>
@@ -1601,8 +1622,30 @@ export function getDbExec(): DbExec {
           }),
         );
       }
+      if (_exec!.atomicBatch) {
+        throw new Error(
+          "This database supports atomic batches, not interactive transactions.",
+        );
+      }
       return explicitTransaction(wrapper.execute)(fn);
     },
+    async atomicBatch(statements) {
+      if (!_initPromise) _initPromise = initClient();
+      try {
+        await _initPromise;
+      } catch (err) {
+        _initPromise = undefined;
+        _exec = undefined;
+        throw err;
+      }
+      if (!_exec!.atomicBatch) {
+        throw new Error("This database does not support atomic batches.");
+      }
+      const batch = (items: typeof statements) =>
+        _exec!.atomicBatch!(items.map((item) => sanitize(item)));
+      Object.assign(proxy, { atomicBatch: batch });
+      return batch(statements);
+    },
   };
   return proxy;
 }