History

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
  1. apps/cloud/drizzle.config.ts +17 −0
  2. apps/cloud/package.json +2 −0
  3. apps/cloud/src/mcp-session.e2e.node.test.ts +221 −0
  4. apps/cloud/src/mcp-session.ts +1 −4
  5. apps/cloud/src/services/db.schema.test.ts +61 −0
  6. apps/cloud/src/services/db.ts +6 −2
  7. notes/rls.md +80 −0
apps/cloud/drizzle.config.ts+17−0
@@ -1,7 +1,24 @@
1import { defineConfig } from "drizzle-kit";1import { 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.
9const DEFAULT_DEV_URL = "postgresql://postgres:postgres@127.0.0.1:5433/postgres";
10
11const 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
3export default defineConfig({17export default defineConfig({
4 schema: ["./src/services/schema.ts", "./src/services/executor-schema.ts"],18 schema: ["./src/services/schema.ts", "./src/services/executor-schema.ts"],
5 out: "./drizzle",19 out: "./drizzle",
6 dialect: "postgresql",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 "dev:vite": "CLOUDFLARE_INCLUDE_PROCESS_ENV=true op run --env-file=.env.op -- portless run --name executor-cloud vite dev",10 "dev:vite": "CLOUDFLARE_INCLUDE_PROCESS_ENV=true op run --env-file=.env.op -- portless run --name executor-cloud vite dev",
11 "db:schema": "node --import jiti/register ../../packages/core/cli/src/index.ts generate --config ./executor.config.ts --output ./src/services/executor-schema.ts",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 "db:generate": "drizzle-kit generate",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 "build": "vite build",15 "build": "vite build",
14 "preview": "vite preview",16 "preview": "vite preview",
15 "deploy": "op run --env-file=.env.production -- bun run build && op run --env-file=.env.production -- wrangler deploy",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
17import { describe, expect, it } from "@effect/vitest";
18import { Effect } from "effect";
19import { Client } from "@modelcontextprotocol/sdk/client/index.js";
20import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
21import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
22import type { ClientCapabilities } from "@modelcontextprotocol/sdk/types.js";
23
24import { createExecutorMcpServer } from "@executor/host-mcp";
25import { createExecutionEngine } from "@executor/execution";
26import {
27 ElicitationResponse,
28 FormElicitation,
29 Scope,
30 ScopeId,
31 collectSchemas,
32 createExecutor,
33 definePlugin,
34} from "@executor/sdk";
35import {
36 makePostgresAdapter,
37 makePostgresBlobStore,
38} from "@executor/storage-postgres";
39import { openApiPlugin } from "@executor/plugin-openapi";
40import { mcpPlugin } from "@executor/plugin-mcp";
41import { graphqlPlugin } from "@executor/plugin-graphql";
42import { workosVaultPlugin } from "@executor/plugin-workos-vault";
43
44import { DbService } from "./services/db";
45import { 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
52const EMPTY_INPUT_SCHEMA = {
53 type: "object",
54 properties: {},
55 additionalProperties: false,
56} as const;
57
58const 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
94const ELICITATION_CAPS: ClientCapabilities = {
95 elicitation: { form: {}, url: {} },
96};
97
98type BuildOptions = { readonly withElicitingPlugin?: boolean };
99
100const 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
127type 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.
132const 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
157const nextOrgId = (() => {
158 let seq = 0;
159 return () => `org_mcp_e2e_${++seq}_${crypto.randomUUID().slice(0, 8)}`;
160})();
161
162// ---------------------------------------------------------------------------
163// Tests
164// ---------------------------------------------------------------------------
165
166describe("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
22import { server } from "./env";22import { server } from "./env";
23import { AutumnService } from "./services/autumn";23import { AutumnService } from "./services/autumn";
24import { createScopedExecutor } from "./services/executor";24import { createScopedExecutor } from "./services/executor";
25import { DbService } from "./services/db";25import { DbService, combinedSchema } from "./services/db";
26import * as cloudSchema from "./services/schema";
27 26
28// ---------------------------------------------------------------------------27// ---------------------------------------------------------------------------
29// Types28// Types
@@ -54,8 +53,6 @@ const jsonRpcError = (status: number, code: number, message: string) => · 24 unmodified lines
54 headers: { "content-type": "application/json" },53 headers: { "content-type": "application/json" },
55 });54 });
56 55
57const combinedSchema = { ...cloudSchema };
58
59type DbHandle = DbServiceShape & { end: () => Promise<void> };56type 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
16import { describe, expect, it } from "vitest";
17import { drizzle } from "drizzle-orm/postgres-js";
18import postgres from "postgres";
19
20import * as cloudSchema from "./schema";
21import * as executorSchema from "./executor-schema";
22import { combinedSchema } from "./db";
23
24describe("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
20import * as executorSchema from "./executor-schema";20import * as executorSchema from "./executor-schema";
21import { server } from "../env";21import { server } from "../env";
22 22
23const 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.
27export const combinedSchema = { ...cloudSchema, ...executorSchema };
24 28
25// eslint-disable-next-line @typescript-eslint/no-explicit-any29// eslint-disable-next-line @typescript-eslint/no-explicit-any
26export type DrizzleDb = PgDatabase<any, any, any>;30export type DrizzleDb = PgDatabase<any, any, any>;
@@ -64,7 +68,7 @@ export class DbService extends Context.Tag("@executor/cloud/DbService")< · 37 unmodified lines
64 Effect.acquireRelease(68 Effect.acquireRelease(
65 Effect.sync((): DbServiceShape => {69 Effect.sync((): DbServiceShape => {
66 const sql = makeSql();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 ({ sql }) =>73 ({ sql }) =>
70 // Fire-and-forget: the Terminate round-trip sometimes hangs, and74 // 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
6enforced 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
9write. Only the worker talks to the DB; clients authenticate to the
10worker 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
12already locked.
13
14The only reason to turn RLS on is **defense-in-depth** — a belt against
15a future bug where a drizzle query forgets to filter by `scope_id`.
16Without RLS, such a query silently returns cross-tenant rows. With RLS,
17it returns zero rows.
18
19## When RLS matters
20
21RLS earns its weight when untrusted clients connect directly to Postgres
22with per-user credentials (the Supabase / PostgREST model). Not our
23shape.
24
25Consider 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
35Scope merging (`ScopeStack` in `notes/scopes.md`) means reads fan out
36across multiple scopes while writes target exactly one. RLS policies
37mirror that asymmetry by reading two settings per request:
38
39```sql
40CREATE POLICY scope_read ON source FOR SELECT
41 USING (scope_id = ANY(current_setting('app.scope_chain', true)::text[]));
42
43CREATE POLICY scope_write ON source FOR INSERT
44 WITH CHECK (scope_id = current_setting('app.write_scope', true));
45
46CREATE 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
50ALTER TABLE source ENABLE ROW LEVEL SECURITY;
51```
52
53And app-side, once per request inside a transaction (so `SET LOCAL`
54takes):
55
56```ts
57yield* sql`SET LOCAL app.scope_chain = ${toPgArray(chain)}`;
58yield* sql`SET LOCAL app.write_scope = ${writeScope}`;
59```
60
61`current_setting('…', true)` returns null when the setting is missing,
62which is what you want — a connection that forgets to configure the
63scope 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
671. 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.
712. 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.
743. 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
78Ballpark effort: a day including the migration + the Effect layer +
79tests. Worth doing only if we hit a scope-leak bug in prod or change
80the access model.