527 lines · 17.4 KB · TypeScript
RawDownload
import { describe, expect, it } from "@effect/vitest";import { Cause, Effect, Exit, Schema } from "effect"; import { createExecutor } from "./executor";import { StorageError, type FumaDb } from "./fuma-runtime";import { Owner } from "./ids";import { definePlugin } from "./plugin";import {  definePluginStorageCollection,  type PluginStorageCollectionFacade,  type PluginStorageCollectionQueryInput,  type PluginStorageCollectionWhere,} from "./plugin-storage";import { makeTestConfig, makeTestExecutor } from "./testing"; const ToolCall = Schema.Struct({  runId: Schema.String,  toolId: Schema.String,  userId: Schema.NullOr(Schema.String),  clientName: Schema.NullOr(Schema.String),  status: Schema.Literals(["ok", "failed", "blocked"]),  startedAt: Schema.String,  durationMs: Schema.Number,}); const toolCalls = definePluginStorageCollection("toolCalls", ToolCall, {  indexes: ["runId", "toolId", "status", "clientName", "startedAt", ["toolId", "startedAt"]],}); type ToolCall = typeof ToolCall.Type; const assertPluginStorageTypes = (storage: PluginStorageCollectionFacade<typeof toolCalls>) => {  const validQuery = storage.query({ where: { toolId: "shell" } });   // @ts-expect-error durationMs is part of the data shape but is not declared as an index.  const invalidWhereQuery = storage.query({ where: { durationMs: 100 } });   // prettier-ignore  // @ts-expect-error orderBy is also restricted to declared index fields.  const invalidOrderQuery = storage.query({ orderBy: [{ field: "durationMs" }] });   // @ts-expect-error indexes must point at fields in the collection schema.  definePluginStorageCollection("bad", ToolCall, { indexes: ["missing"] });   void validQuery;  void invalidWhereQuery;  void invalidOrderQuery;};void assertPluginStorageTypes; const uncheckedToolCallWhere = (  where: Readonly<Record<string, unknown>>,): PluginStorageCollectionWhere<typeof toolCalls> =>  where as PluginStorageCollectionWhere<typeof toolCalls>; const executionHistoryPlugin = definePlugin(() => ({  id: "executionHistory" as const,  pluginStorage: { toolCalls },  storage: ({ pluginStorage }) => ({    toolCalls: pluginStorage.collection(toolCalls),  }),  extension: (ctx) => ({    record: (owner: Owner, key: string, data: ToolCall) =>      ctx.storage.toolCalls.put({ owner, key, data }),    recordMany: (      owner: Owner,      rows: readonly { readonly key: string; readonly data: ToolCall }[],    ) =>      ctx.pluginStorage.putMany({        owner,        entries: rows.map((row) => ({          collection: toolCalls.name,          key: row.key,          data: row.data,        })),      }),    removeMany: (owner: Owner, keys: readonly string[]) =>      ctx.pluginStorage.removeMany({        owner,        entries: keys.map((key) => ({ collection: toolCalls.name, key })),      }),    get: (key: string) => ctx.storage.toolCalls.get({ key }),    getForOwner: (owner: Owner, key: string) => ctx.storage.toolCalls.getForOwner({ owner, key }),    query: (input?: PluginStorageCollectionQueryInput<typeof toolCalls>) =>      ctx.storage.toolCalls.query(input),    count: (      input?: Omit<        PluginStorageCollectionQueryInput<typeof toolCalls>,        "orderBy" | "limit" | "offset"      >,    ) => ctx.storage.toolCalls.count(input),    queryUnindexed: () =>      ctx.storage.toolCalls.query({        where: uncheckedToolCallWhere({ durationMs: 100 }),      }),  }),}))(); const call = (input: {  readonly runId: string;  readonly toolId: string;  readonly status: ToolCall["status"];  readonly startedAt: string;  readonly clientName?: string | null;  readonly userId?: string | null;  readonly durationMs?: number;}): ToolCall => ({  runId: input.runId,  toolId: input.toolId,  userId: input.userId ?? null,  clientName: input.clientName ?? null,  status: input.status,  startedAt: input.startedAt,  durationMs: input.durationMs ?? 0,}); // A FumaDB that commits the FIRST row of a multi-row `plugin_storage` bulk// write and then fails. Injecting the fault mid-write, rather than before it,// is what makes the two rollback cases below meaningful: the row is really on// disk when the failure lands, so only the enclosing transaction can take it// back. The fault is unconditional (not armed by entering a transaction) so// that dropping the transaction is a visible failure and not a silently// disarmed test.const failPluginStorageBulkWriteAfterFirstRow = (db: FumaDb): FumaDb => {  const wrap = (source: FumaDb): FumaDb =>    new Proxy(source, {      get(target, property, receiver) {        if (property === "withContext") {          const withContext = target.withContext;          return withContext === undefined            ? undefined            : (context: unknown) => wrap(withContext(context));        }        if (property === "transaction") {          const transaction: FumaDb["transaction"] = (run) =>            target.transaction((transactionDb) => run(wrap(transactionDb)));          return transaction;        }        if (property === "upsertMany") {          const upsertMany: FumaDb["upsertMany"] = async (table, options) => {            if (table !== "plugin_storage" || options.values.length < 2) {              return target.upsertMany(table, options);            }             await target.upsertMany(table, {              ...options,              values: options.values.slice(0, 1),            });            // oxlint-disable-next-line executor/no-promise-reject -- boundary: fault-injecting FumaDB adapter must reject to exercise transaction rollback            return Promise.reject(              new StorageError({                message: "Injected plugin storage bulk-write failure.",                cause: undefined,              }),            );          };          return upsertMany;        }        return Reflect.get(target, property, receiver);      },    });   return wrap(db);}; describe("plugin storage collections", () => {  it.effect("queries declared indexes through the executor's SQLite FumaDB target", () =>    Effect.gen(function* () {      const executor = yield* makeTestExecutor({        backend: "sqlite",        plugins: [executionHistoryPlugin] as const,      });       yield* executor.executionHistory.record(        "org",        "call-1",        call({          runId: "run-a",          toolId: "browser",          status: "failed",          clientName: "codex",          startedAt: "2026-05-29T10:00:00.000Z",          durationMs: 320,        }),      );      yield* executor.executionHistory.record(        "org",        "call-2",        call({          runId: "run-a",          toolId: "shell",          status: "ok",          clientName: "codex",          startedAt: "2026-05-29T10:01:00.000Z",          durationMs: 42,        }),      );      yield* executor.executionHistory.record(        "org",        "call-3",        call({          runId: "run-b",          toolId: "shell",          status: "failed",          clientName: "codex",          startedAt: "2026-05-29T10:02:00.000Z",          durationMs: 77,        }),      );       const failed = yield* executor.executionHistory.query({        where: {          clientName: "codex",          status: "failed",        },        orderBy: [{ field: "startedAt", direction: "desc" }],        limit: 10,      });      expect(failed.map((entry) => entry.key)).toEqual(["call-3", "call-1"]);      expect(failed.map((entry) => entry.data.toolId)).toEqual(["shell", "browser"]);       const shellCount = yield* executor.executionHistory.count({        where: { toolId: "shell" },      });      expect(shellCount).toBe(2);    }),  );   it.effect("bulk puts and removes plugin storage rows", () =>    Effect.gen(function* () {      const executor = yield* makeTestExecutor({        backend: "sqlite",        plugins: [executionHistoryPlugin] as const,      });      const rows = Array.from({ length: 95 }, (_, index) => ({        key: `call-${String(index).padStart(3, "0")}`,        data: call({          runId: "run-bulk",          toolId: index % 2 === 0 ? "browser" : "shell",          status: "ok",          startedAt: new Date(Date.UTC(2026, 4, 29, 12, index)).toISOString(),        }),      }));       yield* executor.executionHistory.recordMany("org", rows);      yield* executor.executionHistory.recordMany("org", [        {          key: "call-000",          data: call({            runId: "run-bulk",            toolId: "browser",            status: "failed",            startedAt: "2026-05-29T12:00:00.000Z",          }),        },      ]);       const stored = yield* executor.executionHistory.query({        where: { runId: "run-bulk" },        orderBy: [{ field: "startedAt" }],      });      expect(stored).toHaveLength(95);      expect(stored[0]?.key).toBe("call-000");      expect(stored[0]?.data.status).toBe("failed");       yield* executor.executionHistory.removeMany(        "org",        rows.map((row) => row.key),      );      const remaining = yield* executor.executionHistory.query({ where: { runId: "run-bulk" } });      expect(remaining).toEqual([]);    }),  );   it.effect("stores and overwrites every row when a bulk write spans multiple batches", () =>    Effect.gen(function* () {      const executor = yield* makeTestExecutor({        backend: "sqlite",        plugins: [executionHistoryPlugin] as const,      });      // A plugin_storage row binds ~9 values, so 300 rows exceed one      // 999-bound-variable statement budget and must span several batches.      const entries = (status: "ok" | "failed") =>        Array.from({ length: 300 }, (_, index) => ({          key: `batched-call-${String(index).padStart(3, "0")}`,          data: call({            runId: "run-batched",            toolId: "browser",            status,            startedAt: new Date(Date.UTC(2026, 4, 29, 12, 0, index)).toISOString(),          }),        }));       yield* executor.executionHistory.recordMany("org", entries("ok"));      const total = yield* executor.executionHistory.count({        where: { runId: "run-batched" },      });      expect(total).toBe(300);       yield* executor.executionHistory.recordMany("org", entries("failed"));      const failed = yield* executor.executionHistory.count({        where: { runId: "run-batched", status: "failed" },      });      expect(failed).toBe(300);      const totalAfterOverwrite = yield* executor.executionHistory.count({        where: { runId: "run-batched" },      });      expect(totalAfterOverwrite).toBe(300);    }),  );   it.effect("rolls back every plugin storage row when a bulk write fails", () =>    Effect.gen(function* () {      const config = makeTestConfig({        backend: "sqlite",        plugins: [executionHistoryPlugin] as const,      });      const executor = yield* Effect.acquireRelease(        createExecutor({          ...config,          db: failPluginStorageBulkWriteAfterFirstRow(config.db),        }),        (instance) =>          instance            .close()            .pipe(              Effect.ignore,              Effect.andThen(Effect.promise(() => config.testDb.close()).pipe(Effect.ignore)),            ),      );       const exit = yield* Effect.exit(        executor.executionHistory.recordMany("org", [          {            key: "call-first",            data: call({              runId: "run-rollback",              toolId: "browser",              status: "ok",              startedAt: "2026-05-29T12:00:00.000Z",            }),          },          {            key: "call-second",            data: call({              runId: "run-rollback",              toolId: "shell",              status: "ok",              startedAt: "2026-05-29T12:01:00.000Z",            }),          },        ]),      );      expect(Exit.isFailure(exit)).toBe(true);       const stored = yield* executor.executionHistory.query({        where: { runId: "run-rollback" },      });      expect(stored).toEqual([]);    }),  );   // The hazard this whole change exists to remove: the previous implementation  // deleted every target key and only then re-created the rows, so a failure  // between the two halves destroyed data the caller never meant to touch. An  // upsert inside a transaction cannot lose a row it did not successfully  // replace, so the ORIGINAL values must still be readable after the failure —  // not merely absent-and-consistent like the rolled-back insert above.  it.effect("leaves pre-existing rows intact when a bulk overwrite fails mid-batch", () =>    Effect.gen(function* () {      const config = makeTestConfig({        backend: "sqlite",        plugins: [executionHistoryPlugin] as const,      });      const executor = yield* Effect.acquireRelease(        createExecutor({          ...config,          db: failPluginStorageBulkWriteAfterFirstRow(config.db),        }),        (instance) =>          instance            .close()            .pipe(              Effect.ignore,              Effect.andThen(Effect.promise(() => config.testDb.close()).pipe(Effect.ignore)),            ),      );       // Seeded one row at a time, so the seeding itself never goes through the      // bulk path the fault injector breaks.      const original = [        {          key: "call-first",          data: call({            runId: "run-preexisting",            toolId: "browser",            status: "ok",            startedAt: "2026-05-29T12:00:00.000Z",          }),        },        {          key: "call-second",          data: call({            runId: "run-preexisting",            toolId: "shell",            status: "ok",            startedAt: "2026-05-29T12:01:00.000Z",          }),        },      ];      for (const row of original) {        yield* executor.executionHistory.record("org", row.key, row.data);      }       const exit = yield* Effect.exit(        executor.executionHistory.recordMany(          "org",          original.map((row) => ({            key: row.key,            data: { ...row.data, toolId: "overwritten", status: "failed" as const },          })),        ),      );      expect(Exit.isFailure(exit)).toBe(true);       const first = yield* executor.executionHistory.get("call-first");      const second = yield* executor.executionHistory.get("call-second");      expect(first?.data).toEqual(original[0]!.data);      expect(second?.data).toEqual(original[1]!.data);    }),  );   it.effect(    "bulk puts large plugin storage row sets in bounded batches",    () =>      Effect.gen(function* () {        const executor = yield* makeTestExecutor({          backend: "sqlite",          plugins: [executionHistoryPlugin] as const,        });        const rows = Array.from({ length: 7_000 }, (_, index) => ({          key: `large-call-${String(index).padStart(5, "0")}`,          data: call({            runId: "run-large-bulk",            toolId: index % 2 === 0 ? "browser" : "shell",            status: "ok",            startedAt: new Date(Date.UTC(2026, 4, 29, 12, 0, index)).toISOString(),          }),        }));         yield* executor.executionHistory.recordMany("org", rows);         const count = yield* executor.executionHistory.count({          where: { runId: "run-large-bulk" },        });        expect(count).toBe(rows.length);      }),    15_000,  );   it.effect("user rows shadow org rows on read; both share one plugin_storage table", () =>    Effect.gen(function* () {      // One executor bound to a subject sees both org and user owner rows; a      // user-owned row shadows an org-owned row under the same key on read.      const executor = yield* makeTestExecutor({        backend: "sqlite",        plugins: [executionHistoryPlugin] as const,      });       yield* executor.executionHistory.record(        "org",        "shared",        call({          runId: "run-scope",          toolId: "shell",          status: "ok",          startedAt: "2026-05-29T11:00:00.000Z",        }),      );      yield* executor.executionHistory.record(        "user",        "shared",        call({          runId: "run-scope",          toolId: "browser",          status: "failed",          startedAt: "2026-05-29T11:01:00.000Z",        }),      );       const visibleShared = yield* executor.executionHistory.get("shared");      expect(visibleShared?.owner).toBe("user");      expect(visibleShared?.data.toolId).toBe("browser");       const scopedRows = yield* executor.executionHistory.query({        where: { runId: "run-scope" },        orderBy: [{ field: "startedAt" }],      });      expect(        scopedRows.map((entry) => [entry.key, String(entry.owner), entry.data.toolId]),

Lines 1–500 of 527.