Design: improve select and regenerate flow

.changeset/wise-cats-resolve.md

@@ -0,0 +1,5 @@
+---
+"@agent-native/core": patch
+---
+
+Add portable transactional multi-key application-state compare-and-set operations, including atomic D1 batch support and missing-key creation for race-safe UI workflows.

packages/core/src/application-state/index.ts

@@ -5,8 +5,10 @@ export {
   appStatePut,
   appStateDelete,
   appStateCompareAndSet,
+  appStateCompareAndSetMany,
   appStateList,
   appStateDeleteByPrefix,
+  type AppStateCompareAndSetOperation,
 } from "./store.js";
 
 // Emitter (for SSE wiring)
@@ -35,6 +37,7 @@ export {
   writeAppState,
   deleteAppState,
   compareAndSetAppState,
+  compareAndSetManyAppState,
   listAppState,
   deleteAppStateByPrefix,
   readAppStateForCurrentTab,

packages/core/src/application-state/script-helpers.spec.ts

@@ -5,6 +5,7 @@ const mockAppStateGet = vi.fn();
 const mockAppStatePut = vi.fn();
 const mockAppStateDelete = vi.fn();
 const mockAppStateCompareAndSet = vi.fn();
+const mockAppStateCompareAndSetMany = vi.fn();
 const mockAppStateList = vi.fn();
 const mockAppStateDeleteByPrefix = vi.fn();
 
@@ -13,6 +14,8 @@ vi.mock("./store.js", () => ({
   appStatePut: (...args: any[]) => mockAppStatePut(...args),
   appStateDelete: (...args: any[]) => mockAppStateDelete(...args),
   appStateCompareAndSet: (...args: any[]) => mockAppStateCompareAndSet(...args),
+  appStateCompareAndSetMany: (...args: any[]) =>
+    mockAppStateCompareAndSetMany(...args),
   appStateList: (...args: any[]) => mockAppStateList(...args),
   appStateDeleteByPrefix: (...args: any[]) =>
     mockAppStateDeleteByPrefix(...args),
@@ -144,6 +147,33 @@ describe("application-state script-helpers", () => {
     });
   });
 
+  describe("compareAndSetManyAppState", () => {
+    it("delegates the complete transition with the resolved session ID", async () => {
+      process.env.AGENT_USER_EMAIL = "alice@test.com";
+      const { compareAndSetManyAppState } = await import("./script-helpers.js");
+      mockAppStateCompareAndSetMany.mockResolvedValue(true);
+      const operations = [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: null,
+        },
+        {
+          key: "proposal",
+          expectedValue: { proposalId: "p1" },
+          nextValue: null,
+        },
+      ];
+
+      await expect(compareAndSetManyAppState(operations)).resolves.toBe(true);
+      expect(mockAppStateCompareAndSetMany).toHaveBeenCalledWith(
+        "alice@test.com",
+        operations,
+        { requestSource: "agent" },
+      );
+    });
+  });
+
   describe("listAppState", () => {
     it("delegates to appStateList with prefix", async () => {
       process.env.AGENT_USER_EMAIL = "alice@test.com";

packages/core/src/application-state/script-helpers.ts

@@ -18,8 +18,10 @@ import {
   appStatePut,
   appStateDelete,
   appStateCompareAndSet,
+  appStateCompareAndSetMany,
   appStateList,
   appStateDeleteByPrefix,
+  type AppStateCompareAndSetOperation,
 } from "./store.js";
 
 /**
@@ -74,7 +76,7 @@ export async function deleteAppState(key: string): Promise<boolean> {
 
 export async function compareAndSetAppState(
   key: string,
-  expectedValue: Record<string, unknown>,
+  expectedValue: Record<string, unknown> | null,
   nextValue: Record<string, unknown> | null,
 ): Promise<boolean> {
   const sessionId = await resolveSessionId();
@@ -83,6 +85,15 @@ export async function compareAndSetAppState(
   });
 }
 
+export async function compareAndSetManyAppState(
+  operations: readonly AppStateCompareAndSetOperation[],
+): Promise<boolean> {
+  const sessionId = await resolveSessionId();
+  return appStateCompareAndSetMany(sessionId, operations, {
+    requestSource: "agent",
+  });
+}
+
 export async function listAppState(
   prefix: string,
 ): Promise<Array<{ key: string; value: Record<string, unknown> }>> {

packages/core/src/application-state/store.spec.ts

@@ -17,14 +17,39 @@ const rawClient = {
     const info = stmt.run(...args);
     return { rows: [], rowsAffected: info.changes };
   }),
+  transaction: vi.fn(async <T>(fn: (tx: typeof rawClient) => Promise<T>) => {
+    sqlite.exec("BEGIN IMMEDIATE");
+    try {
+      const result = await fn(rawClient);
+      sqlite.exec("COMMIT");
+      return result;
+    } catch (error) {
+      sqlite.exec("ROLLBACK");
+      throw error;
+    }
+  }),
 };
+const atomicBatch = vi.fn(
+  async (statements: readonly (string | { sql: string; args?: unknown[] })[]) =>
+    rawClient.transaction(async (tx) => {
+      const results = [];
+      for (const statement of statements) {
+        results.push(await tx.execute(statement));
+      }
+      return results;
+    }),
+);
 
 const emitAppStateChange = vi.fn();
 const emitAppStateDelete = vi.fn();
-const dbMockState = vi.hoisted(() => ({ localDatabase: true }));
+const dbMockState = vi.hoisted(() => ({
+  localDatabase: true,
+  dialect: "sqlite" as "sqlite" | "postgres" | "d1",
+}));
 
 vi.mock("../db/client.js", () => ({
-  getDbExec: () => rawClient,
+  getDbExec: () => ({ ...rawClient, atomicBatch }),
+  getDialect: () => dbMockState.dialect,
   intType: () => "INTEGER",
   isConnectionError: () => false,
   isLocalDatabase: () => dbMockState.localDatabase,
@@ -41,6 +66,7 @@ const {
   appStateGet,
   appStateGetMany,
   appStateCompareAndSet,
+  appStateCompareAndSetMany,
   appStateList,
   appStateDeleteByPrefix,
 } = await import("./store.js");
@@ -62,6 +88,7 @@ afterEach(() => {
   sqlite.close();
   vi.clearAllMocks();
   dbMockState.localDatabase = true;
+  dbMockState.dialect = "sqlite";
 });
 
 describe("application-state store", () => {
@@ -205,6 +232,163 @@ describe("application-state store", () => {
     );
   });
 
+  it("atomically creates a value only while its key is absent", async () => {
+    await expect(
+      appStateCompareAndSet(SESSION, "rewrite", null, { repromptId: "r1" }),
+    ).resolves.toBe(true);
+    await expect(
+      appStateCompareAndSet(SESSION, "rewrite", null, { repromptId: "r2" }),
+    ).resolves.toBe(false);
+    expect(await appStateGet(SESSION, "rewrite")).toEqual({ repromptId: "r1" });
+  });
+
+  it("rolls back every key when one operation in a multi-key CAS misses", async () => {
+    await appStatePut(SESSION, "pending", { repromptId: "r1" });
+    await appStatePut(SESSION, "proposal", { proposalId: "p1" });
+
+    await expect(
+      appStateCompareAndSetMany(SESSION, [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: null,
+        },
+        {
+          key: "proposal",
+          expectedValue: { proposalId: "stale" },
+          nextValue: null,
+        },
+      ]),
+    ).resolves.toBe(false);
+    expect(await appStateGet(SESSION, "pending")).toEqual({ repromptId: "r1" });
+    expect(await appStateGet(SESSION, "proposal")).toEqual({
+      proposalId: "p1",
+    });
+
+    await expect(
+      appStateCompareAndSetMany(SESSION, [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: null,
+        },
+        {
+          key: "proposal",
+          expectedValue: { proposalId: "p1" },
+          nextValue: null,
+        },
+      ]),
+    ).resolves.toBe(true);
+    expect(await appStateGet(SESSION, "pending")).toBeNull();
+    expect(await appStateGet(SESSION, "proposal")).toBeNull();
+  });
+
+  it("uses an atomic D1 batch and rolls back every mutation on a mismatch", async () => {
+    dbMockState.dialect = "d1";
+    await appStatePut(SESSION, "pending", { repromptId: "r1" });
+    await appStatePut(SESSION, "proposal", { proposalId: "p1" });
+    atomicBatch.mockClear();
+    emitAppStateChange.mockClear();
+    emitAppStateDelete.mockClear();
+
+    await expect(
+      appStateCompareAndSetMany(SESSION, [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: null,
+        },
+        {
+          key: "proposal",
+          expectedValue: { proposalId: "stale" },
+          nextValue: null,
+        },
+      ]),
+    ).resolves.toBe(false);
+    expect(atomicBatch).toHaveBeenCalledTimes(1);
+    const statements = atomicBatch.mock.calls[0]![0];
+    expect(statements.slice(1, 3)).toEqual([
+      expect.objectContaining({ sql: expect.stringContaining("WHERE NOT") }),
+      expect.objectContaining({ sql: expect.stringContaining("WHERE NOT") }),
+    ]);
+    expect(
+      statements.some((statement) =>
+        typeof statement === "string"
+          ? /^(?:BEGIN|COMMIT|ROLLBACK)/i.test(statement)
+          : /^(?:BEGIN|COMMIT|ROLLBACK)/i.test(statement.sql),
+      ),
+    ).toBe(false);
+    expect(await appStateGet(SESSION, "pending")).toEqual({ repromptId: "r1" });
+    expect(await appStateGet(SESSION, "proposal")).toEqual({
+      proposalId: "p1",
+    });
+    expect(emitAppStateChange).not.toHaveBeenCalled();
+    expect(emitAppStateDelete).not.toHaveBeenCalled();
+
+    await expect(
+      appStateCompareAndSetMany(SESSION, [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: null,
+        },
+        {
+          key: "proposal",
+          expectedValue: { proposalId: "p1" },
+          nextValue: null,
+        },
+      ]),
+    ).resolves.toBe(true);
+    expect(await appStateGet(SESSION, "pending")).toBeNull();
+    expect(await appStateGet(SESSION, "proposal")).toBeNull();
+  });
+
+  it("atomically applies mixed D1 update, insert, and delete operations", async () => {
+    dbMockState.dialect = "d1";
+    await appStatePut(SESSION, "pending", { repromptId: "r1" });
+    await appStatePut(SESSION, "prior", { proposalId: "p1" });
+
+    await expect(
+      appStateCompareAndSetMany(SESSION, [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: { repromptId: "r2" },
+        },
+        {
+          key: "current",
+          expectedValue: null,
+          nextValue: { proposalId: "p2" },
+        },
+        {
+          key: "prior",
+          expectedValue: { proposalId: "p1" },
+          nextValue: null,
+        },
+      ]),
+    ).resolves.toBe(true);
+    expect(await appStateGet(SESSION, "pending")).toEqual({ repromptId: "r2" });
+    expect(await appStateGet(SESSION, "current")).toEqual({ proposalId: "p2" });
+    expect(await appStateGet(SESSION, "prior")).toBeNull();
+  });
+
+  it("propagates non-guard D1 batch failures", async () => {
+    dbMockState.dialect = "d1";
+    await appStatePut(SESSION, "pending", { repromptId: "r1" });
+    atomicBatch.mockRejectedValueOnce(new Error("D1 batch unavailable"));
+
+    await expect(
+      appStateCompareAndSetMany(SESSION, [
+        {
+          key: "pending",
+          expectedValue: { repromptId: "r1" },
+          nextValue: null,
+        },
+      ]),
+    ).rejects.toThrow("D1 batch unavailable");
+    expect(await appStateGet(SESSION, "pending")).toEqual({ repromptId: "r1" });
+  });
+
   it("rejects oversized hosted application_state values", async () => {
     dbMockState.localDatabase = false;