refactor: batch KV set and delete operations

packages/core/execution/src/tool-invoker.test.ts

@@ -366,52 +366,45 @@ describe("pause/resume with multiple elicitations", () => {
   // Regression: use separate top-level runPromise calls to match HTTP/CLI
   // pause/resume, and a single-elicit tool so no later pause can mask a dead
   // sandbox fiber.
-  it(
-    "resume returns across separate runPromise boundaries for a single-elicit tool (HTTP-like)",
-    async () => {
-      const executor = await Effect.runPromise(makeElicitingExecutor());
-      const engine = createExecutionEngine({ executor });
-
-      const code = "return await tools.api.singleApproval({});";
-
-      const outcome1 = await engine.executeWithPause(code);
-      expect(outcome1.status).toBe("paused");
-      const paused1 = outcome1 as Extract<typeof outcome1, { status: "paused" }>;
-      expect(paused1.execution.elicitationContext.request.message).toBe("Only approval");
-
-      // `execution.fiber` is on `InternalPausedExecution`; the exported
-      // `PausedExecution` type doesn't carry it. Cast to read.
-      const sandboxFiber = (
-        paused1.execution as unknown as {
-          readonly fiber: Fiber.Fiber<unknown, unknown>;
-        }
-      ).fiber;
-      const exitProbe = await Effect.runPromise(
-        Effect.race(
-          Fiber.await(sandboxFiber),
-          Effect.map(Effect.sleep("50 millis"), () => "still-running" as const),
-        ),
-      );
-      expect(exitProbe).toBe("still-running");
-
-      const outcome2 = await Promise.race([
-        engine.resume(paused1.execution.id, { action: "accept" }),
-        new Promise<never>((_, reject) =>
-          setTimeout(
-            () => reject(new Error("resume hung across runPromise boundaries")),
-            2000,
-          ),
-        ),
-      ]);
-
-      expect(outcome2).not.toBeNull();
-      const resumed = outcome2 as NonNullable<typeof outcome2>;
-      expect(resumed.status).toBe("completed");
-      if (resumed.status === "completed") {
-        expect(resumed.result.error).toBeUndefined();
-        expect(resumed.result.result).toMatchObject({ ok: true });
+  it("resume returns across separate runPromise boundaries for a single-elicit tool (HTTP-like)", async () => {
+    const executor = await Effect.runPromise(makeElicitingExecutor());
+    const engine = createExecutionEngine({ executor });
+
+    const code = "return await tools.api.singleApproval({});";
+
+    const outcome1 = await engine.executeWithPause(code);
+    expect(outcome1.status).toBe("paused");
+    const paused1 = outcome1 as Extract<typeof outcome1, { status: "paused" }>;
+    expect(paused1.execution.elicitationContext.request.message).toBe("Only approval");
+
+    // `execution.fiber` is on `InternalPausedExecution`; the exported
+    // `PausedExecution` type doesn't carry it. Cast to read.
+    const sandboxFiber = (
+      paused1.execution as unknown as {
+        readonly fiber: Fiber.Fiber<unknown, unknown>;
       }
-    },
-    10000,
-  );
+    ).fiber;
+    const exitProbe = await Effect.runPromise(
+      Effect.race(
+        Fiber.await(sandboxFiber),
+        Effect.map(Effect.sleep("50 millis"), () => "still-running" as const),
+      ),
+    );
+    expect(exitProbe).toBe("still-running");
+
+    const outcome2 = await Promise.race([
+      engine.resume(paused1.execution.id, { action: "accept" }),
+      new Promise<never>((_, reject) =>
+        setTimeout(() => reject(new Error("resume hung across runPromise boundaries")), 2000),
+      ),
+    ]);
+
+    expect(outcome2).not.toBeNull();
+    const resumed = outcome2 as NonNullable<typeof outcome2>;
+    expect(resumed.status).toBe("completed");
+    if (resumed.status === "completed") {
+      expect(resumed.result.error).toBeUndefined();
+      expect(resumed.result.result).toMatchObject({ ok: true });
+    }
+  }, 10000);
 });

packages/core/sdk/src/index.ts

@@ -101,4 +101,4 @@ export { makeInMemoryPolicyEngine } from "./in-memory/policy-engine";
 
 // Testing
 export { makeTestConfig } from "./testing";
-export { type Kv, type ScopedKv, scopeKv, makeInMemoryScopedKv } from "./plugin-kv";
+export { type Kv, type KvEntry, type ScopedKv, scopeKv, makeInMemoryScopedKv } from "./plugin-kv";

packages/core/sdk/src/plugin-kv.ts

@@ -8,14 +8,21 @@
 
 import { Effect } from "effect";
 
+export interface KvEntry {
+  readonly key: string;
+  readonly value: string;
+}
+
 /**
  * Global KV — requires a namespace on every call.
  * Implementations: makeSqliteKv, makeInMemoryKv
  */
 export interface Kv {
   readonly get: (namespace: string, key: string) => Effect.Effect<string | null>;
-  readonly set: (namespace: string, key: string, value: string) => Effect.Effect<void>;
-  readonly delete: (namespace: string, key: string) => Effect.Effect<boolean>;
+  /** Batch upsert — inserts or updates one or more key-value pairs. */
+  readonly set: (namespace: string, entries: readonly KvEntry[]) => Effect.Effect<void>;
+  /** Batch delete — removes one or more keys. */
+  readonly delete: (namespace: string, keys: readonly string[]) => Effect.Effect<number>;
   readonly list: (namespace: string) => Effect.Effect<readonly { key: string; value: string }[]>;
   readonly deleteAll: (namespace: string) => Effect.Effect<number>;
   readonly withTransaction?: <A, E>(
@@ -29,8 +36,10 @@ export interface Kv {
  */
 export interface ScopedKv {
   readonly get: (key: string) => Effect.Effect<string | null>;
-  readonly set: (key: string, value: string) => Effect.Effect<void>;
-  readonly delete: (key: string) => Effect.Effect<boolean>;
+  /** Batch upsert — inserts or updates one or more key-value pairs. */
+  readonly set: (entries: readonly KvEntry[]) => Effect.Effect<void>;
+  /** Batch delete — removes one or more keys. */
+  readonly delete: (keys: readonly string[]) => Effect.Effect<number>;
   readonly list: () => Effect.Effect<readonly { key: string; value: string }[]>;
   readonly deleteAll: () => Effect.Effect<number>;
   readonly withTransaction?: <A, E>(
@@ -43,8 +52,8 @@ export interface ScopedKv {
  */
 export const scopeKv = (kv: Kv, namespace: string): ScopedKv => ({
   get: (key) => kv.get(namespace, key),
-  set: (key, value) => kv.set(namespace, key, value),
-  delete: (key) => kv.delete(namespace, key),
+  set: (entries) => kv.set(namespace, entries),
+  delete: (keys) => kv.delete(namespace, keys),
   list: () => kv.list(namespace),
   deleteAll: () => kv.deleteAll(namespace),
   withTransaction: kv.withTransaction,
@@ -57,11 +66,16 @@ export const makeInMemoryScopedKv = (): ScopedKv => {
   const store = new Map<string, string>();
   return {
     get: (key) => Effect.succeed(store.get(key) ?? null),
-    set: (key, value) =>
+    set: (entries) =>
+      Effect.sync(() => {
+        for (const { key, value } of entries) store.set(key, value);
+      }),
+    delete: (keys) =>
       Effect.sync(() => {
-        store.set(key, value);
+        let count = 0;
+        for (const key of keys) if (store.delete(key)) count++;
+        return count;
       }),
-    delete: (key) => Effect.sync(() => store.delete(key)),
     list: () => Effect.sync(() => [...store.entries()].map(([key, value]) => ({ key, value }))),
     deleteAll: () =>
       Effect.sync(() => {

packages/core/storage-file/src/index.ts

@@ -85,8 +85,8 @@ export const makeKvConfig = <const TPlugins extends readonly ExecutorPlugin<stri
  */
 export const makeScopedKv = (kv: Kv, folder: string): Kv => ({
   get: (namespace, key) => kv.get(`${folder}::${namespace}`, key),
-  set: (namespace, key, value) => kv.set(`${folder}::${namespace}`, key, value),
-  delete: (namespace, key) => kv.delete(`${folder}::${namespace}`, key),
+  set: (namespace, entries) => kv.set(`${folder}::${namespace}`, entries),
+  delete: (namespace, keys) => kv.delete(`${folder}::${namespace}`, keys),
   list: (namespace) => kv.list(`${folder}::${namespace}`),
   deleteAll: (namespace) => kv.deleteAll(`${folder}::${namespace}`),
   withTransaction: kv.withTransaction,

packages/core/storage-file/src/plugin-kv.ts

@@ -29,22 +29,30 @@ export const makeSqliteKv = (sql: SqlClient.SqlClient): Kv => ({
       }),
     ),
 
-  set: (namespace, key, value) =>
+  set: (namespace, entries) =>
     absorbSql(
-      sql`
-      INSERT OR REPLACE INTO kv (namespace, key, value)
-      VALUES (${namespace}, ${key}, ${value})
-    `.pipe(Effect.asVoid),
+      Effect.gen(function* () {
+        for (const { key, value } of entries) {
+          yield* sql`
+            INSERT OR REPLACE INTO kv (namespace, key, value)
+            VALUES (${namespace}, ${key}, ${value})
+          `;
+        }
+      }),
     ),
 
-  delete: (namespace, key) =>
+  delete: (namespace, keys) =>
     absorbSql(
       Effect.gen(function* () {
-        const before = yield* sql<{ c: number }>`
-        SELECT COUNT(*) as c FROM kv WHERE namespace = ${namespace} AND key = ${key}
-      `;
-        yield* sql`DELETE FROM kv WHERE namespace = ${namespace} AND key = ${key}`;
-        return (before[0]?.c ?? 0) > 0;
+        let count = 0;
+        for (const key of keys) {
+          const before = yield* sql<{ c: number }>`
+            SELECT COUNT(*) as c FROM kv WHERE namespace = ${namespace} AND key = ${key}
+          `;
+          yield* sql`DELETE FROM kv WHERE namespace = ${namespace} AND key = ${key}`;
+          if ((before[0]?.c ?? 0) > 0) count++;
+        }
+        return count;
       }),
     ),
 
@@ -110,12 +118,19 @@ export const makeInMemoryKv = (): Kv => {
   return {
     get: (namespace, key) => Effect.succeed(bucket(namespace).get(key) ?? null),
 
-    set: (namespace, key, value) =>
+    set: (namespace, entries) =>
       Effect.sync(() => {
-        bucket(namespace).set(key, value);
+        const b = bucket(namespace);
+        for (const { key, value } of entries) b.set(key, value);
       }),
 
-    delete: (namespace, key) => Effect.sync(() => bucket(namespace).delete(key)),
+    delete: (namespace, keys) =>
+      Effect.sync(() => {
+        const b = bucket(namespace);
+        let count = 0;
+        for (const key of keys) if (b.delete(key)) count++;
+        return count;
+      }),
 
     list: (namespace) =>
       Effect.sync(() => [...bucket(namespace).entries()].map(([key, value]) => ({ key, value }))),

packages/core/storage-file/src/policy-engine.ts

@@ -26,7 +26,8 @@ export const makeKvPolicyEngine = (policiesKv: ScopedKv, metaKv: ScopedKv) => {
       return raw ? parseInt(raw, 10) : 0;
     });
 
-  const setCounter = (n: number): Effect.Effect<void> => metaKv.set("policy_counter", String(n));
+  const setCounter = (n: number): Effect.Effect<void> =>
+    metaKv.set([{ key: "policy_counter", value: String(n) }]);
 
   return {
     list: (scopeId: ScopeId) =>
@@ -43,15 +44,15 @@ export const makeKvPolicyEngine = (policiesKv: ScopedKv, metaKv: ScopedKv) => {
         yield* setCounter(counter);
         const id = PolicyId.make(`policy-${counter}`);
         const full = new Policy({ ...policy, id, createdAt: new Date() });
-        yield* policiesKv.set(id, encodePolicy(full));
+        yield* policiesKv.set([{ key: id, value: encodePolicy(full) }]);
         return full;
       }),
 
     remove: (policyId: PolicyId) =>
       Effect.gen(function* () {
         const raw = yield* policiesKv.get(policyId);
         if (!raw) return false;
-        yield* policiesKv.delete(policyId);
+        yield* policiesKv.delete([policyId]);
         return true;
       }),
   };

packages/core/storage-file/src/secret-store.ts

@@ -150,7 +150,7 @@ export const makeKvSecretStore = (refsKv: ScopedKv) => {
           createdAt: new Date(),
         });
 
-        yield* refsKv.set(input.id, encodeRef(ref));
+        yield* refsKv.set([{ key: input.id, value: encodeRef(ref) }]);
         return ref;
       }),
 
@@ -164,7 +164,7 @@ export const makeKvSecretStore = (refsKv: ScopedKv) => {
         const provider = findWritableProvider(providerKey);
         if (provider?.delete) yield* provider.delete(secretId);
 
-        yield* refsKv.delete(secretId);
+        yield* refsKv.delete([secretId]);
         return true;
       }),
 

packages/core/storage-file/src/tool-registry.ts

@@ -116,12 +116,13 @@ export const makeKvToolRegistry = (toolsKv: ScopedKv, defsKv: ScopedKv) => {
     registerDefinitions: (newDefs: Record<string, unknown>) =>
       withKvTransaction(
         defsKv,
-        Effect.gen(function* () {
-          for (const [name, schema] of Object.entries(newDefs)) {
+        defsKv.set(
+          Object.entries(newDefs).map(([name, schema]) => ({
+            key: name,
             // @effect-diagnostics-next-line preferSchemaOverJson:off
-            yield* defsKv.set(name, JSON.stringify(schema));
-          }
-        }),
+            value: JSON.stringify(schema),
+          })),
+        ),
       ),
 
     registerRuntimeDefinitions: (newDefs: Record<string, unknown>) =>
@@ -178,11 +179,7 @@ export const makeKvToolRegistry = (toolsKv: ScopedKv, defsKv: ScopedKv) => {
     register: (newTools: readonly ToolRegistration[]) =>
       withKvTransaction(
         toolsKv,
-        Effect.gen(function* () {
-          for (const t of newTools) {
-            yield* toolsKv.set(t.id, encodeTool(t));
-          }
-        }),
+        toolsKv.set(newTools.map((t) => ({ key: t.id, value: encodeTool(t) }))),
       ),
 
     registerRuntime: (newTools: readonly ToolRegistration[]) =>
@@ -210,18 +207,15 @@ export const makeKvToolRegistry = (toolsKv: ScopedKv, defsKv: ScopedKv) => {
         for (const id of toolIds) {
           runtimeTools.delete(id);
           runtimeHandlers.delete(id);
-          yield* toolsKv.delete(id);
         }
+        yield* toolsKv.delete([...toolIds]);
       }),
 
     unregisterBySource: (sourceId: string) =>
       Effect.gen(function* () {
         const allTools = yield* getAllTools();
-        for (const t of allTools) {
-          if (t.sourceId === sourceId) {
-            yield* toolsKv.delete(t.id);
-          }
-        }
+        const idsToDelete = allTools.filter((t) => t.sourceId === sourceId).map((t) => t.id);
+        if (idsToDelete.length > 0) yield* toolsKv.delete(idsToDelete);
         for (const [id, t] of runtimeTools) {
           if (t.sourceId === sourceId) {
             runtimeTools.delete(id);

packages/core/storage-postgres/src/index.test.ts

@@ -335,13 +335,13 @@ describe("Executor with Postgres storage", () => {
       const kv = makePgKv(db, TEST_ORG_ID);
       const scoped = scopeKv(kv, "my-plugin");
 
-      yield* scoped.set("k1", "v1");
+      yield* scoped.set([{ key: "k1", value: "v1" }]);
       expect(yield* scoped.get("k1")).toBe("v1");
 
       const items = yield* scoped.list();
       expect(items).toHaveLength(1);
 
-      yield* scoped.delete("k1");
+      yield* scoped.delete(["k1"]);
       expect(yield* scoped.get("k1")).toBeNull();
     }),
   );
@@ -351,8 +351,8 @@ describe("Executor with Postgres storage", () => {
       const kv1 = makePgKv(db, "org-a");
       const kv2 = makePgKv(db, "org-b");
 
-      yield* kv1.set("ns", "key", "org-a-value");
-      yield* kv2.set("ns", "key", "org-b-value");
+      yield* kv1.set("ns", [{ key: "key", value: "org-a-value" }]);
+      yield* kv2.set("ns", [{ key: "key", value: "org-b-value" }]);
 
       expect(yield* kv1.get("ns", "key")).toBe("org-a-value");
       expect(yield* kv2.get("ns", "key")).toBe("org-b-value");