// ---------------------------------------------------------------------------// MCP Session Durable Object — holds MCP server + engine per session// --------------------------------------------------------------------------- import * as Cloudflare from "alchemy/Cloudflare/Workers/Runtime";import { Cause, Data, Effect, Layer } from "effect";import * as Tracer from "effect/Tracer";import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";import { drizzle } from "drizzle-orm/postgres-js";import postgres from "postgres"; import { createExecutorMcpServer } from "@executor-js/host-mcp";import { buildExecuteDescription } from "@executor-js/execution";import type { DrizzleDb, DbServiceShape } from "./services/db"; // Import directly from core-shared-services, NOT from ./api/layers.ts.// The full layers module pulls in `auth/handlers.ts` → `@tanstack/react-start/server`,// which uses a `#tanstack-start-entry` subpath specifier that breaks module// load under vitest-pool-workers. The DO only needs the core two services// (WorkOSAuth + AutumnService), so we import them from the tight module.import { CoreSharedServices } from "./api/core-shared-services";import { UserStoreService } from "./auth/context";import { resolveOrganization } from "./auth/resolve-organization";import { DbService, combinedSchema, resolveConnectionString } from "./services/db";import { makeExecutionStack } from "./services/execution-stack";import {  makeMcpWorkerTransport,  type McpTransportState,  type McpWorkerTransport,} from "./services/mcp-worker-transport";import { DoTelemetryLive } from "./services/telemetry";import { captureCause } from "./observability"; // ---------------------------------------------------------------------------// Types// --------------------------------------------------------------------------- export type McpSessionInit = {  organizationId: string;  userId: string;}; export type IncomingTraceHeaders = {  readonly traceparent?: string;  readonly tracestate?: string;  readonly baggage?: string;}; export type McpSessionShape = {  readonly init: (    token: McpSessionInit,    incoming?: IncomingTraceHeaders,  ) => Effect.Effect<void, unknown>;  readonly handleRequest: (request: Request) => Effect.Effect<Response, unknown>;  readonly clearSession: (incoming?: IncomingTraceHeaders) => Effect.Effect<void, unknown>;  readonly alarm: () => Effect.Effect<void, unknown>;}; const HEARTBEAT_MS = 30 * 1000;const SESSION_TIMEOUT_MS = 5 * 60 * 1000;const LONG_LIVED_DB_IDLE_TIMEOUT_SECONDS = 5;const LONG_LIVED_DB_MAX_LIFETIME_SECONDS = 120;const TRANSPORT_STATE_KEY = "transport";const SESSION_META_KEY = "session-meta";const LAST_ACTIVITY_KEY = "last-activity-ms";const INTERNAL_ACCOUNT_ID_HEADER = "x-executor-mcp-account-id";const INTERNAL_ORGANIZATION_ID_HEADER = "x-executor-mcp-organization-id"; // ---------------------------------------------------------------------------// Errors// --------------------------------------------------------------------------- class OrganizationNotFoundError extends Data.TaggedError("OrganizationNotFoundError")<{  readonly organizationId: string;}> {} class McpServerCloseError extends Data.TaggedError("McpServerCloseError")<{  readonly cause: unknown;}> {} // ---------------------------------------------------------------------------// 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 sessionOwnerMismatch = () =>  jsonRpcError(403, -32003, "MCP session does not belong to the current bearer"); // W3C propagation across the worker -> DO boundary. mcp.ts injects the// worker's `traceparent` and forwards incoming trace headers on forwarded// requests. We convert the parent span into Effect's native ExternalSpan so// the DO's root span stays in the same logical trace without the OpenTelemetry// SDK bridge.const TRACEPARENT_PATTERN = /^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/; const parseTraceparent = (traceparent: string | null | undefined): Tracer.ExternalSpan | null => {  const value = traceparent;  if (!value) return null;  const match = TRACEPARENT_PATTERN.exec(value);  if (!match) return null;  return Tracer.externalSpan({    traceId: match[2]!,    spanId: match[3]!,    sampled: (parseInt(match[4]!, 16) & 1) === 1,  });}; const withIncomingParent = <A, E, R>(  incoming: IncomingTraceHeaders | null | undefined,  effect: Effect.Effect<A, E, R>,): Effect.Effect<A, E, R> => {  const parsed = parseTraceparent(incoming?.traceparent);  return parsed ? Effect.withParentSpan(effect, parsed) : effect;}; type DbHandle = DbServiceShape & { end: () => Promise<void> };type SessionMeta = {  readonly organizationId: string;  readonly organizationName: string;  readonly userId: string;}; /** * Base DB handle factory for MCP session runtimes. * * The DO keeps one postgres.js client for the MCP session runtime. postgres.js * closes idle sockets quickly, while the runtime object stays alive so the MCP * server can preserve session-local protocol state across requests. */const makeDbHandle = (options: {  readonly idleTimeout: number;  readonly maxLifetime: number;  readonly env: Env;}): DbHandle => {  const connectionString = resolveConnectionString(options.env);  const sql = postgres(connectionString, {    max: 1,    idle_timeout: options.idleTimeout,    max_lifetime: options.maxLifetime,    connect_timeout: 10,    fetch_types: false,    prepare: true,    onnotice: () => undefined,  });  return {    sql,    db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb,    // oxlint-disable-next-line executor/no-promise-catch -- boundary: postgres.js close is best-effort during DO/runtime cleanup    end: () => sql.end({ timeout: 0 }).catch(() => undefined),  };}; const makeLongLivedDb = (workerEnv: Env): DbHandle =>  makeDbHandle({    idleTimeout: LONG_LIVED_DB_IDLE_TIMEOUT_SECONDS,    maxLifetime: LONG_LIVED_DB_MAX_LIFETIME_SECONDS,    env: workerEnv,  }); const makeEphemeralDb = (workerEnv: Env): DbHandle =>  makeDbHandle({ idleTimeout: 0, maxLifetime: 60, env: workerEnv }); const makeResolveOrganizationServices = (dbHandle: DbHandle) => {  const DbLive = Layer.succeed(DbService)({ sql: dbHandle.sql, db: dbHandle.db });  const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));  return Layer.mergeAll(DbLive, UserStoreLive, CoreSharedServices);}; // Session services DON'T re-provide `DoTelemetryLive` — that would install a// second WebSdk tracer in the nested Effect scope, disconnecting every// child span from the outer `McpSessionDO.init` / `McpSessionDO.handleRequest`// trace. Tracer comes from the outermost `Effect.provide(DoTelemetryLive)`// at the DO method boundary.const makeSessionServices = (dbHandle: DbHandle) => makeResolveOrganizationServices(dbHandle); const resolveSessionMeta = Effect.fn("McpSessionDO.resolveSessionMeta")(function* (  organizationId: string,  userId: string,) {  const org = yield* resolveOrganization(organizationId);  if (!org) {    return yield* new OrganizationNotFoundError({ organizationId });  }  return {    organizationId: org.id,    organizationName: org.name,    userId,  } satisfies SessionMeta;}); // ---------------------------------------------------------------------------// Durable Object// --------------------------------------------------------------------------- type AlchemyDurableObjectState = Cloudflare.DurableObjectState["Service"]; const makeMcpSession = (state: AlchemyDurableObjectState, workerEnv: Env): McpSessionShape => {  const instanceCreatedAt = Date.now();  let mcpServer: McpServer | null = null;  let transport: McpWorkerTransport | null = null;  let initialized = false;  let lastActivityMs = 0;  let dbHandle: DbHandle | null = null;  let sessionMeta: SessionMeta | null = null;  let transportJsonResponseMode: boolean | null = null;  // Updated at the start of each handleRequest so the host-mcp server's  // parentSpan getter can anchor deferred MCP SDK callbacks to the request span.  let currentRequestSpan: Tracer.AnySpan | null = null;   const provideDoTelemetry = <A, E>(effect: Effect.Effect<A, E, Cloudflare.WorkerEnvironment>) =>    effect.pipe(      Effect.provide(DoTelemetryLive),      Effect.provideService(Cloudflare.WorkerEnvironment, workerEnv),    );   const makeStorage = () => ({    get: (): Effect.Effect<McpTransportState | undefined, unknown> =>      state.storage.get<McpTransportState>(TRANSPORT_STATE_KEY),    set: (value: McpTransportState): Effect.Effect<void, unknown> =>      state.storage.put(TRANSPORT_STATE_KEY, value),  });   const loadSessionMeta = (): Effect.Effect<SessionMeta | null, unknown> =>    Effect.gen(function* () {      if (sessionMeta) return sessionMeta;      const stored = yield* state.storage.get<SessionMeta>(SESSION_META_KEY);      sessionMeta = stored ?? null;      return sessionMeta;    }).pipe(Effect.withSpan("mcp.session.load_meta"));   const saveSessionMeta = (next: SessionMeta): Effect.Effect<void, unknown> =>    Effect.gen(function* () {      sessionMeta = next;      yield* state.storage.put(SESSION_META_KEY, next);    });   const markActivity = (now = Date.now()): Effect.Effect<void, unknown> =>    Effect.gen(function* () {      lastActivityMs = now;      yield* Effect.all(        [state.storage.put(LAST_ACTIVITY_KEY, now), state.storage.setAlarm(now + HEARTBEAT_MS)],        { concurrency: "unbounded" },      );    });   const loadLastActivity = (): Effect.Effect<number, unknown> =>    Effect.gen(function* () {      if (lastActivityMs > 0) return lastActivityMs;      const stored = yield* state.storage.get<number>(LAST_ACTIVITY_KEY);      lastActivityMs = stored ?? 0;      return lastActivityMs;    });   const entryAttrs = (methodEnteredAt: number): Record<string, unknown> => {    const now = Date.now();    return {      "mcp.do.instance_age_ms": now - instanceCreatedAt,      "mcp.do.method_entry_delay_ms": now - methodEnteredAt,      "mcp.session.session_id": state.id.toString(),      "mcp.session.initialized": initialized,      "mcp.session.has_transport": !!transport,      "mcp.session.has_meta_memory": !!sessionMeta,    };  };   const clearSessionState = (): Effect.Effect<void> =>    Effect.gen(function* () {      sessionMeta = null;      initialized = false;      lastActivityMs = 0;      transportJsonResponseMode = null;       yield* Effect.all(        [          state.storage.delete(TRANSPORT_STATE_KEY).pipe(Effect.asVoid),          state.storage.delete(SESSION_META_KEY).pipe(Effect.asVoid),          state.storage.delete(LAST_ACTIVITY_KEY).pipe(Effect.asVoid),          state.storage.deleteAlarm(),        ],        { concurrency: "unbounded" },      );    }).pipe(Effect.withSpan("mcp.session.clear_state"));   const createConnectedRuntime = (    meta: SessionMeta,    options: { readonly dbHandle: DbHandle; readonly enableJsonResponse?: boolean },  ) =>    Effect.gen(function* () {      const { executor, engine } = yield* makeExecutionStack(        meta.userId,        meta.organizationId,        meta.organizationName,      );      const description = yield* buildExecuteDescription(executor);      const server = yield* createExecutorMcpServer({        engine,        description,        parentSpan: () => currentRequestSpan ?? undefined,        debug: workerEnv.EXECUTOR_MCP_DEBUG === "true",      }).pipe(Effect.withSpan("McpSessionDO.createExecutorMcpServer"));      const nextTransport = yield* makeMcpWorkerTransport({        sessionIdGenerator: () => state.id.toString(),        storage: makeStorage(),        enableJsonResponse: options.enableJsonResponse,      });      transportJsonResponseMode = options.enableJsonResponse ?? false;      yield* nextTransport.connect(server);      return { mcpServer: server, transport: nextTransport };    }).pipe(      Effect.withSpan("McpSessionDO.createRuntime"),      Effect.provide(makeSessionServices(options.dbHandle)),    );   const closeRuntime = (): Effect.Effect<void, unknown> =>    Effect.gen(function* () {      if (transport) {        yield* transport.close();        transport = null;      }      if (mcpServer) {        const server = mcpServer;        yield* Effect.ignore(          Effect.tryPromise({            try: () => server.close(),            catch: (cause) => new McpServerCloseError({ cause }),          }),        );        mcpServer = null;      }      if (dbHandle) {        const handle = dbHandle;        yield* Effect.promise(() => handle.end());        dbHandle = null;      }      initialized = false;      transportJsonResponseMode = null;    });   const installRuntime = (    meta: SessionMeta,    options: { readonly dbHandle: DbHandle; readonly enableJsonResponse: boolean },  ) =>    Effect.gen(function* () {      const runtime = yield* createConnectedRuntime(meta, options);      dbHandle = options.dbHandle;      mcpServer = runtime.mcpServer;      transport = runtime.transport;      initialized = true;    });   const restoreRuntimeFromStorage = (    request: Request,  ): Effect.Effect<"restored" | "missing_meta", unknown, Cloudflare.WorkerEnvironment> =>    Effect.gen(function* () {      if (initialized && transport) return "restored" as const;       const meta = yield* loadSessionMeta();      if (!meta) {        yield* Effect.annotateCurrentSpan({ "mcp.session.restore.outcome": "missing_meta" });        return "missing_meta" as const;      }       yield* closeRuntime();      const nextDbHandle = makeLongLivedDb(workerEnv);      yield* installRuntime(meta, { dbHandle: nextDbHandle, enableJsonResponse: true });      yield* markActivity().pipe(Effect.withSpan("McpSessionDO.markActivity"));      yield* Effect.annotateCurrentSpan({ "mcp.session.restore.outcome": "restored" });      return "restored" as const;    }).pipe(      Effect.withSpan("McpSessionDO.restoreRuntime", {        attributes: {          "mcp.request.method": request.method,          "mcp.request.session_id_present": !!request.headers.get("mcp-session-id"),        },      }),    );   const ensureJsonResponseTransportForPost = (    request: Request,  ): Effect.Effect<void, unknown, Cloudflare.WorkerEnvironment> =>    Effect.gen(function* () {      if (request.method !== "POST" || transportJsonResponseMode === true) return;       const meta = yield* loadSessionMeta();      if (!meta) return;       yield* closeRuntime();      const nextDbHandle = makeLongLivedDb(workerEnv);      yield* installRuntime(meta, { dbHandle: nextDbHandle, enableJsonResponse: true });      yield* Effect.annotateCurrentSpan({ "mcp.session.transport_upgraded_json_response": true });    }).pipe(Effect.withSpan("McpSessionDO.ensureJsonResponseTransportForPost"));   const validateSessionOwner = (request: Request): Effect.Effect<Response | null, unknown> =>    Effect.gen(function* () {      const meta = yield* loadSessionMeta();      if (!meta) return null;       const accountId = request.headers.get(INTERNAL_ACCOUNT_ID_HEADER);      const organizationId = request.headers.get(INTERNAL_ORGANIZATION_ID_HEADER);      const matches = accountId === meta.userId && organizationId === meta.organizationId;      yield* Effect.annotateCurrentSpan({ "mcp.session.owner_match": matches });      return matches ? null : sessionOwnerMismatch();    }).pipe(Effect.withSpan("mcp.session.validate_owner"));   const resolveAndStoreSessionMeta = Effect.fn("mcp.session.resolve_and_store_meta")(function* (    token: McpSessionInit,  ) {    const handle = makeEphemeralDb(workerEnv);    return yield* resolveSessionMeta(token.organizationId, token.userId).pipe(      Effect.provide(makeResolveOrganizationServices(handle)),      Effect.tap((meta) => saveSessionMeta(meta).pipe(Effect.withSpan("mcp.session.save_meta"))),      Effect.ensuring(Effect.promise(() => handle.end())),    );  });   const cleanup = (): Effect.Effect<void, unknown> =>    Effect.gen(function* () {      yield* closeRuntime();      yield* clearSessionState();    });   const doInit = (token: McpSessionInit) =>    Effect.gen(function* () {      const meta = yield* resolveAndStoreSessionMeta(token);      const handle = makeLongLivedDb(workerEnv);      const runtime = yield* createConnectedRuntime(meta, {        dbHandle: handle,        enableJsonResponse: true,      });      dbHandle = handle;      mcpServer = runtime.mcpServer;      transport = runtime.transport;      initialized = true;      yield* markActivity().pipe(Effect.withSpan("McpSessionDO.markActivity"));    }).pipe(      Effect.tapCause((cause) =>        Effect.sync(() => {          console.error("[mcp-session] init failed:", cause);        }),      ),      Effect.catchCause((cause) =>        Effect.gen(function* () {          yield* cleanup();          return yield* Effect.failCause(cause);        }),      ),    );   const init = (token: McpSessionInit, incoming?: IncomingTraceHeaders) => {    const methodEnteredAt = Date.now();    if (initialized) return Effect.void;    return Effect.gen(function* () {      yield* Effect.annotateCurrentSpan(entryAttrs(methodEnteredAt));      yield* doInit(token);    }).pipe(      Effect.withSpan("McpSessionDO.init", {        attributes: { "mcp.auth.organization_id": token.organizationId },      }),      (eff) => withIncomingParent(incoming, eff),      provideDoTelemetry,    );  };   const dispatchAuthorizedRequest = (    request: Request,  ): Effect.Effect<Response, unknown, Cloudflare.WorkerEnvironment> => {    if (!initialized || !transport) {      if (request.method === "DELETE") {        return clearSessionState().pipe(          Effect.as(new Response(null, { status: 204 })),          Effect.withSpan("mcp.session.stale_delete"),        );      }      return Effect.gen(function* () {        const restored = yield* restoreRuntimeFromStorage(request);        if (restored === "restored") return yield* dispatchAuthorizedRequest(request);        return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");      });    }     return Effect.gen(function* () {      yield* ensureJsonResponseTransportForPost(request);      const activeTransport = transport;      if (!activeTransport) {        return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");      }       yield* markActivity().pipe(Effect.withSpan("McpSessionDO.markActivity"));      const response = yield* activeTransport.handleRequest(request).pipe(        Effect.withSpan("McpSessionDO.transport.handleRequest", {          attributes: {            "mcp.request.method": request.method,            "mcp.request.content_type": request.headers.get("content-type") ?? "",            "mcp.request.content_length": request.headers.get("content-length") ?? "",          },

Lines 1–500 of 606. Load more