History

Merge pull request #281 from RhysSullivan/rs/mcp-do-shared-layer

Rhys Sullivan committed 7 months ago

Commit b5741222 parents 599ae33 8177e398 files · +176 −95Browse files

Unified · Split

8 files
  1. apps/cloud/src/api/error-response.ts +6 −1
  2. apps/cloud/src/api/layers.ts +12 −2
  3. apps/cloud/src/api/protected.ts +2 −15
  4. apps/cloud/src/mcp-session.ts +37 −42
  5. apps/cloud/src/mcp.ts +2 −0
  6. apps/cloud/src/server.ts +52 −29
  7. apps/cloud/src/services/execution-stack.ts +29 −0
  8. apps/cloud/src/services/telemetry.ts +36 −6
apps/cloud/src/api/error-response.ts+6−1
@@ -1,3 +1,4 @@
1+import * as Sentry from "@sentry/cloudflare";
12 import { Data } from "effect";
23 import { HttpServerResponse } from "@effect/platform";
34
@@ -20,12 +21,16 @@ export const isServerError = (error: unknown): boolean => toHttpResponseError(er · 16 unmodified lines
2021
2122 export const toErrorResponse = (error: unknown): Response => {
2223 const mapped = toHttpResponseError(error);
24+ if (mapped.status >= 500) Sentry.captureException(error);
2325 return Response.json({ error: mapped.message, code: mapped.code }, { status: mapped.status });
2426 };
2527
2628 export const toErrorServerResponse = (error: unknown): HttpServerResponse.HttpServerResponse => {
27− console.error("[api] toErrorServerResponse error:", error instanceof Error ? error.stack : error);
2829 const mapped = toHttpResponseError(error);
30+ if (mapped.status >= 500) {
31+ console.error("[api] toErrorServerResponse error:", error instanceof Error ? error.stack : error);
32+ Sentry.captureException(error);
33+ }
2934 return HttpServerResponse.unsafeJson(
3035 { error: mapped.message, code: mapped.code },
3136 { status: mapped.status },
apps/cloud/src/api/layers.ts+12−2
@@ -30,11 +30,21 @@ const ProtectedCloudApi = CoreExecutorApi.add(OpenApiGroup) · 29 unmodified lines
3030 const DbLive = DbService.Live;
3131 const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));
3232
33+/**
34+ * Services that are independent of how the DB or tracer is provisioned —
35+ * both the stateless HTTP path (per-request DB via Hyperdrive) and the MCP
36+ * session DO (long-lived DB + isolate-local tracer SDK) merge this with
37+ * their own `DbLive` + `UserStoreLive` + telemetry layer.
38+ */
39+export const CoreSharedServices = Layer.mergeAll(
40+ WorkOSAuth.Default,
41+ AutumnService.Default,
42+);
43+
3344 export const SharedServices = Layer.mergeAll(
3445 DbLive,
3546 UserStoreLive,
36− WorkOSAuth.Default,
37− AutumnService.Default,
47+ CoreSharedServices,
3848 HttpServer.layerContext,
3949 TelemetryLive,
4050 );
apps/cloud/src/api/protected.ts+2−15
@@ -1,21 +1,15 @@
1−import { env } from "cloudflare:workers";
21 import { HttpApiBuilder, HttpApiSwagger, HttpServerRequest } from "@effect/platform";
32 import { Effect, Layer } from "effect";
43
54 import { ExecutorService, ExecutionEngineService } from "@executor/api/server";
6−import { createExecutionEngine } from "@executor/execution";
7−import { makeDynamicWorkerExecutor } from "@executor/runtime-dynamic-worker";
85 import { OpenApiExtensionService } from "@executor/plugin-openapi/api";
96 import { McpExtensionService } from "@executor/plugin-mcp/api";
107 import { GraphqlExtensionService } from "@executor/plugin-graphql/api";
118
129 import { authorizeOrganization } from "../auth/authorize-organization";
1310 import { WorkOSAuth } from "../auth/workos";
14−import { AutumnService } from "../services/autumn";
15−import { createScopedExecutor } from "../services/executor";
16−import { makeTrackExecutionUsage } from "./autumn";
11+import { makeExecutionStack } from "../services/execution-stack";
1712 import { HttpResponseError, isServerError, toErrorServerResponse } from "./error-response";
18−import { withExecutionUsageTracking } from "./execution-usage";
1913 import { ProtectedCloudApiLive, RouterConfig, SharedServices } from "./layers";
2014
2115 const lookupOrgForRequest = (request: HttpServerRequest.HttpServerRequest) =>
@@ -38,14 +32,7 @@ const lookupOrgForRequest = (request: HttpServerRequest.HttpServerRequest) => · 16 unmodified lines
3832
3933 const createProtectedApp = (organizationId: string, organizationName: string) =>
4034 Effect.gen(function* () {
41− const executor = yield* createScopedExecutor(organizationId, organizationName);
42− const codeExecutor = makeDynamicWorkerExecutor({ loader: env.LOADER });
43− const autumn = yield* AutumnService;
44− const engine = withExecutionUsageTracking(
45− organizationId,
46− createExecutionEngine({ executor, codeExecutor }),
47− makeTrackExecutionUsage(autumn),
48− );
35+ const { executor, engine } = yield* makeExecutionStack(organizationId, organizationName);
4936
5037 const requestServices = Layer.mergeAll(
5138 Layer.succeed(ExecutorService, executor),
apps/cloud/src/mcp-session.ts+37−42
@@ -10,19 +10,15 @@ import { drizzle } from "drizzle-orm/postgres-js"; · 9 unmodified lines
1010 import postgres from "postgres";
1111
1212 import { createExecutorMcpServer } from "@executor/host-mcp";
13−import { createExecutionEngine } from "@executor/execution";
14−import { makeDynamicWorkerExecutor } from "@executor/runtime-dynamic-worker";
1513 import type { DrizzleDb, DbServiceShape } from "./services/db";
1614
17−import { makeTrackExecutionUsage } from "./api/autumn";
18−import { withExecutionUsageTracking } from "./api/execution-usage";
15+import { CoreSharedServices } from "./api/layers";
1916 import { UserStoreService } from "./auth/context";
2017 import { resolveOrganization } from "./auth/resolve-organization";
21−import { WorkOSAuth } from "./auth/workos";
2218 import { server } from "./env";
23−import { AutumnService } from "./services/autumn";
24−import { createScopedExecutor } from "./services/executor";
2519 import { DbService, combinedSchema } from "./services/db";
20+import { makeExecutionStack } from "./services/execution-stack";
21+import { DoTelemetryLive } from "./services/telemetry";
2622
2723 // ---------------------------------------------------------------------------
2824 // Types
@@ -105,47 +101,46 @@ export class McpSessionDO extends DurableObject { · 76 unmodified lines
105101 async init(token: McpSessionInit): Promise<void> {
106102 if (this.initialized) return;
107103
108− // Create a long-lived DB connection for the DO's lifetime
109104 this.dbHandle = makeLongLivedDb();
110−
111− const { sql, db } = this.dbHandle;
112− const DbLive = Layer.succeed(DbService, { sql, db });
113− const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));
114− const Services = Layer.mergeAll(
115− DbLive,
116− UserStoreLive,
117− WorkOSAuth.Default,
118− AutumnService.Default,
119− );
120−
121− const program = Effect.gen(function* () {
122− const org = yield* resolveOrganization(token.organizationId);
123− if (!org)
124− return yield* new OrganizationNotFoundError({ organizationId: token.organizationId });
125−
126− const executor = yield* createScopedExecutor(org.id, org.name);
127− const codeExecutor = makeDynamicWorkerExecutor({ loader: env.LOADER });
128− const autumn = yield* AutumnService;
129− const engine = withExecutionUsageTracking(
130− org.id,
131− createExecutionEngine({ executor, codeExecutor }),
132− makeTrackExecutionUsage(autumn),
105+ try {
106+ const { sql, db } = this.dbHandle;
107+ const DbLive = Layer.succeed(DbService, { sql, db });
108+ const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));
109+ const Services = Layer.mergeAll(
110+ DbLive,
111+ UserStoreLive,
112+ CoreSharedServices,
113+ DoTelemetryLive,
133114 );
134− return yield* Effect.promise(() => createExecutorMcpServer({ engine }));
135− }).pipe(Effect.provide(Services));
136115
137− this.mcpServer = await Effect.runPromise(program);
116+ const program = Effect.gen(function* () {
117+ const org = yield* resolveOrganization(token.organizationId);
118+ if (!org)
119+ return yield* new OrganizationNotFoundError({ organizationId: token.organizationId });
138120
139− this.transport = new WorkerTransport({
140− sessionIdGenerator: () => this.ctx.id.toString(),
141− storage: this.makeStorage(),
142− });
121+ const { engine } = yield* makeExecutionStack(org.id, org.name);
122+ return yield* Effect.promise(() => createExecutorMcpServer({ engine }));
123+ }).pipe(Effect.withSpan("McpSessionDO.init"), Effect.provide(Services));
143124
144− await this.mcpServer.connect(this.transport);
145− this.initialized = true;
146− this.lastActivityMs = Date.now();
125+ this.mcpServer = await Effect.runPromise(program);
147126
148− await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS);
127+ this.transport = new WorkerTransport({
128+ sessionIdGenerator: () => this.ctx.id.toString(),
129+ storage: this.makeStorage(),
130+ });
131+
132+ await this.mcpServer.connect(this.transport);
133+ this.initialized = true;
134+ this.lastActivityMs = Date.now();
135+
136+ await this.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS);
137+ } catch (err) {
138+ // Partial init leaves dangling resources (DB socket, maybe mcpServer).
139+ // Clean up before rethrowing so the DO isn't stuck in a half-built state.
140+ console.error("[mcp-session] init failed:", err instanceof Error ? err.stack : err);
141+ await this.cleanup();
142+ throw err;
143+ }
149144 }
150145
151146 async handleRequest(request: Request): Promise<Response> {
apps/cloud/src/mcp.ts+2−0
@@ -3,6 +3,7 @@ · 2 unmodified lines
33 // ---------------------------------------------------------------------------
44
55 import { env } from "cloudflare:workers";
6+import * as Sentry from "@sentry/cloudflare";
67 import { createRemoteJWKSet, jwtVerify } from "jose";
78
89 import { server } from "./env";
@@ -130,6 +131,7 @@ const handleMcpRequest_POST = async (request: Request, token: VerifiedToken): Pr · 121 unmodified lines
130131 return await stub.handleRequest(request);
131132 } catch (err) {
132133 console.error("[mcp] POST handler error:", err instanceof Error ? err.stack : err);
134+ Sentry.captureException(err);
133135 return jsonRpcError(500, -32603, "Internal server error");
134136 }
135137 };
apps/cloud/src/server.ts+52−29
@@ -1,45 +1,68 @@
11 import * as Sentry from "@sentry/cloudflare";
22 import handler from "@tanstack/react-start/server-entry";
3−import { instrument, type ResolveConfigFn } from "@microlabs/otel-cf-workers";
3+import { instrument, type TraceConfig } from "@microlabs/otel-cf-workers";
44
5−// Export Durable Objects as named exports
6−export { McpSessionDO } from "./mcp-session";
5+import { McpSessionDO as McpSessionDOBase } from "./mcp-session";
6+import { server } from "./env";
77
88 // ---------------------------------------------------------------------------
9−// OTEL config — `otel-cf-workers` owns the global TracerProvider and flushes
10−// via `ctx.waitUntil` at the end of each request. `TelemetryLive` in
11−// `services/telemetry.ts` plugs Effect's tracer into that same provider.
9+// OTEL config for the main fetch handler — `otel-cf-workers` owns the global
10+// TracerProvider and flushes via `ctx.waitUntil` at the end of each request.
11+// The DO runs in a separate isolate and uses its own self-contained WebSdk
12+// (see `services/telemetry.ts#DoTelemetryLive`); `instrumentDO` from
13+// otel-cf-workers is NOT used because it breaks `this` binding on
14+// `WorkerTransport`'s stream primitives and crashes every MCP request with
15+// DOMException "Illegal invocation".
1216 // ---------------------------------------------------------------------------
1317
14−type OtelEnv = {
15− AXIOM_TOKEN?: string;
16− AXIOM_DATASET?: string;
17−};
18−
19−const resolveOtelConfig: ResolveConfigFn<OtelEnv> = (env) => ({
18+const otelConfig: TraceConfig = {
2019 service: { name: "executor-cloud", version: "1.0.0" },
2120 exporter: {
2221 url: "https://api.axiom.co/v1/traces",
2322 headers: {
24− Authorization: `Bearer ${env.AXIOM_TOKEN ?? ""}`,
25− "X-Axiom-Dataset": env.AXIOM_DATASET ?? "executor-cloud",
23+ Authorization: `Bearer ${server.AXIOM_TOKEN}`,
24+ "X-Axiom-Dataset": server.AXIOM_DATASET,
2625 },
2726 },
27+};
28+
29+// otel-cf-workers owns the global TracerProvider. Sentry's OTEL compat shim
30+// registers a ProxyTracerProvider of its own, which prevents otel-cf-workers
31+// from finding its WorkerTracer and breaks the whole request path with
32+// "global tracer is not of type WorkerTracer".
33+//
34+// The `_env` parameter is unused — `server` from `./env` already gives us
35+// typed access to every secret. It's only in the signature so Sentry's
36+// generics infer the DO's `Env` type correctly.
37+const sentryOptions = (_env: Env) => ({
38+ dsn: server.SENTRY_DSN,
39+ tracesSampleRate: 0,
40+ enableLogs: true,
41+ sendDefaultPii: true,
42+ skipOpenTelemetrySetup: true,
43+ // Our DO methods (init/handleRequest/alarm) live on the prototype, not on
44+ // the instance. Sentry's default DO auto-wrap only visits own properties,
45+ // which misses prototype methods — so errors thrown inside init() never
46+ // reach Sentry. This flag opts into prototype-method instrumentation.
47+ instrumentPrototypeMethods: true,
2848 });
2949
30−const instrumentedHandler = instrument({ fetch: handler.fetch }, resolveOtelConfig);
31−
32−export default Sentry.withSentry(
33− (env: Record<string, string>) => ({
34− dsn: env.SENTRY_DSN,
35− tracesSampleRate: 0,
36− enableLogs: true,
37− sendDefaultPii: true,
38− // otel-cf-workers owns the global TracerProvider. Sentry's OTEL compat
39− // shim registers a ProxyTracerProvider of its own, which prevents
40− // otel-cf-workers from finding its WorkerTracer and breaks the whole
41− // request path with "global tracer is not of type WorkerTracer".
42− skipOpenTelemetrySetup: true,
43− }),
44− instrumentedHandler,
50+// ---------------------------------------------------------------------------
51+// Durable Object — wrapped with Sentry so DO errors land in Sentry (inits the
52+// client inside the DO isolate, which plain `Sentry.captureException` cannot
53+// do on its own). We deliberately do NOT wrap with otel-cf-workers'
54+// `instrumentDO` (see note above).
55+// ---------------------------------------------------------------------------
56+
57+export const McpSessionDO = Sentry.instrumentDurableObjectWithSentry(
58+ sentryOptions,
59+ McpSessionDOBase,
4560 );
61+
62+// ---------------------------------------------------------------------------
63+// Worker fetch handler
64+// ---------------------------------------------------------------------------
65+
66+const instrumentedHandler = instrument({ fetch: handler.fetch }, otelConfig);
67+
68+export default Sentry.withSentry(sentryOptions, instrumentedHandler);
apps/cloud/src/services/execution-stack.ts+29−0
@@ -0,0 +1,29 @@
1+// ---------------------------------------------------------------------------
2+// Shared execution stack — the wiring that turns an organization into a
3+// runnable executor + engine. Used by the protected HTTP API (per-request)
4+// and the MCP session DO (per-session) so changes to the stack flow to both.
5+// ---------------------------------------------------------------------------
6+
7+import { env } from "cloudflare:workers";
8+import { Effect } from "effect";
9+
10+import { createExecutionEngine } from "@executor/execution";
11+import { makeDynamicWorkerExecutor } from "@executor/runtime-dynamic-worker";
12+
13+import { makeTrackExecutionUsage } from "../api/autumn";
14+import { withExecutionUsageTracking } from "../api/execution-usage";
15+import { AutumnService } from "./autumn";
16+import { createScopedExecutor } from "./executor";
17+
18+export const makeExecutionStack = (organizationId: string, organizationName: string) =>
19+ Effect.gen(function* () {
20+ const executor = yield* createScopedExecutor(organizationId, organizationName);
21+ const codeExecutor = makeDynamicWorkerExecutor({ loader: env.LOADER });
22+ const autumn = yield* AutumnService;
23+ const engine = withExecutionUsageTracking(
24+ organizationId,
25+ createExecutionEngine({ executor, codeExecutor }),
26+ makeTrackExecutionUsage(autumn),
27+ );
28+ return { executor, engine };
29+ });
apps/cloud/src/services/telemetry.ts+36−6
@@ -1,15 +1,45 @@
11 // ---------------------------------------------------------------------------
2−// Effect → global OTEL bridge
2+// Effect → OTEL → Axiom bridge
33 // ---------------------------------------------------------------------------
4−// The global TracerProvider is owned by `@microlabs/otel-cf-workers` (see
5−// `server.ts`). This layer plugs Effect's tracer into that provider so every
6−// `Effect.withSpan(...)` becomes a real OTLP span exported to Axiom, with
7−// flushing handled reliably by the instrument() wrapper via ctx.waitUntil.
4+//
5+// Two callers, two setups:
6+//
7+// - `TelemetryLive` (fetch path): reads the global `TracerProvider` that
8+// `@microlabs/otel-cf-workers`' `instrument(...)` installs in `server.ts`.
9+// Flushing is handled by `instrument()` via `ctx.waitUntil` at request end.
10+//
11+// - `DoTelemetryLive` (Durable Object path): provisions its own
12+// `WebSdk`-backed tracer via `Effect`. The DO runs in a separate isolate
13+// and we deliberately avoid `instrumentDO` (it wraps DO methods in a way
14+// that breaks `this` binding on `WorkerTransport`'s stream primitives —
15+// every MCP request 500s with "Illegal invocation"). The DO uses a
16+// `SimpleSpanProcessor` so spans export immediately; there's no
17+// `ctx.waitUntil` to rely on for batching.
818 // ---------------------------------------------------------------------------
919
10−import { Resource, Tracer as OtelTracer } from "@effect/opentelemetry";
20+import { Resource, Tracer as OtelTracer, WebSdk } from "@effect/opentelemetry";
21+import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
22+import { SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base";
1123 import { Layer } from "effect";
1224
25+import { server } from "../env";
26+
1327 export const TelemetryLive: Layer.Layer<never> = OtelTracer.layerGlobal.pipe(
1428 Layer.provide(Resource.layer({ serviceName: "executor-cloud", serviceVersion: "1.0.0" })),
1529 );
30+
31+const makeDoOtelExporter = () =>
32+ new OTLPTraceExporter({
33+ url: "https://api.axiom.co/v1/traces",
34+ headers: {
35+ Authorization: `Bearer ${server.AXIOM_TOKEN}`,
36+ "X-Axiom-Dataset": server.AXIOM_DATASET,
37+ },
38+ });
39+
40+export const DoTelemetryLive: Layer.Layer<never> = server.AXIOM_TOKEN
41+ ? WebSdk.layer(() => ({
42+ resource: { serviceName: "executor-cloud", serviceVersion: "1.0.0" },
43+ spanProcessor: new SimpleSpanProcessor(makeDoOtelExporter()),
44+ }))
45+ : Layer.empty;