Share cloud drizzle schema across DbService + MCP session

mcp-session.ts was building its own combinedSchema from cloudSchema only, dropping every executor-schema table — every /mcp request failed with [storage-drizzle] unknown model "source". Export combinedSchema from services/db.ts and consume it in both DbService.Live and the MCP Durable Object so drift can't happen. Regression coverage: - db.schema.test.ts asserts combinedSchema (and drizzle's _.fullSchema) contain every cloud + executor export - mcp-session.e2e.node.test.ts drives a real @modelcontextprotocol/sdk Client through listTools / callTool / form elicitation round-trips

Rhys Sullivan committed 6 months ago

Commit 38d662b1 parent 807df1d7 files · +388 −6Browse files

Unified · Split

7 files

apps/cloud/drizzle.config.ts +17 −0
apps/cloud/package.json +2 −0
apps/cloud/src/mcp-session.e2e.node.test.ts +221 −0
apps/cloud/src/mcp-session.ts +1 −4
apps/cloud/src/services/db.schema.test.ts +61 −0
apps/cloud/src/services/db.ts +6 −2
notes/rls.md +80 −0

apps/cloud/drizzle.config.ts+17−0

@@ -1,7 +1,24 @@
1	1	import { defineConfig } from "drizzle-kit";
2	2
	3	+// drizzle-kit studio reads `dbCredentials.url`; everything else (generate,
	4	+// migrate) ignores it. Default to the local PGlite socket started by
	5	+// `bun run dev:db`; override via `DATABASE_URL` for prod studio sessions.
	6	+// drizzle-kit uses node-postgres (`pg`) for studio and the `ssl` option in
	7	+// dbCredentials doesn't reliably reach the pool — append `sslmode=require`
	8	+// directly to the URL instead, which `pg` honours.
	9	+const DEFAULT_DEV_URL = "postgresql://postgres:postgres@127.0.0.1:5433/postgres";
	10	+
	11	+const withSslMode = (url: string): string => {
	12	+ if (url.includes("127.0.0.1") \|\| url.includes("localhost")) return url;
	13	+ if (/[?&]sslmode=/.test(url)) return url;
	14	+ return url + (url.includes("?") ? "&" : "?") + "sslmode=require";
	15	+};
	16	+
3	17	export default defineConfig({
4	18	schema: ["./src/services/schema.ts", "./src/services/executor-schema.ts"],
5	19	out: "./drizzle",
6	20	dialect: "postgresql",
	21	+ dbCredentials: {
	22	+ url: withSslMode(process.env.DATABASE_URL ?? DEFAULT_DEV_URL),
	23	+ },
7	24	});

apps/cloud/package.json+2−0

@@ -10,6 +10,8 @@ · 9 unmodified lines
10	10	"dev:vite": "CLOUDFLARE_INCLUDE_PROCESS_ENV=true op run --env-file=.env.op -- portless run --name executor-cloud vite dev",
11	11	"db:schema": "node --import jiti/register ../../packages/core/cli/src/index.ts generate --config ./executor.config.ts --output ./src/services/executor-schema.ts",
12	12	"db:generate": "drizzle-kit generate",
	13	+ "db:studio": "drizzle-kit studio",
	14	+ "db:studio:prod": "op run --env-file=.env.production -- drizzle-kit studio",
13	15	"build": "vite build",
14	16	"preview": "vite preview",
15	17	"deploy": "op run --env-file=.env.production -- bun run build && op run --env-file=.env.production -- wrangler deploy",

apps/cloud/src/mcp-session.e2e.node.test.ts+221−0

@@ -0,0 +1,221 @@
	1	+// End-to-end coverage for the cloud MCP server.
	2	+//
	3	+// The `McpSessionDO` in mcp-session.ts wires several things that previously
	4	+// had zero integration coverage:
	5	+// - `createScopedExecutor` against a real drizzle adapter (the 2026-04-16
	6	+// prod outage was a schema spread bug here; see services/db.schema.test.ts)
	7	+// - `createExecutionEngine` with an in-process code executor
	8	+// - `createExecutorMcpServer` for the MCP request surface
	9	+// - Real `@modelcontextprotocol/sdk` Client → server round-trips
	10	+//
	11	+// This test replicates the DO's init path (minus the WorkerTransport and
	12	+// Durable Object routing, which are thin CF plumbing) and drives it with a
	13	+// real MCP Client over in-memory transports. If any of the wiring drifts —
	14	+// schema, plugin list, engine contract, MCP handshake — these tests fail
	15	+// before prod does.
	16	+
	17	+import { describe, expect, it } from "@effect/vitest";
	18	+import { Effect } from "effect";
	19	+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
	20	+import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
	21	+import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
	22	+import type { ClientCapabilities } from "@modelcontextprotocol/sdk/types.js";
	23	+
	24	+import { createExecutorMcpServer } from "@executor/host-mcp";
	25	+import { createExecutionEngine } from "@executor/execution";
	26	+import {
	27	+ ElicitationResponse,
	28	+ FormElicitation,
	29	+ Scope,
	30	+ ScopeId,
	31	+ collectSchemas,
	32	+ createExecutor,
	33	+ definePlugin,
	34	+} from "@executor/sdk";
	35	+import {
	36	+ makePostgresAdapter,
	37	+ makePostgresBlobStore,
	38	+} from "@executor/storage-postgres";
	39	+import { openApiPlugin } from "@executor/plugin-openapi";
	40	+import { mcpPlugin } from "@executor/plugin-mcp";
	41	+import { graphqlPlugin } from "@executor/plugin-graphql";
	42	+import { workosVaultPlugin } from "@executor/plugin-workos-vault";
	43	+
	44	+import { DbService } from "./services/db";
	45	+import { makeFakeVaultClient } from "./services/__test-harness__/api-harness";
	46	+
	47	+// ---------------------------------------------------------------------------
	48	+// Test-only plugin: exposes one in-memory tool that elicits once. Lets the
	49	+// eliciting test drive the real engine + sandbox rather than a stub engine.
	50	+// ---------------------------------------------------------------------------
	51	+
	52	+const EMPTY_INPUT_SCHEMA = {
	53	+ type: "object",
	54	+ properties: {},
	55	+ additionalProperties: false,
	56	+} as const;
	57	+
	58	+const elicitingTestPlugin = definePlugin(() => ({
	59	+ id: "eliciting-test" as const,
	60	+ storage: () => ({}),
	61	+ staticSources: () => [
	62	+ {
	63	+ id: "e2e",
	64	+ kind: "in-memory",
	65	+ name: "E2E Test",
	66	+ tools: [
	67	+ {
	68	+ name: "needsApproval",
	69	+ description: "Tool that asks the caller to approve before returning.",
	70	+ inputSchema: EMPTY_INPUT_SCHEMA,
	71	+ handler: ({ elicit }: { elicit: (r: FormElicitation) => Effect.Effect<typeof ElicitationResponse.Type, unknown> }) =>
	72	+ Effect.gen(function* () {
	73	+ const response = yield* elicit(
	74	+ new FormElicitation({
	75	+ message: "Approve?",
	76	+ requestedSchema: {
	77	+ type: "object",
	78	+ properties: { approved: { type: "boolean" } },
	79	+ },
	80	+ }),
	81	+ );
	82	+ return { action: response.action, content: response.content ?? null };
	83	+ }).pipe(Effect.orDie),
	84	+ },
	85	+ ],
	86	+ },
	87	+ ],
	88	+}));
	89	+
	90	+// ---------------------------------------------------------------------------
	91	+// Session harness — mirrors McpSessionDO.init() minus the WorkerTransport
	92	+// ---------------------------------------------------------------------------
	93	+
	94	+const ELICITATION_CAPS: ClientCapabilities = {
	95	+ elicitation: { form: {}, url: {} },
	96	+};
	97	+
	98	+type BuildOptions = { readonly withElicitingPlugin?: boolean };
	99	+
	100	+const buildScopedExecutor = (
	101	+ scopeId: string,
	102	+ scopeName: string,
	103	+ options: BuildOptions = {},
	104	+) =>
	105	+ Effect.gen(function* () {
	106	+ const { db } = yield* DbService;
	107	+ const basePlugins = [
	108	+ openApiPlugin(),
	109	+ mcpPlugin({ dangerouslyAllowStdioMCP: false }),
	110	+ graphqlPlugin(),
	111	+ workosVaultPlugin({ client: makeFakeVaultClient() }),
	112	+ ] as const;
	113	+ const plugins = options.withElicitingPlugin
	114	+ ? ([...basePlugins, elicitingTestPlugin()] as const)
	115	+ : basePlugins;
	116	+ const schema = collectSchemas(plugins);
	117	+ const adapter = makePostgresAdapter({ db, schema });
	118	+ const blobs = makePostgresBlobStore({ db });
	119	+ const scope = new Scope({
	120	+ id: ScopeId.make(scopeId),
	121	+ name: scopeName,
	122	+ createdAt: new Date(),
	123	+ });
	124	+ return yield* createExecutor({ scope, adapter, blobs, plugins });
	125	+ });
	126	+
	127	+type McpSession = { readonly client: Client };
	128	+
	129	+// Builds a scope, wires a real execution engine + MCP server, and yields
	130	+// them connected to an in-memory MCP client. Shaped as an acquireRelease so
	131	+// the transport teardown is guaranteed when the test scope closes.
	132	+const openSession = (
	133	+ orgId: string,
	134	+ options: BuildOptions & { readonly caps?: ClientCapabilities } = {},
	135	+) =>
	136	+ Effect.acquireRelease(
	137	+ Effect.gen(function* () {
	138	+ const executor = yield* buildScopedExecutor(orgId, `Org ${orgId}`, options);
	139	+ const engine = createExecutionEngine({ executor });
	140	+ const mcpServer = yield* Effect.promise(() => createExecutorMcpServer({ engine }));
	141	+ const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
	142	+ const client = new Client(
	143	+ { name: "cloud-e2e-test", version: "1.0.0" },
	144	+ { capabilities: options.caps ?? ELICITATION_CAPS },
	145	+ );
	146	+ yield* Effect.promise(() => mcpServer.connect(serverTransport));
	147	+ yield* Effect.promise(() => client.connect(clientTransport));
	148	+ return { client, clientTransport, serverTransport };
	149	+ }),
	150	+ ({ clientTransport, serverTransport }) =>
	151	+ Effect.promise(async () => {
	152	+ await clientTransport.close().catch(() => undefined);
	153	+ await serverTransport.close().catch(() => undefined);
	154	+ }),
	155	+ ).pipe(Effect.map(({ client }) => ({ client })));
	156	+
	157	+const nextOrgId = (() => {
	158	+ let seq = 0;
	159	+ return () => `org_mcp_e2e_${++seq}_${crypto.randomUUID().slice(0, 8)}`;
	160	+})();
	161	+
	162	+// ---------------------------------------------------------------------------
	163	+// Tests
	164	+// ---------------------------------------------------------------------------
	165	+
	166	+describe("cloud MCP session end-to-end", () => {
	167	+ it.effect("initializes and exposes the execute tool to the MCP client", () =>
	168	+ Effect.gen(function* () {
	169	+ const { client } = yield* openSession(nextOrgId());
	170	+ const tools = yield* Effect.promise(() => client.listTools());
	171	+ const names = tools.tools.map((t) => t.name);
	172	+ expect(names).toContain("execute");
	173	+ }).pipe(Effect.provide(DbService.Live), Effect.scoped),
	174	+ );
	175	+
	176	+ it.effect("runs user code via the execute tool end-to-end", () =>
	177	+ Effect.gen(function* () {
	178	+ const { client } = yield* openSession(nextOrgId());
	179	+ const result = yield* Effect.promise(() =>
	180	+ client.callTool({ name: "execute", arguments: { code: "return 1 + 2" } }),
	181	+ );
	182	+ expect(result.isError).toBeFalsy();
	183	+ const text = (result.content as Array<{ type: string; text: string }>)[0]!.text;
	184	+ expect(text).toContain("3");
	185	+ }).pipe(Effect.provide(DbService.Live), Effect.scoped),
	186	+ );
	187	+
	188	+ // Isolates the drizzle adapter path so a schema spread drift surfaces as
	189	+ // a raw "unknown model" error. The prod outage on 2026-04-16 would have
	190	+ // thrown at `executor.sources.list()` when the MCP session's drizzle
	191	+ // instance lost the executor-schema tables.
	192	+ it.effect("exercises the drizzle adapter directly via executor.sources.list", () =>
	193	+ Effect.gen(function* () {
	194	+ const executor = yield* buildScopedExecutor(nextOrgId(), "drizzle-probe");
	195	+ const sources = yield* executor.sources.list();
	196	+ expect(Array.isArray(sources)).toBe(true);
	197	+ }).pipe(Effect.provide(DbService.Live), Effect.scoped),
	198	+ );
	199	+
	200	+ it.effect("bridges a form elicitation from engine to client and back", () =>
	201	+ Effect.gen(function* () {
	202	+ const { client } = yield* openSession(nextOrgId(), { withElicitingPlugin: true });
	203	+
	204	+ client.setRequestHandler(ElicitRequestSchema, async () => ({
	205	+ action: "accept" as const,
	206	+ content: { approved: true },
	207	+ }));
	208	+
	209	+ const result = yield* Effect.promise(() =>
	210	+ client.callTool({
	211	+ name: "execute",
	212	+ arguments: { code: "return await tools.e2e.needsApproval({});" },
	213	+ }),
	214	+ );
	215	+ expect(result.isError).toBeFalsy();
	216	+ const text = (result.content as Array<{ type: string; text: string }>)[0]!.text;
	217	+ expect(text).toContain("accept");
	218	+ expect(text).toContain("approved");
	219	+ }).pipe(Effect.provide(DbService.Live), Effect.scoped),
	220	+ );
	221	+});

apps/cloud/src/mcp-session.ts+1−4

@@ -22,8 +22,7 @@ import { WorkOSAuth } from "./auth/workos"; · 21 unmodified lines
22	22	import { server } from "./env";
23	23	import { AutumnService } from "./services/autumn";
24	24	import { createScopedExecutor } from "./services/executor";
25		−import { DbService } from "./services/db";
26		−import * as cloudSchema from "./services/schema";
	25	+import { DbService, combinedSchema } from "./services/db";
27	26
28	27	// ---------------------------------------------------------------------------
29	28	// Types
@@ -54,8 +53,6 @@ const jsonRpcError = (status: number, code: number, message: string) => · 24 unmodified lines
54	53	headers: { "content-type": "application/json" },
55	54	});
56	55
57		−const combinedSchema = { ...cloudSchema };
58		−
59	56	type DbHandle = DbServiceShape & { end: () => Promise<void> };
60	57
61	58	/**

apps/cloud/src/services/db.schema.test.ts+61−0

@@ -0,0 +1,61 @@
	1	+// Regression: every cloud drizzle instance must be constructed with the
	2	+// merge of cloud + executor schemas. If ANY call site builds its own
	3	+// `{ schema }` without spreading both, `db._.fullSchema` comes back
	4	+// missing tables and the drizzle adapter throws
	5	+// `unknown model "source"` at the first request that touches scoped data.
	6	+//
	7	+// We hit this in prod (MCP endpoint) when `mcp-session.ts` built
	8	+// `combinedSchema = { ...cloudSchema }` and forgot to spread the executor
	9	+// schema — web-app requests worked because they went through
	10	+// `DbService.Live`, but the MCP Durable Object constructed its own.
	11	+//
	12	+// This test asserts the only sanctioned constant (`combinedSchema` from
	13	+// `./db`) actually contains every executor-schema export AND that drizzle
	14	+// surfaces them all under `db._.fullSchema`.
	15	+
	16	+import { describe, expect, it } from "vitest";
	17	+import { drizzle } from "drizzle-orm/postgres-js";
	18	+import postgres from "postgres";
	19	+
	20	+import * as cloudSchema from "./schema";
	21	+import * as executorSchema from "./executor-schema";
	22	+import { combinedSchema } from "./db";
	23	+
	24	+describe("combinedSchema", () => {
	25	+ it("spreads every cloud + executor schema export", () => {
	26	+ const expected = new Set([
	27	+ ...Object.keys(cloudSchema),
	28	+ ...Object.keys(executorSchema),
	29	+ ]);
	30	+ for (const key of expected) {
	31	+ expect(combinedSchema, `combinedSchema missing "${key}"`).toHaveProperty(key);
	32	+ }
	33	+ });
	34	+
	35	+ // Executor-schema drives the scope-sharded tables (source/tool/etc). If
	36	+ // any of these go missing the drizzle adapter's `getTable` lookup throws.
	37	+ it("includes scope-sharded executor tables", () => {
	38	+ for (const key of Object.keys(executorSchema)) {
	39	+ expect(combinedSchema, `combinedSchema missing "${key}"`).toHaveProperty(key);
	40	+ }
	41	+ });
	42	+
	43	+ // The prod bug was actually at the drizzle layer: spread + __exportAll
	44	+ // getters could theoretically drop tables if evaluated before their
	45	+ // declarations. Construct a drizzle instance and walk its fullSchema
	46	+ // to catch that class of bug too.
	47	+ it("drizzle(combinedSchema) exposes every table under _.fullSchema", () => {
	48	+ // postgres() lazily connects — safe to build with a dummy url, we
	49	+ // never .query() so no socket is opened.
	50	+ const sql = postgres("postgres://u:p@127.0.0.1:1/x", { max: 1 });
	51	+ try {
	52	+ const db = drizzle(sql, { schema: combinedSchema });
	53	+ const fullSchema = (db as unknown as { _: { fullSchema: Record<string, unknown> } })._.fullSchema;
	54	+ for (const key of Object.keys(executorSchema)) {
	55	+ expect(fullSchema, `fullSchema missing "${key}"`).toHaveProperty(key);
	56	+ }
	57	+ } finally {
	58	+ sql.end({ timeout: 0 }).catch(() => undefined);
	59	+ }
	60	+ });
	61	+});

apps/cloud/src/services/db.ts+6−2

@@ -20,7 +20,11 @@ import * as cloudSchema from "./schema"; · 19 unmodified lines
20	20	import * as executorSchema from "./executor-schema";
21	21	import { server } from "../env";
22	22
23		−const schema = { ...cloudSchema, ...executorSchema };
	23	+// Exported so every drizzle() call in the cloud app shares one schema
	24	+// object. Historically `mcp-session.ts` built its own and forgot to spread
	25	+// `executorSchema`, producing runtime "unknown model source" errors that
	26	+// only surfaced in prod. See apps/cloud/src/services/db.schema.test.ts.
	27	+export const combinedSchema = { ...cloudSchema, ...executorSchema };
24	28
25	29	// eslint-disable-next-line @typescript-eslint/no-explicit-any
26	30	export type DrizzleDb = PgDatabase<any, any, any>;
@@ -64,7 +68,7 @@ export class DbService extends Context.Tag("@executor/cloud/DbService")< · 37 unmodified lines
64	68	Effect.acquireRelease(
65	69	Effect.sync((): DbServiceShape => {
66	70	const sql = makeSql();
67		− return { sql, db: drizzle(sql, { schema }) as DrizzleDb };
	71	+ return { sql, db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb };
68	72	}),
69	73	({ sql }) =>
70	74	// Fire-and-forget: the Terminate round-trip sometimes hangs, and

notes/rls.md+80−0

@@ -0,0 +1,80 @@
	1	+# Postgres RLS — optional, not required
	2	+
	3	+## TL;DR
	4	+
	5	+RLS is not needed for the current architecture. Tenant isolation is
	6	+enforced at the application layer: every query against a scoped table
	7	+(see `notes/scopes.md`) is routed through the scope adapter which ANDs
	8	+`scope_id = <current scope>` into every read and stamps it on every
	9	+write. Only the worker talks to the DB; clients authenticate to the
	10	+worker via WorkOS JWT, never to Postgres directly. One DB role
	11	+(`postgres`) owns every row. RLS would be a second lock on a door that's
	12	+already locked.
	13	+
	14	+The only reason to turn RLS on is defense-in-depth — a belt against
	15	+a future bug where a drizzle query forgets to filter by `scope_id`.
	16	+Without RLS, such a query silently returns cross-tenant rows. With RLS,
	17	+it returns zero rows.
	18	+
	19	+## When RLS matters
	20	+
	21	+RLS earns its weight when untrusted clients connect directly to Postgres
	22	+with per-user credentials (the Supabase / PostgREST model). Not our
	23	+shape.
	24	+
	25	+Consider turning it on if any of these change:
	26	+
	27	+- Clients start hitting Postgres directly (JWT-authed session, row-level
	28	+ policies as the only authorisation layer).
	29	+- We run multiple DB roles that need row-level isolation between them.
	30	+- We decide the cost of a missed `scope_id` filter in app code is high
	31	+ enough to want a DB-level backstop.
	32	+
	33	+## How it would work with scope merging
	34	+
	35	+Scope merging (`ScopeStack` in `notes/scopes.md`) means reads fan out
	36	+across multiple scopes while writes target exactly one. RLS policies
	37	+mirror that asymmetry by reading two settings per request:
	38	+
	39	+```sql
	40	+CREATE POLICY scope_read ON source FOR SELECT
	41	+ USING (scope_id = ANY(current_setting('app.scope_chain', true)::text[]));
	42	+
	43	+CREATE POLICY scope_write ON source FOR INSERT
	44	+ WITH CHECK (scope_id = current_setting('app.write_scope', true));
	45	+
	46	+CREATE POLICY scope_update ON source FOR UPDATE
	47	+ USING (scope_id = current_setting('app.write_scope', true))
	48	+ WITH CHECK (scope_id = current_setting('app.write_scope', true));
	49	+
	50	+ALTER TABLE source ENABLE ROW LEVEL SECURITY;
	51	+```
	52	+
	53	+And app-side, once per request inside a transaction (so `SET LOCAL`
	54	+takes):
	55	+
	56	+```ts
	57	+yield* sql`SET LOCAL app.scope_chain = ${toPgArray(chain)}`;
	58	+yield* sql`SET LOCAL app.write_scope = ${writeScope}`;
	59	+```
	60	+
	61	+`current_setting('…', true)` returns null when the setting is missing,
	62	+which is what you want — a connection that forgets to configure the
	63	+scope sees zero rows, not an error that might be handled into a 500.
	64	+
	65	+## Rollout sketch (if we ever decide to enable it)
	66	+
	67	+1. Add a migration that `ENABLE ROW LEVEL SECURITY` on every scoped
	68	+ table and installs the read/write policies above. The list of scoped
	69	+ tables is whatever `executor-schema.ts` exports with a `scope_id`
	70	+ column — same set the scope adapter already knows about.
	71	+2. Add an Effect layer that wraps the DbService to issue the two `SET
	72	+ LOCAL`s at connection acquire. `postgres.js`'s `sql.begin` already
	73	+ gives us the transaction boundary.
	74	+3. Leave `BYPASSRLS` off the worker's role so we actually get the
	75	+ enforcement; grant it on the role used by the migration / admin
	76	+ scripts.
	77	+
	78	+Ballpark effort: a day including the migration + the Effect layer +
	79	+tests. Worth doing only if we hit a scope-leak bug in prod or change
	80	+the access model.