// ---------------------------------------------------------------------------// MCP Session Durable Object — holds MCP server + engine per session// --------------------------------------------------------------------------- import { DurableObject, env } from "cloudflare:workers";import { Data, Effect, Layer } from "effect";import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";import { WorkerTransport, type TransportState } from "agents/mcp";import { drizzle } from "drizzle-orm/postgres-js";import postgres from "postgres"; import { createExecutorMcpServer } from "@executor/host-mcp";import { makeDynamicWorkerExecutor } from "@executor/runtime-dynamic-worker";import type { DrizzleDb } from "@executor/storage-postgres";import * as sharedSchema from "@executor/storage-postgres/schema"; import { UserStoreService } from "./auth/context";import { server } from "./env";import { createOrgExecutor } from "./services/executor";import { DbService } from "./services/db";import * as cloudSchema from "./services/schema"; // ---------------------------------------------------------------------------// Types// --------------------------------------------------------------------------- export type McpSessionInit = { organizationId: string;}; const HEARTBEAT_MS = 30 * 1000;const SESSION_TIMEOUT_MS = 5 * 60 * 1000; // ---------------------------------------------------------------------------// Errors// --------------------------------------------------------------------------- class OrganizationNotFoundError extends Data.TaggedError("OrganizationNotFoundError")<{ readonly organizationId: string;}> {} // ---------------------------------------------------------------------------// Helpers// --------------------------------------------------------------------------- const jsonRpcError = (status: number, code: number, message: string) => new Response(JSON.stringify({ jsonrpc: "2.0", error: { code, message }, id: null }), { status, headers: { "content-type": "application/json" }, }); const combinedSchema = { ...sharedSchema, ...cloudSchema }; /** * Create a long-lived DB connection for the DO lifetime. * * Unlike the per-request `DbService.Live` used in `api.ts`, this connection * stays open across requests within the DO. DOs have a single-threaded * execution model — there's no cross-request socket reuse issue because * only one request runs at a time. The connection is closed when the DO * cleans up (timeout or eviction). */const makeLongLivedDb = (): { db: DrizzleDb; end: () => Promise<void> } => { const connectionString = env.HYPERDRIVE?.connectionString ?? server.DATABASE_URL; const sql = postgres(connectionString, { max: 1, idle_timeout: 20, max_lifetime: 300, connect_timeout: 10, onnotice: () => undefined, }); return { db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb, end: () => sql.end({ timeout: 0 }).catch(() => undefined), };}; // ---------------------------------------------------------------------------// Durable Object// --------------------------------------------------------------------------- export class McpSessionDO extends DurableObject { private mcpServer: McpServer | null = null; private transport: WorkerTransport | null = null; private initialized = false; private lastActivityMs = 0; private dbHandle: { db: DrizzleDb; end: () => Promise<void> } | null = null; private makeStorage() { return { get: async (): Promise<TransportState | undefined> => { return await this.ctx.storage.get<TransportState>("transport"); }, set: async (state: TransportState): Promise<void> => { await this.ctx.storage.put("transport", state); }, }; } async init(token: McpSessionInit): Promise<void> { if (this.initialized) return; // Create a long-lived DB connection for the DO's lifetime this.dbHandle = makeLongLivedDb(); const DbLive = Layer.succeed(DbService, this.dbHandle.db); const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive)); const Services = Layer.mergeAll(DbLive, UserStoreLive); const program = Effect.gen(function* () { const users = yield* UserStoreService; const org = yield* users.use((store) => store.getOrganization(token.organizationId)); if (!org) return yield* new OrganizationNotFoundError({ organizationId: token.organizationId }); const executor = yield* createOrgExecutor(org.id, org.name); const codeExecutor = makeDynamicWorkerExecutor({ loader: env.LOADER }); return yield* Effect.promise(() => createExecutorMcpServer({ executor, codeExecutor })); }).pipe(Effect.provide(Services)); this.mcpServer = await Effect.runPromise(program); this.transport = new WorkerTransport({ sessionIdGenerator: () => this.ctx.id.toString(), storage: this.makeStorage(), }); await this.mcpServer.connect(this.transport); this.initialized = true; this.lastActivityMs = Date.now(); await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS); } async handleRequest(request: Request): Promise<Response> { if (!this.initialized || !this.transport) { return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect"); } this.lastActivityMs = Date.now(); try { return await this.transport.handleRequest(request); } catch (err) { console.error("[mcp-session] handleRequest error:", err instanceof Error ? err.stack : err); return jsonRpcError(500, -32603, "Internal error"); } } async alarm(): Promise<void> { const idleMs = Date.now() - this.lastActivityMs; if (idleMs >= SESSION_TIMEOUT_MS) { await this.cleanup(); return; } await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS); } private async cleanup(): Promise<void> { if (this.transport) { await this.transport.close().catch(() => undefined); this.transport = null; } if (this.mcpServer) { await this.mcpServer.close().catch(() => undefined); this.mcpServer = null; } if (this.dbHandle) { await this.dbHandle.end(); this.dbHandle = null; } this.initialized = false; }}
// ---------------------------------------------------------------------------
// MCP Session Durable Object — holds MCP server + engine per session
// ---------------------------------------------------------------------------
import { DurableObject, env } from "cloudflare:workers";
import { Data, Effect, Layer } from "effect";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { WorkerTransport, type TransportState } from "agents/mcp";
import { drizzle } from "drizzle-orm/postgres-js";
import postgres from "postgres";
import { createExecutorMcpServer } from "@executor/host-mcp";
import { makeDynamicWorkerExecutor } from "@executor/runtime-dynamic-worker";
import type { DrizzleDb } from "@executor/storage-postgres";
import * as sharedSchema from "@executor/storage-postgres/schema";
import { UserStoreService } from "./auth/context";
import { server } from "./env";
import { createOrgExecutor } from "./services/executor";
import { DbService } from "./services/db";
import * as cloudSchema from "./services/schema";
// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------
export type McpSessionInit = {
organizationId: string;
};
const HEARTBEAT_MS = 30 * 1000;
const SESSION_TIMEOUT_MS = 5 * 60 * 1000;
// ---------------------------------------------------------------------------
// Errors
// ---------------------------------------------------------------------------
class OrganizationNotFoundError extends Data.TaggedError("OrganizationNotFoundError")<{
readonly organizationId: string;
}> {}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
const jsonRpcError = (status: number, code: number, message: string) =>
new Response(JSON.stringify({ jsonrpc: "2.0", error: { code, message }, id: null }), {
status,
headers: { "content-type": "application/json" },
});
const combinedSchema = { ...sharedSchema, ...cloudSchema };
/**
* Create a long-lived DB connection for the DO lifetime.
*
* Unlike the per-request `DbService.Live` used in `api.ts`, this connection
* stays open across requests within the DO. DOs have a single-threaded
* execution model — there's no cross-request socket reuse issue because
* only one request runs at a time. The connection is closed when the DO
* cleans up (timeout or eviction).
*/
const makeLongLivedDb = (): { db: DrizzleDb; end: () => Promise<void> } => {
const connectionString = env.HYPERDRIVE?.connectionString ?? server.DATABASE_URL;
const sql = postgres(connectionString, {
max: 1,
idle_timeout: 20,
max_lifetime: 300,
connect_timeout: 10,
onnotice: () => undefined,
});
return {
db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb,
end: () => sql.end({ timeout: 0 }).catch(() => undefined),
};
};
// ---------------------------------------------------------------------------
// Durable Object
// ---------------------------------------------------------------------------
export class McpSessionDO extends DurableObject {
private mcpServer: McpServer | null = null;
private transport: WorkerTransport | null = null;
private initialized = false;
private lastActivityMs = 0;
private dbHandle: { db: DrizzleDb; end: () => Promise<void> } | null = null;
private makeStorage() {
return {
get: async (): Promise<TransportState | undefined> => {
return await this.ctx.storage.get<TransportState>("transport");
},
set: async (state: TransportState): Promise<void> => {
await this.ctx.storage.put("transport", state);
},
};
}
async init(token: McpSessionInit): Promise<void> {
if (this.initialized) return;
// Create a long-lived DB connection for the DO's lifetime
this.dbHandle = makeLongLivedDb();
const DbLive = Layer.succeed(DbService, this.dbHandle.db);
const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));
const Services = Layer.mergeAll(DbLive, UserStoreLive);
const program = Effect.gen(function* () {
const users = yield* UserStoreService;
const org = yield* users.use((store) => store.getOrganization(token.organizationId));
if (!org)
return yield* new OrganizationNotFoundError({ organizationId: token.organizationId });
const executor = yield* createOrgExecutor(org.id, org.name);
const codeExecutor = makeDynamicWorkerExecutor({ loader: env.LOADER });
return yield* Effect.promise(() => createExecutorMcpServer({ executor, codeExecutor }));
}).pipe(Effect.provide(Services));
this.mcpServer = await Effect.runPromise(program);
this.transport = new WorkerTransport({
sessionIdGenerator: () => this.ctx.id.toString(),
storage: this.makeStorage(),
});
await this.mcpServer.connect(this.transport);
this.initialized = true;
this.lastActivityMs = Date.now();
await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS);
}
async handleRequest(request: Request): Promise<Response> {
if (!this.initialized || !this.transport) {
return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");
}
this.lastActivityMs = Date.now();
try {
return await this.transport.handleRequest(request);
} catch (err) {
console.error("[mcp-session] handleRequest error:", err instanceof Error ? err.stack : err);
return jsonRpcError(500, -32603, "Internal error");
}
}
async alarm(): Promise<void> {
const idleMs = Date.now() - this.lastActivityMs;
if (idleMs >= SESSION_TIMEOUT_MS) {
await this.cleanup();
return;
}
await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS);
}
private async cleanup(): Promise<void> {
if (this.transport) {
await this.transport.close().catch(() => undefined);
this.transport = null;
}
if (this.mcpServer) {
await this.mcpServer.close().catch(() => undefined);
this.mcpServer = null;
}
if (this.dbHandle) {
await this.dbHandle.end();
this.dbHandle = null;
}
this.initialized = false;
}
}