diff --git a/.changeset/wise-cats-resolve.md b/.changeset/wise-cats-resolve.md new file mode 100644 index 0000000000..58b1d890ae --- /dev/null +++ b/.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. diff --git a/packages/core/src/application-state/index.ts b/packages/core/src/application-state/index.ts index b677ccb008..575352e641 100644 --- a/packages/core/src/application-state/index.ts +++ b/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, diff --git a/packages/core/src/application-state/script-helpers.spec.ts b/packages/core/src/application-state/script-helpers.spec.ts index 8dc559bc54..61be03bf85 100644 --- a/packages/core/src/application-state/script-helpers.spec.ts +++ b/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"; diff --git a/packages/core/src/application-state/script-helpers.ts b/packages/core/src/application-state/script-helpers.ts index 76ebf0bd60..aceb89bd2a 100644 --- a/packages/core/src/application-state/script-helpers.ts +++ b/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 { export async function compareAndSetAppState( key: string, - expectedValue: Record, + expectedValue: Record | null, nextValue: Record | null, ): Promise { const sessionId = await resolveSessionId(); @@ -83,6 +85,15 @@ export async function compareAndSetAppState( }); } +export async function compareAndSetManyAppState( + operations: readonly AppStateCompareAndSetOperation[], +): Promise { + const sessionId = await resolveSessionId(); + return appStateCompareAndSetMany(sessionId, operations, { + requestSource: "agent", + }); +} + export async function listAppState( prefix: string, ): Promise }>> { diff --git a/packages/core/src/application-state/store.spec.ts b/packages/core/src/application-state/store.spec.ts index ee82f627d0..84f29b424c 100644 --- a/packages/core/src/application-state/store.spec.ts +++ b/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 (fn: (tx: typeof rawClient) => Promise) => { + 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; diff --git a/packages/core/src/application-state/store.ts b/packages/core/src/application-state/store.ts index 1d38d8a69d..de7341c2b9 100644 --- a/packages/core/src/application-state/store.ts +++ b/packages/core/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, + expectedValue: Record | null, nextValue: Record | null, options?: StoreWriteOptions, ): Promise { 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 | null; + nextValue: Record | null; +} + +const APP_STATE_CAS_MISMATCH = Symbol("app-state-cas-mismatch"); + +async function executeAppStateCompareAndSet( + client: DbExec, + sessionId: string, + key: string, + expectedValue: Record | null, + nextValue: Record | null, +): Promise { + 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 | null, + nextValue: Record | 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 { + 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 { + 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( diff --git a/packages/core/src/db/client.spec.ts b/packages/core/src/db/client.spec.ts index d22c0336f4..86ab70f1a2 100644 --- a/packages/core/src/db/client.spec.ts +++ b/packages/core/src/db/client.spec.ts @@ -130,6 +130,52 @@ describe("db/client dialect detection", () => { }); }); +describe("db/client D1 execution", () => { + it("uses D1 batch for atomic statements instead of interactive SQL transactions", async () => { + const prepared: Array<{ sql: string; args: unknown[] }> = []; + const binding = { + prepare: vi.fn((sql: string) => { + const statement = { + sql, + args: [] as unknown[], + bind(...args: unknown[]) { + statement.args = args; + return statement; + }, + all: vi.fn(), + }; + prepared.push(statement); + return statement; + }), + batch: vi.fn(async () => [ + { results: [{ matched: 1 }], meta: { changes: 0 } }, + { results: [], meta: { changes: 1 } }, + ]), + }; + const { createDbExec } = await import("./client.js"); + const client = await createDbExec({ d1Binding: binding }); + + expect(client.transaction).toBeUndefined(); + await expect( + client.atomicBatch?.([ + { sql: "SELECT value FROM state WHERE key = ?", args: ["pending"] }, + { sql: "DELETE FROM state WHERE key = ?", args: ["proposal"] }, + ]), + ).resolves.toEqual([ + { rows: [{ matched: 1 }], rowsAffected: 0 }, + { rows: [], rowsAffected: 1 }, + ]); + expect(binding.batch).toHaveBeenCalledTimes(1); + expect(prepared.map(({ sql, args }) => ({ sql, args }))).toEqual([ + { + sql: "SELECT value FROM state WHERE key = ?", + args: ["pending"], + }, + { sql: "DELETE FROM state WHERE key = ?", args: ["proposal"] }, + ]); + }); +}); + describe("pgliteDataDirFromUrl", () => { afterEach(() => { vi.unstubAllEnvs(); diff --git a/packages/core/src/db/client.ts b/packages/core/src/db/client.ts index b3c555b2f6..4d4c5dd051 100644 --- a/packages/core/src/db/client.ts +++ b/packages/core/src/db/client.ts @@ -24,6 +24,9 @@ export interface DbExec { sql: string | { sql: string; args?: unknown[] }, ): Promise<{ rows: any[]; rowsAffected: number }>; transaction?(fn: (tx: DbExec) => Promise): Promise; + atomicBatch?( + statements: readonly (string | { sql: string; args?: unknown[] })[], + ): Promise>; /** * Release the underlying connection/pool held by this exec. * Only non-singleton execs created via `createDbExec()` (e.g. the migration @@ -1033,7 +1036,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, + })); + }, }; } @@ -1497,6 +1510,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) => @@ -1521,6 +1538,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) => @@ -1540,8 +1561,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; } diff --git a/templates/design/AGENTS.md b/templates/design/AGENTS.md index 1443f09c0a..7340a171cd 100644 --- a/templates/design/AGENTS.md +++ b/templates/design/AGENTS.md @@ -277,11 +277,15 @@ patterns live in `.agents/skills/`. `view-screen`; not rendered as canvas overlays). - `design-reprompt-pending::` is the client-captured source selection, instruction, base hash, and authoritative current request id for - a scoped regenerate request. + a scoped regenerate request. The frontend starts requests through the + compare-and-set `begin-node-rewrite-request` action and must not overwrite a + `resolving` acceptance reservation. - `design-reprompt-proposal:::` is one request-specific preview-only subtree proposal. Candidate payloads have a - 256 KiB aggregate serialized limit. Resolution and cancellation use atomic - compare-and-set cleanup so an older request cannot erase a newer one. + 256 KiB aggregate serialized limit. Proposal publication, resolution, and + cancellation use atomic multi-key compare-and-set transitions. Acceptance + reserves its matching pending request before writing design content, so a + newer request and an older acceptance cannot both win. `view-screen` lists only proposals paired to the current pending request as `pendingCandidateReviews`. - `show-design-questions` opens focused pre-generation questions in the main diff --git a/templates/design/actions/begin-node-rewrite-request.spec.ts b/templates/design/actions/begin-node-rewrite-request.spec.ts new file mode 100644 index 0000000000..69ae3eb1dc --- /dev/null +++ b/templates/design/actions/begin-node-rewrite-request.spec.ts @@ -0,0 +1,91 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + assertAccess: vi.fn(), + readAppState: vi.fn(), + compareAndSetAppState: vi.fn(), +})); + +vi.mock("@agent-native/core", () => ({ + defineAction: (config: unknown) => config, +})); + +vi.mock("@agent-native/core/application-state", () => ({ + readAppState: mocks.readAppState, + compareAndSetAppState: mocks.compareAndSetAppState, +})); + +vi.mock("@agent-native/core/sharing", () => ({ + assertAccess: mocks.assertAccess, +})); + +import { designRepromptPendingStateKey } from "../shared/node-rewrite.js"; +import action from "./begin-node-rewrite-request.js"; + +const request = { + repromptId: "reprompt_2", + designId: "design_1", + fileId: "file_1", + target: { nodeId: "hero" }, + baseVersionHash: "hash_base", + instruction: "Make it darker", + createdAt: "2026-07-17T00:00:00.000Z", +}; + +describe("begin-node-rewrite-request", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.readAppState.mockResolvedValue(null); + mocks.compareAndSetAppState.mockResolvedValue(true); + }); + + it("atomically creates the first pending request", async () => { + await expect(action.run(request)).resolves.toEqual({ + started: true, + repromptId: "reprompt_2", + }); + expect(mocks.compareAndSetAppState).toHaveBeenCalledWith( + designRepromptPendingStateKey("design_1", "file_1"), + null, + expect.objectContaining({ status: "pending", repromptId: "reprompt_2" }), + ); + }); + + it("refuses to replace an acceptance reservation", async () => { + mocks.readAppState.mockResolvedValue({ + ...request, + repromptId: "reprompt_1", + status: "resolving", + claimId: "claim_1", + proposalId: "proposal_1", + resolution: "accept", + }); + + await expect(action.run(request)).rejects.toThrow( + "currently being accepted", + ); + expect(mocks.compareAndSetAppState).not.toHaveBeenCalled(); + }); + + it("requires a refinement to replace the proposal it was based on", async () => { + mocks.readAppState.mockResolvedValue({ + ...request, + repromptId: "reprompt_3", + status: "pending", + }); + + await expect( + action.run({ + ...request, + priorProposalId: "proposal_1", + priorRepromptId: "reprompt_1", + }), + ).rejects.toThrow("candidates changed before refinement"); + }); + + it("fails when another request wins the same compare-and-set", async () => { + mocks.compareAndSetAppState.mockResolvedValue(false); + + await expect(action.run(request)).rejects.toThrow("changed concurrently"); + }); +}); diff --git a/templates/design/actions/begin-node-rewrite-request.ts b/templates/design/actions/begin-node-rewrite-request.ts new file mode 100644 index 0000000000..d3dc29a77e --- /dev/null +++ b/templates/design/actions/begin-node-rewrite-request.ts @@ -0,0 +1,78 @@ +import { defineAction } from "@agent-native/core"; +import { + compareAndSetAppState, + readAppState, +} from "@agent-native/core/application-state"; +import { assertAccess } from "@agent-native/core/sharing"; +import { z } from "zod"; + +import { + designRepromptPendingStateKey, + isNodeRewriteResolutionClaim, + isPendingDesignReprompt, + type PendingDesignReprompt, +} from "../shared/node-rewrite.js"; + +const targetSchema = z + .object({ + nodeId: z.string().min(1).optional(), + selector: z.string().min(1).optional(), + }) + .refine((target) => target.nodeId || target.selector, { + message: "target.nodeId or target.selector is required", + }); + +export default defineAction({ + agentTool: false, + description: + "Atomically starts or supersedes one pending node rewrite unless its current proposal is being accepted.", + schema: z.object({ + repromptId: z.string().min(1), + designId: z.string().min(1), + fileId: z.string().min(1), + target: targetSchema, + baseVersionHash: z.string().min(1), + instruction: z.string().trim().min(1), + createdAt: z.string().min(1), + priorProposalId: z.string().min(1).optional(), + priorRepromptId: z.string().min(1).optional(), + }), + run: async (request) => { + await assertAccess("design", request.designId, "editor"); + const pendingKey = designRepromptPendingStateKey( + request.designId, + request.fileId, + ); + const current = await readAppState(pendingKey); + if (isNodeRewriteResolutionClaim(current)) { + throw new Error( + "This candidate is currently being accepted. Wait for it to finish before regenerating.", + ); + } + if ( + request.priorRepromptId && + (!isPendingDesignReprompt(current) || + current.repromptId !== request.priorRepromptId) + ) { + throw new Error( + "The candidates changed before refinement started. Review the latest candidates instead.", + ); + } + + const pending: PendingDesignReprompt = { + ...request, + status: "pending", + }; + const started = await compareAndSetAppState( + pendingKey, + current, + pending as unknown as Record, + ); + if (!started) { + throw new Error( + "The regeneration request changed concurrently. Try again with the latest candidates.", + ); + } + return { started: true, repromptId: request.repromptId }; + }, +}); diff --git a/templates/design/actions/cancel-node-rewrite-request.spec.ts b/templates/design/actions/cancel-node-rewrite-request.spec.ts index c63ee8ca88..d88246010a 100644 --- a/templates/design/actions/cancel-node-rewrite-request.spec.ts +++ b/templates/design/actions/cancel-node-rewrite-request.spec.ts @@ -4,6 +4,7 @@ const mocks = vi.hoisted(() => ({ assertAccess: vi.fn(), readAppState: vi.fn(), compareAndSetAppState: vi.fn(), + compareAndSetManyAppState: vi.fn(), })); vi.mock("@agent-native/core", () => ({ @@ -13,6 +14,7 @@ vi.mock("@agent-native/core", () => ({ vi.mock("@agent-native/core/application-state", () => ({ readAppState: mocks.readAppState, compareAndSetAppState: mocks.compareAndSetAppState, + compareAndSetManyAppState: mocks.compareAndSetManyAppState, })); vi.mock("@agent-native/core/sharing", () => ({ @@ -35,6 +37,7 @@ describe("cancel-node-rewrite-request", () => { beforeEach(() => { vi.clearAllMocks(); mocks.compareAndSetAppState.mockResolvedValue(true); + mocks.compareAndSetManyAppState.mockResolvedValue(true); }); it("does not cancel a newer request", async () => { @@ -71,4 +74,43 @@ describe("cancel-node-rewrite-request", () => { null, ); }); + + it("clears a pending request and its proposal in one transition", async () => { + const proposal = { + proposalId: "proposal_1", + repromptId: "reprompt_1", + designId: "design_1", + fileId: "file_1", + baseVersionHash: "hash_base", + target: { nodeId: "hero" }, + resolvedTarget: { nodeId: "hero", selector: "[data-hero]" }, + variants: [{ html: "
", summary: "Updated" }], + }; + mocks.readAppState + .mockResolvedValueOnce(pending) + .mockResolvedValueOnce(proposal); + + await expect( + action.run({ + designId: "design_1", + fileId: "file_1", + repromptId: "reprompt_1", + }), + ).resolves.toMatchObject({ + cancelled: true, + proposalCancelled: true, + superseded: false, + }); + expect(mocks.compareAndSetManyAppState).toHaveBeenCalledWith([ + expect.objectContaining({ + key: "design-reprompt-pending:design_1:file_1", + nextValue: null, + }), + expect.objectContaining({ + key: "design-reprompt-proposal:design_1:file_1:reprompt_1", + nextValue: null, + }), + ]); + expect(mocks.compareAndSetAppState).not.toHaveBeenCalled(); + }); }); diff --git a/templates/design/actions/cancel-node-rewrite-request.ts b/templates/design/actions/cancel-node-rewrite-request.ts index 2c04a8561d..9a2105c988 100644 --- a/templates/design/actions/cancel-node-rewrite-request.ts +++ b/templates/design/actions/cancel-node-rewrite-request.ts @@ -1,6 +1,7 @@ import { defineAction } from "@agent-native/core"; import { compareAndSetAppState, + compareAndSetManyAppState, readAppState, } from "@agent-native/core/application-state"; import { assertAccess } from "@agent-native/core/sharing"; @@ -39,15 +40,24 @@ export default defineAction({ repromptId, ); const proposal = await readAppState(proposalKey); - const [pendingCancelled, proposalCancelled] = await Promise.all([ - compareAndSetAppState(pendingKey, pending, null), - isNodeRewriteProposal(proposal) - ? compareAndSetAppState(proposalKey, proposal, null) - : Promise.resolve(false), - ]); + const proposalCancelled = isNodeRewriteProposal(proposal); + const pendingCancelled = proposalCancelled + ? await compareAndSetManyAppState([ + { + key: pendingKey, + expectedValue: pending, + nextValue: null, + }, + { + key: proposalKey, + expectedValue: proposal, + nextValue: null, + }, + ]) + : await compareAndSetAppState(pendingKey, pending, null); return { cancelled: pendingCancelled, - proposalCancelled, + proposalCancelled: proposalCancelled && pendingCancelled, superseded: !pendingCancelled, }; }, diff --git a/templates/design/actions/propose-node-rewrite.spec.ts b/templates/design/actions/propose-node-rewrite.spec.ts index 2345176d7e..cabe6e3f8d 100644 --- a/templates/design/actions/propose-node-rewrite.spec.ts +++ b/templates/design/actions/propose-node-rewrite.spec.ts @@ -26,8 +26,7 @@ const mocks = vi.hoisted(() => { selectChain, readLiveSourceFile: vi.fn(), readAppState: vi.fn(), - writeAppState: vi.fn(), - compareAndSetAppState: vi.fn(), + compareAndSetManyAppState: vi.fn(), assertAccess: vi.fn(), nanoid: vi.fn(), }; @@ -38,9 +37,8 @@ vi.mock("@agent-native/core", () => ({ })); vi.mock("@agent-native/core/application-state", () => ({ - compareAndSetAppState: mocks.compareAndSetAppState, + compareAndSetManyAppState: mocks.compareAndSetManyAppState, readAppState: mocks.readAppState, - writeAppState: mocks.writeAppState, })); vi.mock("@agent-native/core/sharing", () => ({ @@ -103,7 +101,7 @@ describe("propose-node-rewrite", () => { createdAt: "2026-07-16T00:00:00.000Z", }); mocks.nanoid.mockReturnValue("proposal_1"); - mocks.compareAndSetAppState.mockResolvedValue(true); + mocks.compareAndSetManyAppState.mockResolvedValue(true); }); it("uses client-valid per-design and per-file application-state keys", () => { @@ -155,15 +153,21 @@ describe("propose-node-rewrite", () => { expect(mocks.readAppState).toHaveBeenCalledWith( designRepromptPendingStateKey("design_1", "file_1"), ); - expect(mocks.writeAppState).toHaveBeenCalledWith( - designRepromptProposalStateKey("design_1", "file_1", "reprompt_1"), + expect(mocks.compareAndSetManyAppState).toHaveBeenCalledWith([ expect.objectContaining({ - proposalId: "node-rewrite-proposal_1", - repromptId: "reprompt_1", - baseVersionHash: "hash_base", - chosenIndex: 0, + key: designRepromptPendingStateKey("design_1", "file_1"), }), - ); + expect.objectContaining({ + key: designRepromptProposalStateKey("design_1", "file_1", "reprompt_1"), + expectedValue: null, + nextValue: expect.objectContaining({ + proposalId: "node-rewrite-proposal_1", + repromptId: "reprompt_1", + baseVersionHash: "hash_base", + chosenIndex: 0, + }), + }), + ]); expect(result.bridgeMessages).toEqual([ expect.objectContaining({ type: "node-html-preview", @@ -184,7 +188,7 @@ describe("propose-node-rewrite", () => { variants: [{ html: "
Changed
", summary: "Changed main" }], }), ).rejects.toThrow("does not match the pending selected subtree"); - expect(mocks.writeAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); }); it("fails with re-anchor guidance when the selected node is gone", async () => { @@ -205,7 +209,7 @@ describe("propose-node-rewrite", () => { ], }), ).rejects.toThrow(/Target missing.+re-anchor/i); - expect(mocks.writeAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); }); it("rejects a variant containing more than one root subtree", async () => { @@ -223,7 +227,7 @@ describe("propose-node-rewrite", () => { ], }), ).rejects.toThrow("exactly one root element"); - expect(mocks.writeAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); }); it("rejects candidate payloads above the bounded application-state budget", async () => { @@ -243,15 +247,11 @@ describe("propose-node-rewrite", () => { ], }), ).rejects.toThrow("too large to preview safely"); - expect(mocks.writeAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); }); it("removes stale candidates when a newer request wins during proposal creation", async () => { - const firstPending = await mocks.readAppState(); - mocks.readAppState - .mockReset() - .mockResolvedValueOnce(firstPending) - .mockResolvedValueOnce({ ...firstPending, repromptId: "reprompt_2" }); + mocks.compareAndSetManyAppState.mockResolvedValue(false); await expect( action.run({ @@ -265,19 +265,6 @@ describe("propose-node-rewrite", () => { }), ).rejects.toThrow("superseded by a newer request"); - const proposalKey = designRepromptProposalStateKey( - "design_1", - "file_1", - "reprompt_1", - ); - expect(mocks.writeAppState).toHaveBeenCalledWith( - proposalKey, - expect.objectContaining({ repromptId: "reprompt_1" }), - ); - expect(mocks.compareAndSetAppState).toHaveBeenCalledWith( - proposalKey, - expect.objectContaining({ repromptId: "reprompt_1" }), - null, - ); + expect(mocks.compareAndSetManyAppState).toHaveBeenCalledTimes(1); }); }); diff --git a/templates/design/actions/propose-node-rewrite.ts b/templates/design/actions/propose-node-rewrite.ts index 960608ce2c..4ddd1109eb 100644 --- a/templates/design/actions/propose-node-rewrite.ts +++ b/templates/design/actions/propose-node-rewrite.ts @@ -1,8 +1,8 @@ import { defineAction } from "@agent-native/core"; import { - compareAndSetAppState, + compareAndSetManyAppState, readAppState, - writeAppState, + type AppStateCompareAndSetOperation, } from "@agent-native/core/application-state"; import { accessFilter, assertAccess } from "@agent-native/core/sharing"; import { and, eq } from "drizzle-orm"; @@ -68,6 +68,7 @@ const variantSchema = z.object({ }); const pendingRepromptSchema = z.object({ + status: z.literal("pending").optional(), repromptId: z.string().min(1), designId: z.string().min(1), fileId: z.string().min(1), @@ -241,21 +242,18 @@ export default defineAction({ file.id, repromptId, ); - await writeAppState( - proposalKey, - proposal as unknown as Record, - ); - const currentPending = await readAppState(pendingKey); - if (currentPending?.repromptId !== repromptId) { - await compareAndSetAppState( - proposalKey, - proposal as unknown as Record, - null, - ); - throw new Error( - "This regeneration was superseded by a newer request before its candidates were published.", - ); - } + const publishOperations: AppStateCompareAndSetOperation[] = [ + { + key: pendingKey, + expectedValue: pending as unknown as Record, + nextValue: pending as unknown as Record, + }, + { + key: proposalKey, + expectedValue: null, + nextValue: proposal as unknown as Record, + }, + ]; if (pending.priorProposalId && pending.priorRepromptId) { const priorProposalKey = designRepromptProposalStateKey( file.designId, @@ -264,9 +262,19 @@ export default defineAction({ ); const priorProposal = await readAppState(priorProposalKey); if (priorProposal?.proposalId === pending.priorProposalId) { - await compareAndSetAppState(priorProposalKey, priorProposal, null); + publishOperations.push({ + key: priorProposalKey, + expectedValue: priorProposal, + nextValue: null, + }); } } + const published = await compareAndSetManyAppState(publishOperations); + if (!published) { + throw new Error( + "This regeneration was superseded by a newer request before its candidates were published.", + ); + } const bridgeMessages: NodeHtmlPreviewBridgeMessage[] = [ { diff --git a/templates/design/actions/resolve-node-rewrite.interleave.spec.ts b/templates/design/actions/resolve-node-rewrite.interleave.spec.ts index 43a341bb8f..ee13680366 100644 --- a/templates/design/actions/resolve-node-rewrite.interleave.spec.ts +++ b/templates/design/actions/resolve-node-rewrite.interleave.spec.ts @@ -28,6 +28,7 @@ const mocks = vi.hoisted(() => { selectChain, state: new Map>(), live: { content: initialContent, versionHash: "hash_base" }, + readLiveSourceFile: vi.fn(), writeInlineSourceFile: vi.fn(), }; }); @@ -49,16 +50,37 @@ vi.mock("@agent-native/core/application-state", () => ({ compareAndSetAppState: vi.fn( async ( key: string, - expected: Record, + expected: Record | null, next: Record | null, ) => { - const current = mocks.state.get(key); + const current = mocks.state.get(key) ?? null; if (JSON.stringify(current) !== JSON.stringify(expected)) return false; if (next) mocks.state.set(key, next); else mocks.state.delete(key); return true; }, ), + compareAndSetManyAppState: vi.fn( + async ( + operations: Array<{ + key: string; + expectedValue: Record | null; + nextValue: Record | null; + }>, + ) => { + const matches = operations.every( + ({ key, expectedValue }) => + JSON.stringify(mocks.state.get(key) ?? null) === + JSON.stringify(expectedValue), + ); + if (!matches) return false; + for (const { key, nextValue } of operations) { + if (nextValue) mocks.state.set(key, nextValue); + else mocks.state.delete(key); + } + return true; + }, + ), })); vi.mock("@agent-native/core/sharing", () => ({ @@ -91,10 +113,7 @@ vi.mock("../server/db/index.js", () => ({ })); vi.mock("../server/source-workspace.js", () => ({ - readLiveSourceFile: vi.fn(async () => ({ - ...mocks.live, - language: "html", - })), + readLiveSourceFile: mocks.readLiveSourceFile, writeInlineSourceFile: mocks.writeInlineSourceFile, })); @@ -102,6 +121,7 @@ import { designRepromptPendingStateKey, designRepromptProposalStateKey, } from "../shared/node-rewrite.js"; +import beginAction from "./begin-node-rewrite-request.js"; import proposeAction from "./propose-node-rewrite.js"; import resolveAction from "./resolve-node-rewrite.js"; @@ -112,6 +132,10 @@ describe("node rewrite propose/accept interleave", () => { mocks.live.content = mocks.initialContent; mocks.live.versionHash = "hash_base"; mocks.selectChain.limit.mockResolvedValue([mocks.file]); + mocks.readLiveSourceFile.mockImplementation(async () => ({ + ...mocks.live, + language: "html", + })); mocks.writeInlineSourceFile.mockResolvedValue({ changed: true, versionHash: "hash_next", @@ -166,4 +190,93 @@ describe("node rewrite propose/accept interleave", () => { mocks.state.has(designRepromptPendingStateKey("design_1", "file_1")), ).toBe(true); }); + + it("does not write an old proposal when a newer request wins before reservation", async () => { + const proposed = await proposeAction.run({ + source: { fileId: "file_1" }, + target: { nodeId: "hero" }, + baseVersionHash: "hash_base", + repromptId: "reprompt_1", + variants: [ + { + html: '
New
', + summary: "Updated hero", + }, + ], + }); + let finishRead!: () => void; + mocks.readLiveSourceFile.mockImplementationOnce( + () => + new Promise((resolve) => { + finishRead = () => + resolve({ ...mocks.live, language: "html" } as never); + }), + ); + + const accepting = resolveAction.run({ + proposalId: proposed.proposalId, + resolution: "accept", + }); + await vi.waitFor(() => expect(finishRead).toBeTypeOf("function")); + await beginAction.run({ + repromptId: "reprompt_2", + designId: "design_1", + fileId: "file_1", + target: { nodeId: "hero" }, + baseVersionHash: "hash_base", + instruction: "Try another direction", + createdAt: "2026-07-17T00:00:00.000Z", + }); + finishRead(); + + await expect(accepting).rejects.toThrow("newer regeneration request"); + expect(mocks.writeInlineSourceFile).not.toHaveBeenCalled(); + }); + + it("lets acceptance reservation defeat a concurrent rejection", async () => { + const proposed = await proposeAction.run({ + source: { fileId: "file_1" }, + target: { nodeId: "hero" }, + baseVersionHash: "hash_base", + repromptId: "reprompt_1", + variants: [ + { + html: '
New
', + summary: "Updated hero", + }, + ], + }); + let finishWrite!: () => void; + mocks.writeInlineSourceFile.mockImplementationOnce( + () => + new Promise((resolve) => { + finishWrite = () => + resolve({ + changed: true, + versionHash: "hash_next", + updatedAt: "2026-07-17T00:01:00.000Z", + }); + }), + ); + + const accepting = resolveAction.run({ + proposalId: proposed.proposalId, + resolution: "accept", + }); + await vi.waitFor(() => expect(finishWrite).toBeTypeOf("function")); + await expect( + resolveAction.run({ + proposalId: proposed.proposalId, + resolution: "reject", + }), + ).rejects.toThrow("newer regeneration request"); + finishWrite(); + + await expect(accepting).resolves.toEqual( + expect.objectContaining({ changed: true, cleanupComplete: true }), + ); + expect( + mocks.state.has(designRepromptPendingStateKey("design_1", "file_1")), + ).toBe(false); + }); }); diff --git a/templates/design/actions/resolve-node-rewrite.spec.ts b/templates/design/actions/resolve-node-rewrite.spec.ts index 815d94cf6b..84328778e5 100644 --- a/templates/design/actions/resolve-node-rewrite.spec.ts +++ b/templates/design/actions/resolve-node-rewrite.spec.ts @@ -27,6 +27,7 @@ const mocks = vi.hoisted(() => { listAppState: vi.fn(), readAppState: vi.fn(), compareAndSetAppState: vi.fn(), + compareAndSetManyAppState: vi.fn(), readLiveSourceFile: vi.fn(), writeInlineSourceFile: vi.fn(), assertAccess: vi.fn(), @@ -39,6 +40,7 @@ vi.mock("@agent-native/core", () => ({ vi.mock("@agent-native/core/application-state", () => ({ compareAndSetAppState: mocks.compareAndSetAppState, + compareAndSetManyAppState: mocks.compareAndSetManyAppState, listAppState: mocks.listAppState, readAppState: mocks.readAppState, })); @@ -126,6 +128,7 @@ describe("resolve-node-rewrite", () => { createdAt: "2026-07-16T00:00:00.000Z", }); mocks.compareAndSetAppState.mockResolvedValue(true); + mocks.compareAndSetManyAppState.mockResolvedValue(true); mocks.readLiveSourceFile.mockResolvedValue({ content: mocks.file.content, versionHash: "hash_base", @@ -169,17 +172,26 @@ describe("resolve-node-rewrite", () => { versionHash: "hash_next", }), ); - expect(mocks.compareAndSetAppState).toHaveBeenCalledTimes(2); - expect(mocks.compareAndSetAppState).toHaveBeenCalledWith( - designRepromptProposalStateKey("design_1", "file_1", "reprompt_1"), - proposal(), - null, - ); + expect(mocks.compareAndSetAppState).toHaveBeenCalledTimes(1); expect(mocks.compareAndSetAppState).toHaveBeenCalledWith( designRepromptPendingStateKey("design_1", "file_1"), expect.objectContaining({ repromptId: "reprompt_1" }), - null, + expect.objectContaining({ + status: "resolving", + proposalId: "proposal_1", + }), ); + expect(mocks.compareAndSetManyAppState).toHaveBeenCalledWith([ + expect.objectContaining({ + key: designRepromptProposalStateKey("design_1", "file_1", "reprompt_1"), + nextValue: null, + }), + expect.objectContaining({ + key: designRepromptPendingStateKey("design_1", "file_1"), + expectedValue: expect.objectContaining({ status: "resolving" }), + nextValue: null, + }), + ]); }); it("rejects without reading or writing design content and clears both records", async () => { @@ -190,7 +202,8 @@ describe("resolve-node-rewrite", () => { expect(mocks.readLiveSourceFile).not.toHaveBeenCalled(); expect(mocks.writeInlineSourceFile).not.toHaveBeenCalled(); - expect(mocks.compareAndSetAppState).toHaveBeenCalledTimes(2); + expect(mocks.compareAndSetAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).toHaveBeenCalledTimes(1); expect(result).toEqual( expect.objectContaining({ resolution: "reject", @@ -220,6 +233,7 @@ describe("resolve-node-rewrite", () => { action.run({ proposalId: "proposal_1", resolution: "reject" }), ).rejects.toThrow("newer regeneration request"); expect(mocks.compareAndSetAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); }); it("fails closed on a version mismatch and leaves proposal state intact", async () => { @@ -234,6 +248,7 @@ describe("resolve-node-rewrite", () => { ).rejects.toThrow("Screen changed since proposal"); expect(mocks.writeInlineSourceFile).not.toHaveBeenCalled(); expect(mocks.compareAndSetAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); }); it("does not clear state when the chosen variant index is invalid", async () => { @@ -246,5 +261,44 @@ describe("resolve-node-rewrite", () => { ).rejects.toThrow("out of range"); expect(mocks.writeInlineSourceFile).not.toHaveBeenCalled(); expect(mocks.compareAndSetAppState).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); + }); + + it("does not write when a newer request wins the acceptance reservation", async () => { + mocks.compareAndSetAppState.mockResolvedValue(false); + + await expect( + action.run({ proposalId: "proposal_1", resolution: "accept" }), + ).rejects.toThrow("newer regeneration request"); + expect(mocks.writeInlineSourceFile).not.toHaveBeenCalled(); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); + }); + + it("releases its acceptance reservation when the content write fails", async () => { + mocks.writeInlineSourceFile.mockRejectedValue(new Error("write failed")); + + await expect( + action.run({ proposalId: "proposal_1", resolution: "accept" }), + ).rejects.toThrow("write failed"); + expect(mocks.compareAndSetAppState).toHaveBeenCalledTimes(2); + expect(mocks.compareAndSetAppState.mock.calls[1]).toEqual([ + designRepromptPendingStateKey("design_1", "file_1"), + expect.objectContaining({ status: "resolving" }), + expect.objectContaining({ repromptId: "reprompt_1" }), + ]); + expect(mocks.compareAndSetManyAppState).not.toHaveBeenCalled(); + }); + + it("reports committed content without claiming incomplete cleanup succeeded", async () => { + mocks.compareAndSetManyAppState.mockResolvedValue(false); + + const result = await action.run({ + proposalId: "proposal_1", + resolution: "accept", + }); + + expect(result).toEqual( + expect.objectContaining({ changed: true, cleanupComplete: false }), + ); }); }); diff --git a/templates/design/actions/resolve-node-rewrite.ts b/templates/design/actions/resolve-node-rewrite.ts index 444cee7832..ae2032813c 100644 --- a/templates/design/actions/resolve-node-rewrite.ts +++ b/templates/design/actions/resolve-node-rewrite.ts @@ -1,11 +1,13 @@ import { defineAction } from "@agent-native/core"; import { compareAndSetAppState, + compareAndSetManyAppState, listAppState, readAppState, } from "@agent-native/core/application-state"; import { accessFilter, assertAccess } from "@agent-native/core/sharing"; import { and, eq } from "drizzle-orm"; +import { nanoid } from "nanoid"; import { z } from "zod"; import { getDb, schema } from "../server/db/index.js"; @@ -22,6 +24,7 @@ import { isPendingDesignReprompt, spliceNodeRewriteVariant, type NodeHtmlPreviewBridgeMessage, + type NodeRewriteResolutionClaim, type NodeRewriteProposal, } from "../shared/node-rewrite.js"; import { designSourceTypeFromData } from "../shared/source-mode.js"; @@ -142,22 +145,19 @@ async function clearProposalState( proposalKey: string, proposal: NodeRewriteProposal, pending: Record, -): Promise { +): Promise { const pendingKey = designRepromptPendingStateKey( proposal.designId, proposal.fileId, ); - const [proposalCleared] = await Promise.all([ - compareAndSetAppState( - proposalKey, - proposal as unknown as Record, - null, - ), - compareAndSetAppState(pendingKey, pending, null), + return compareAndSetManyAppState([ + { + key: proposalKey, + expectedValue: proposal as unknown as Record, + nextValue: null, + }, + { key: pendingKey, expectedValue: pending, nextValue: null }, ]); - if (!proposalCleared) { - throw new Error("Node rewrite proposal changed while it was resolving."); - } } export default defineAction({ @@ -195,7 +195,16 @@ export default defineAction({ if (resolution === "reject") { await assertAccess("design", proposal.designId, "editor"); - await clearProposalState(proposalKey, proposal, pending); + const cleanupComplete = await clearProposalState( + proposalKey, + proposal, + pending, + ); + if (!cleanupComplete) { + throw new Error( + "The proposal changed while it was being rejected. Review the latest candidates instead.", + ); + } return { proposalId, repromptId: proposal.repromptId, @@ -234,13 +243,51 @@ export default defineAction({ source, fileType: file.fileType, }); - const write = await writeInlineSourceFile({ - designId: file.designId, - file, - content: rewrite.content, - expectedVersionHash: proposal.baseVersionHash, - }); - await clearProposalState(proposalKey, proposal, pending); + const claim: NodeRewriteResolutionClaim = { + ...pending, + status: "resolving", + claimId: `node-rewrite-claim-${nanoid()}`, + proposalId, + resolution: "accept", + }; + const claimValue = claim as unknown as Record; + const claimed = await compareAndSetAppState( + pendingKey, + pending, + claimValue, + ); + if (!claimed) { + throw new Error( + "A newer regeneration request replaced this proposal before it could be accepted.", + ); + } + + let write; + try { + write = await writeInlineSourceFile({ + designId: file.designId, + file, + content: rewrite.content, + expectedVersionHash: proposal.baseVersionHash, + }); + } catch (error) { + const released = await compareAndSetAppState( + pendingKey, + claimValue, + pending, + ); + if (!released) { + throw new Error( + "The design write failed and its regeneration reservation could not be restored.", + ); + } + throw error; + } + const cleanupComplete = await clearProposalState( + proposalKey, + proposal, + claimValue, + ); return { proposalId, @@ -252,6 +299,7 @@ export default defineAction({ designId: file.designId, fileId: file.id, versionHash: write.versionHash, + cleanupComplete, }; }, }); diff --git a/templates/design/app/components/visual-editor/NodeRewriteProposal.spec.tsx b/templates/design/app/components/visual-editor/NodeRewriteProposal.spec.tsx index 3322469b91..30511432fa 100644 --- a/templates/design/app/components/visual-editor/NodeRewriteProposal.spec.tsx +++ b/templates/design/app/components/visual-editor/NodeRewriteProposal.spec.tsx @@ -228,14 +228,10 @@ describe("NodeRewriteProposal overview positioning", () => { await Promise.resolve(); }); - expect(mocks.setClientAppState).toHaveBeenCalledWith( - "design-reprompt-pending:design-1:screen-1", + expect(mocks.callAction).toHaveBeenCalledWith( + "begin-node-rewrite-request", expect.objectContaining({ instruction: "Make it warmer" }), ); - expect(mocks.setClientAppState).not.toHaveBeenCalledWith( - "design-reprompt-pending:design-1:screen-1", - null, - ); expect(mocks.callAction).toHaveBeenCalledWith( "cancel-node-rewrite-request", expect.objectContaining({ diff --git a/templates/design/app/components/visual-editor/NodeRewriteProposal.tsx b/templates/design/app/components/visual-editor/NodeRewriteProposal.tsx index 57243ce321..2f9eae6a96 100644 --- a/templates/design/app/components/visual-editor/NodeRewriteProposal.tsx +++ b/templates/design/app/components/visual-editor/NodeRewriteProposal.tsx @@ -1,12 +1,10 @@ import { callAction, - setClientAppState, useActionMutation, useChangeVersion, useT, } from "@agent-native/core/client"; import { - designRepromptPendingStateKey, isNodeRewriteProposal, type NodeHtmlPreviewBridgeMessage, type NodeRewriteProposal, @@ -426,7 +424,6 @@ export function NodeRewriteProposal({ if (!proposal || !instruction || refining || resolveMutation.isPending) return; const repromptId = crypto.randomUUID(); - const pendingKey = designRepromptPendingStateKey(designId, fileId); const pending = { repromptId, designId, @@ -440,7 +437,7 @@ export function NodeRewriteProposal({ }; setRefining(true); try { - await setClientAppState(pendingKey, pending); + await callAction("begin-node-rewrite-request", pending); const submission = formatNodeRepromptSubmission({ ...pending, priorProposalId: proposal.proposalId, diff --git a/templates/design/app/components/visual-editor/ReviewCanvasPins.spec.tsx b/templates/design/app/components/visual-editor/ReviewCanvasPins.spec.tsx index e9f68e2d27..92dcce7a4d 100644 --- a/templates/design/app/components/visual-editor/ReviewCanvasPins.spec.tsx +++ b/templates/design/app/components/visual-editor/ReviewCanvasPins.spec.tsx @@ -539,8 +539,8 @@ describe("ReviewCanvasPins persisted thread popover", () => { await act(async () => send?.click()); expect(mocks.createMutate).not.toHaveBeenCalled(); - expect(mocks.setClientAppState).toHaveBeenCalledWith( - "design-reprompt-pending:design-1:screen-1", + expect(mocks.callAction).toHaveBeenCalledWith( + "begin-node-rewrite-request", expect.objectContaining({ designId: "design-1", fileId: "screen-1", @@ -564,8 +564,10 @@ describe("ReviewCanvasPins persisted thread popover", () => { expect(document.body.textContent).not.toContain("review.clickToPin"); const pinsBeforePreview = document.querySelectorAll("[data-review-pin]"); - const appStateCalls = mocks.setClientAppState.mock.calls; - const pending = appStateCalls[appStateCalls.length - 1]?.[1] as { + const beginCall = mocks.callAction.mock.calls.find( + ([name]) => name === "begin-node-rewrite-request", + ); + const pending = beginCall?.[1] as { repromptId?: string; }; await act(async () => { diff --git a/templates/design/app/components/visual-editor/ReviewCanvasPins.tsx b/templates/design/app/components/visual-editor/ReviewCanvasPins.tsx index 0b4a0ffc2d..aebd7f4cd6 100644 --- a/templates/design/app/components/visual-editor/ReviewCanvasPins.tsx +++ b/templates/design/app/components/visual-editor/ReviewCanvasPins.tsx @@ -2,7 +2,6 @@ import { buildReviewThreads, callAction, ReviewCommentComposer, - setClientAppState, useCreateReviewComment, useReplyReviewComment, useResolveReviewThread, @@ -11,10 +10,7 @@ import { type ReviewThread, } from "@agent-native/core/client"; import type { ReviewComment } from "@agent-native/core/review"; -import { - designRepromptPendingStateKey, - type NodeRewriteTarget, -} from "@shared/node-rewrite"; +import type { NodeRewriteTarget } from "@shared/node-rewrite"; import { IconArrowUp, IconChevronDown, @@ -875,10 +871,9 @@ export function ReviewCanvasPins({ element = null; } } - const stateKey = designRepromptPendingStateKey(resourceId, targetId); setAgentSubmitting(true); try { - await setClientAppState(stateKey, pending); + await callAction("begin-node-rewrite-request", pending); const submission = formatNodeRepromptSubmission({ ...pending, subtreeHtml: diff --git a/templates/design/shared/node-rewrite.ts b/templates/design/shared/node-rewrite.ts index 9a96b19398..d49f4806f1 100644 --- a/templates/design/shared/node-rewrite.ts +++ b/templates/design/shared/node-rewrite.ts @@ -15,6 +15,7 @@ export const MAX_NODE_REWRITE_PROPOSAL_BYTES = 256 * 1024; export type NodeRewriteTarget = EditIntentTarget; export interface PendingDesignReprompt { + status?: "pending"; repromptId: string; designId: string; fileId: string; @@ -26,6 +27,16 @@ export interface PendingDesignReprompt { priorRepromptId?: string; } +export interface NodeRewriteResolutionClaim extends Omit< + PendingDesignReprompt, + "status" +> { + status: "resolving"; + claimId: string; + proposalId: string; + resolution: "accept"; +} + export interface NodeRewriteVariant { html: string; summary: string; @@ -69,6 +80,7 @@ export function isPendingDesignReprompt( if (!value || typeof value !== "object" || Array.isArray(value)) return false; const pending = value as Partial; return ( + (pending.status === undefined || pending.status === "pending") && typeof pending.repromptId === "string" && typeof pending.designId === "string" && typeof pending.fileId === "string" && @@ -79,6 +91,26 @@ export function isPendingDesignReprompt( ); } +export function isNodeRewriteResolutionClaim( + value: unknown, +): value is NodeRewriteResolutionClaim { + if (!value || typeof value !== "object" || Array.isArray(value)) return false; + const claim = value as Partial; + return ( + claim.status === "resolving" && + typeof claim.claimId === "string" && + typeof claim.proposalId === "string" && + claim.resolution === "accept" && + typeof claim.repromptId === "string" && + typeof claim.designId === "string" && + typeof claim.fileId === "string" && + typeof claim.baseVersionHash === "string" && + typeof claim.instruction === "string" && + typeof claim.createdAt === "string" && + Boolean(claim.target) + ); +} + export interface NodeHtmlPreviewBridgeMessage { type: "node-html-preview"; proposalId: string;