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
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