// ---------------------------------------------------------------------------// MCP Session Durable Object — holds MCP server + engine per session// --------------------------------------------------------------------------- import { DurableObject, env } from "cloudflare:workers";import { Data, Effect, Layer } from "effect";import * as OtelTracer from "@effect/opentelemetry/Tracer";import * as Sentry from "@sentry/cloudflare";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 { buildExecuteDescription } from "@executor/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 { server } from "./env";import { DbService, combinedSchema } from "./services/db";import { makeExecutionStack } from "./services/execution-stack";import { DoTelemetryLive } from "./services/telemetry"; // ---------------------------------------------------------------------------// Types// --------------------------------------------------------------------------- export type McpSessionInit = {  organizationId: string;}; const HEARTBEAT_MS = 30 * 1000;const SESSION_TIMEOUT_MS = 5 * 60 * 1000;const TRANSPORT_STATE_KEY = "transport";const SESSION_META_KEY = "session-meta"; // ---------------------------------------------------------------------------// 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" },  }); // W3C traceparent propagation across the worker→DO boundary. mcp.ts injects// the worker-side `mcp.request` SpanContext as a `traceparent` header on// forwarded requests (and as a second arg to `init()`). We parse it here and// use `OtelTracer.withSpanContext` to stitch the DO's root span under the// worker span so the entire logical request lives in one Axiom trace.const TRACEPARENT_PATTERN = /^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/; type IncomingSpanContext = {  readonly traceId: string;  readonly spanId: string;  readonly traceFlags: number;}; const parseTraceparent = (value: string | null | undefined): IncomingSpanContext | null => {  if (!value) return null;  const match = TRACEPARENT_PATTERN.exec(value);  if (!match) return null;  return { traceId: match[2]!, spanId: match[3]!, traceFlags: parseInt(match[4]!, 16) };}; const withIncomingParent = <A, E, R>(  traceparent: string | null | undefined,  effect: Effect.Effect<A, E, R>,): Effect.Effect<A, E, R> => {  const parsed = parseTraceparent(traceparent);  return parsed ? OtelTracer.withSpanContext(effect, parsed) : effect;}; type DbHandle = DbServiceShape & { end: () => Promise<void> };type SessionMeta = {  readonly organizationId: string;  readonly organizationName: string;}; /** * Base DB handle factory for MCP session runtimes. * * The production DO keeps one postgres.js socket for the session lifetime. * The workerd test pool rejects that socket on the next request because the * underlying I/O object is request-bound. Tests therefore opt into a * request-scoped runtime and rebuild the server + DB handle per POST/DELETE * while keeping only the MCP transport state in DO storage. */const makeDbHandle = (options: {  readonly idleTimeout: number;  readonly maxLifetime: number;}): DbHandle => {  const connectionString = env.HYPERDRIVE?.connectionString ?? server.DATABASE_URL;  const sql = postgres(connectionString, {    max: 1,    idle_timeout: options.idleTimeout,    max_lifetime: options.maxLifetime,    connect_timeout: 10,    onnotice: () => undefined,  });  return {    sql,    db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb,    end: () => sql.end({ timeout: 0 }).catch(() => undefined),  };}; const makeLongLivedDb = (): DbHandle => makeDbHandle({ idleTimeout: 20, maxLifetime: 300 }); const makeRequestScopedDb = (): DbHandle => makeDbHandle({ idleTimeout: 0, maxLifetime: 60 }); 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,) {    const org = yield* resolveOrganization(organizationId);    if (!org) {      return yield* new OrganizationNotFoundError({ organizationId });    }    return {      organizationId: org.id,      organizationName: org.name,    } satisfies SessionMeta;  }); const requestScopedRuntimeEnabled = server.MCP_SESSION_REQUEST_SCOPED_RUNTIME === "true"; // ---------------------------------------------------------------------------// 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: DbHandle | null = null;  private sessionMeta: SessionMeta | null = null;   private makeStorage() {    return {      get: async (): Promise<TransportState | undefined> => {        return await this.ctx.storage.get<TransportState>(TRANSPORT_STATE_KEY);      },      set: async (state: TransportState): Promise<void> => {        await this.ctx.storage.put(TRANSPORT_STATE_KEY, state);      },    };  }   private async loadSessionMeta(): Promise<SessionMeta | null> {    if (this.sessionMeta) return this.sessionMeta;    const stored = await this.ctx.storage.get<SessionMeta>(SESSION_META_KEY);    this.sessionMeta = stored ?? null;    return this.sessionMeta;  }   private async saveSessionMeta(sessionMeta: SessionMeta): Promise<void> {    this.sessionMeta = sessionMeta;    await this.ctx.storage.put(SESSION_META_KEY, sessionMeta);  }   private async clearSessionState(): Promise<void> {    this.sessionMeta = null;    this.initialized = false;    this.lastActivityMs = 0;     await Promise.all([      this.ctx.storage.delete(TRANSPORT_STATE_KEY).catch(() => false),      this.ctx.storage.delete(SESSION_META_KEY).catch(() => false),    ]);  }   private createConnectedRuntimeEffect(    sessionMeta: SessionMeta,    options: { readonly dbHandle: DbHandle; readonly enableJsonResponse?: boolean },  ) {    const self = this;    return Effect.gen(function* () {      const { executor, engine } = yield* makeExecutionStack(        sessionMeta.organizationId,        sessionMeta.organizationName,      );      // Build the description here so the two postgres queries it runs      // (`executor.sources.list` + `executor.tools.list`) land as      // children of `McpSessionDO.createRuntime`. host-mcp would      // otherwise call `Effect.runPromise(engine.getDescription)` at      // its async MCP-SDK boundary and orphan those sub-spans.      const description = yield* buildExecuteDescription(executor);      const mcpServer = yield* Effect.promise(() =>        createExecutorMcpServer({ engine, description }),      ).pipe(Effect.withSpan("McpSessionDO.createExecutorMcpServer"));      const transport = new WorkerTransport({        sessionIdGenerator: () => self.ctx.id.toString(),        storage: self.makeStorage(),        enableJsonResponse: options.enableJsonResponse,      });      yield* Effect.promise(() => mcpServer.connect(transport)).pipe(        Effect.withSpan("McpSessionDO.transport.connect"),      );      return { mcpServer, transport };    }).pipe(      Effect.withSpan("McpSessionDO.createRuntime"),      Effect.provide(makeSessionServices(options.dbHandle)),    );  }   private resolveAndStoreSessionMetaEffect(token: McpSessionInit) {    const self = this;    return Effect.gen(function* () {      const dbHandle = makeRequestScopedDb();      try {        const sessionMeta = yield* resolveSessionMeta(token.organizationId).pipe(          Effect.provide(makeResolveOrganizationServices(dbHandle)),        );        yield* Effect.promise(() => self.saveSessionMeta(sessionMeta));        return sessionMeta;      } finally {        yield* Effect.promise(() => dbHandle.end());      }    });  }   async init(token: McpSessionInit, traceparent?: string): Promise<void> {    if (this.initialized) return;    return Effect.runPromise(      this.doInitEffect(token).pipe(        Effect.withSpan("McpSessionDO.init", {          attributes: { "mcp.auth.organization_id": token.organizationId },        }),        (eff) => withIncomingParent(traceparent, eff),        Effect.provide(DoTelemetryLive),      ),    );  }   private doInitEffect(token: McpSessionInit) {    const self = this;    // Single Effect chain so every sub-span (resolveSessionMeta,    // createRuntime, createScopedExecutor, createExecutorMcpServer,    // transport.connect, storage.setAlarm) lands as a child of    // `McpSessionDO.init`. The prior implementation called    // `Effect.runPromise` nested inside an async function, which orphaned    // each sub-span into its own root trace and made init opaque —    // dashboard saw one 2.77s span with nothing under it.    return Effect.gen(function* () {      const sessionMeta = yield* self.resolveAndStoreSessionMetaEffect(token);       if (!requestScopedRuntimeEnabled) {        self.dbHandle = makeLongLivedDb();        const runtime = yield* self.createConnectedRuntimeEffect(sessionMeta, {          dbHandle: self.dbHandle,        });        self.mcpServer = runtime.mcpServer;        self.transport = runtime.transport;      }       self.initialized = true;      self.lastActivityMs = Date.now();       yield* Effect.promise(() => self.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS)).pipe(        Effect.withSpan("McpSessionDO.setAlarm"),      );    }).pipe(      Effect.tapErrorCause((cause) =>        Effect.sync(() => {          console.error("[mcp-session] init failed:", cause);        }),      ),      Effect.catchAllCause((cause) =>        Effect.gen(function* () {          yield* Effect.promise(() => self.cleanup());          return yield* Effect.failCause(cause);        }),      ),      Effect.orDie,    );  }   private async handleRequestWithRequestScopedRuntime(request: Request): Promise<Response> {    const sessionMeta = await this.loadSessionMeta();    if (!sessionMeta) {      return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");    }     this.initialized = true;    this.lastActivityMs = Date.now();     let dbHandle: DbHandle | null = null;    let mcpServer: McpServer | null = null;    let transport: WorkerTransport | null = null;     try {      dbHandle = makeRequestScopedDb();      const runtime = await Effect.runPromise(        this.createConnectedRuntimeEffect(sessionMeta, {          dbHandle,          enableJsonResponse: request.method !== "GET",        }).pipe(Effect.provide(DoTelemetryLive)),      );      mcpServer = runtime.mcpServer;      transport = runtime.transport;       const response = await transport.handleRequest(request);      if (request.method === "DELETE") {        await this.clearSessionState();      }      return response;    } catch (err) {      console.error(        "[mcp-session] request-scoped handleRequest error:",        err instanceof Error ? err.stack : err,      );      Sentry.captureException(err);      return jsonRpcError(500, -32603, "Internal error");    } finally {      await transport?.close().catch(() => undefined);      await mcpServer?.close().catch(() => undefined);      await dbHandle?.end();    }  }   async handleRequest(request: Request): Promise<Response> {    // Wrap the dispatch in an Effect span so every DO request — not just    // the rare new-session `init()` — shows up in Axiom. Basic attributes    // only (method, session-id presence, response status); rich client    // fingerprint stays on the edge `mcp.request` span, which shares a    // trace_id with this one.    const traceparent = request.headers.get("traceparent");    const program = Effect.promise(() => this.dispatchRequest(request)).pipe(      Effect.tap((response) =>        Effect.annotateCurrentSpan({          "mcp.response.status_code": response.status,        }),      ),      Effect.withSpan("McpSessionDO.handleRequest", {        attributes: {          "mcp.request.method": request.method,          "mcp.request.session_id_present": !!request.headers.get("mcp-session-id"),        },      }),      (eff) => withIncomingParent(traceparent, eff),      Effect.provide(DoTelemetryLive),    );    return Effect.runPromise(program);  }   private async dispatchRequest(request: Request): Promise<Response> {    if (requestScopedRuntimeEnabled) {      return this.handleRequestWithRequestScopedRuntime(request);    }     if (!this.initialized || !this.transport) {      return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");    }     this.lastActivityMs = Date.now();     try {      const response = await this.transport.handleRequest(request);      if (request.method === "DELETE") {        await this.cleanup();      }      return response;    } catch (err) {      console.error("[mcp-session] handleRequest error:", err instanceof Error ? err.stack : err);      Sentry.captureException(err);      return jsonRpcError(500, -32603, "Internal error");    }  }   async alarm(): Promise<void> {    const program = Effect.promise(() => this.runAlarm()).pipe(      Effect.withSpan("McpSessionDO.alarm"),      Effect.provide(DoTelemetryLive),    );    return Effect.runPromise(program);  }   private async runAlarm(): 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;    }    await this.clearSessionState();  }}