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

Showing the changes to apps/cloud/src (4 of 7 files). Show all

Unified · Split

4 files
  1. apps/cloud/src/mcp-session.e2e.node.test.ts +221 −0
  2. apps/cloud/src/mcp-session.ts +1 −4
  3. apps/cloud/src/services/db.schema.test.ts +61 −0
  4. apps/cloud/src/services/db.ts +6 −2
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
2222 import { server } from "./env";
2323 import { AutumnService } from "./services/autumn";
2424 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";
2726
2827 // ---------------------------------------------------------------------------
2928 // Types
@@ -54,8 +53,6 @@ const jsonRpcError = (status: number, code: number, message: string) => · 24 unmodified lines
5453 headers: { "content-type": "application/json" },
5554 });
5655
57−const combinedSchema = { ...cloudSchema };
58−
5956 type DbHandle = DbServiceShape & { end: () => Promise<void> };
6057
6158 /**
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
2020 import * as executorSchema from "./executor-schema";
2121 import { server } from "../env";
2222
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 };
2428
2529 // eslint-disable-next-line @typescript-eslint/no-explicit-any
2630 export type DrizzleDb = PgDatabase<any, any, any>;
@@ -64,7 +68,7 @@ export class DbService extends Context.Tag("@executor/cloud/DbService")< · 37 unmodified lines
6468 Effect.acquireRelease(
6569 Effect.sync((): DbServiceShape => {
6670 const sql = makeSql();
67− return { sql, db: drizzle(sql, { schema }) as DrizzleDb };
71+ return { sql, db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb };
6872 }),
6973 ({ sql }) =>
7074 // Fire-and-forget: the Terminate round-trip sometimes hangs, and