Unified · Split
| @@ -182,6 +182,7 @@ const resolveSessionMeta = Effect.fn("McpSessionDO.resolveSessionMeta")(function · 181 unmodified lines | |||
| 182 | // --------------------------------------------------------------------------- | 182 | // --------------------------------------------------------------------------- |
| 183 | 183 | ||
| 184 | export class McpSessionDO extends DurableObject { | 184 | export class McpSessionDO extends DurableObject { |
| 185 | private readonly instanceCreatedAt = Date.now(); | ||
| 185 | private mcpServer: McpServer | null = null; | 186 | private mcpServer: McpServer | null = null; |
| 186 | private transport: WorkerTransport | null = null; | 187 | private transport: WorkerTransport | null = null; |
| 187 | private initialized = false; | 188 | private initialized = false; |
| @@ -223,6 +224,17 @@ export class McpSessionDO extends DurableObject { · 35 unmodified lines | |||
| 223 | await this.ctx.storage.put(SESSION_META_KEY, sessionMeta); | 224 | await this.ctx.storage.put(SESSION_META_KEY, sessionMeta); |
| 224 | } | 225 | } |
| 225 | 226 | ||
| 227 | private entryAttrs(methodEnteredAt: number): Record<string, unknown> { | ||
| 228 | const now = Date.now(); | ||
| 229 | return { | ||
| 230 | "mcp.do.instance_age_ms": now - this.instanceCreatedAt, | ||
| 231 | "mcp.do.method_entry_delay_ms": now - methodEnteredAt, | ||
| 232 | "mcp.session.initialized": this.initialized, | ||
| 233 | "mcp.session.has_transport": !!this.transport, | ||
| 234 | "mcp.session.has_meta_memory": !!this.sessionMeta, | ||
| 235 | }; | ||
| 236 | } | ||
| 237 | |||
| 226 | private clearSessionState(): Effect.Effect<void> { | 238 | private clearSessionState(): Effect.Effect<void> { |
| 227 | return Effect.promise(async () => { | 239 | return Effect.promise(async () => { |
| 228 | this.sessionMeta = null; | 240 | this.sessionMeta = null; |
| @@ -274,6 +286,68 @@ export class McpSessionDO extends DurableObject { · 45 unmodified lines | |||
| 274 | ); | 286 | ); |
| 275 | } | 287 | } |
| 276 | 288 | ||
| 289 | private closeRuntime(): Effect.Effect<void> { | ||
| 290 | const self = this; | ||
| 291 | return Effect.promise(async () => { | ||
| 292 | if (self.transport) { | ||
| 293 | await self.transport.close().catch(() => undefined); | ||
| 294 | self.transport = null; | ||
| 295 | } | ||
| 296 | if (self.mcpServer) { | ||
| 297 | await self.mcpServer.close().catch(() => undefined); | ||
| 298 | self.mcpServer = null; | ||
| 299 | } | ||
| 300 | if (self.dbHandle) { | ||
| 301 | await self.dbHandle.end(); | ||
| 302 | self.dbHandle = null; | ||
| 303 | } | ||
| 304 | self.initialized = false; | ||
| 305 | }).pipe(Effect.orDie); | ||
| 306 | } | ||
| 307 | |||
| 308 | private restoreRuntimeFromStorage( | ||
| 309 | request: Request, | ||
| 310 | ): Effect.Effect<"restored" | "missing_meta"> { | ||
| 311 | const self = this; | ||
| 312 | return Effect.gen(function* () { | ||
| 313 | if (self.initialized && self.transport) return "restored" as const; | ||
| 314 | |||
| 315 | const sessionMeta = yield* self.loadSessionMeta(); | ||
| 316 | if (!sessionMeta) { | ||
| 317 | yield* Effect.annotateCurrentSpan({ | ||
| 318 | "mcp.session.restore.outcome": "missing_meta", | ||
| 319 | }); | ||
| 320 | return "missing_meta" as const; | ||
| 321 | } | ||
| 322 | |||
| 323 | yield* self.closeRuntime(); | ||
| 324 | self.dbHandle = makeLongLivedDb(); | ||
| 325 | const runtime = yield* self.createConnectedRuntime(sessionMeta, { | ||
| 326 | dbHandle: self.dbHandle, | ||
| 327 | enableJsonResponse: request.method !== "GET", | ||
| 328 | }); | ||
| 329 | self.mcpServer = runtime.mcpServer; | ||
| 330 | self.transport = runtime.transport; | ||
| 331 | self.initialized = true; | ||
| 332 | self.lastActivityMs = Date.now(); | ||
| 333 | yield* Effect.promise(() => self.ctx.storage.setAlarm(Date.now() + HEARTBEAT_MS)).pipe( | ||
| 334 | Effect.withSpan("McpSessionDO.setAlarm"), | ||
| 335 | ); | ||
| 336 | yield* Effect.annotateCurrentSpan({ | ||
| 337 | "mcp.session.restore.outcome": "restored", | ||
| 338 | }); | ||
| 339 | return "restored" as const; | ||
| 340 | }).pipe( | ||
| 341 | Effect.withSpan("McpSessionDO.restoreRuntime", { | ||
| 342 | attributes: { | ||
| 343 | "mcp.request.method": request.method, | ||
| 344 | "mcp.request.session_id_present": !!request.headers.get("mcp-session-id"), | ||
| 345 | }, | ||
| 346 | }), | ||
| 347 | Effect.orDie, | ||
| 348 | ); | ||
| 349 | } | ||
| 350 | |||
| 277 | private resolveAndStoreSessionMeta(token: McpSessionInit) { | 351 | private resolveAndStoreSessionMeta(token: McpSessionInit) { |
| 278 | const self = this; | 352 | const self = this; |
| 279 | return Effect.gen(function* () { | 353 | return Effect.gen(function* () { |
| @@ -294,9 +368,14 @@ export class McpSessionDO extends DurableObject { · 14 unmodified lines | |||
| 294 | } | 368 | } |
| 295 | 369 | ||
| 296 | async init(token: McpSessionInit, incoming?: IncomingTraceHeaders): Promise<void> { | 370 | async init(token: McpSessionInit, incoming?: IncomingTraceHeaders): Promise<void> { |
| 371 | const methodEnteredAt = Date.now(); | ||
| 297 | if (this.initialized) return; | 372 | if (this.initialized) return; |
| 373 | const self = this; | ||
| 298 | return Effect.runPromise( | 374 | return Effect.runPromise( |
| 299 | this.doInit(token).pipe( | 375 | Effect.gen(function* () { |
| 376 | yield* Effect.annotateCurrentSpan(self.entryAttrs(methodEnteredAt)); | ||
| 377 | yield* self.doInit(token); | ||
| 378 | }).pipe( | ||
| 300 | Effect.withSpan("McpSessionDO.init", { | 379 | Effect.withSpan("McpSessionDO.init", { |
| 301 | attributes: { "mcp.auth.organization_id": token.organizationId }, | 380 | attributes: { "mcp.auth.organization_id": token.organizationId }, |
| 302 | }), | 381 | }), |
| @@ -356,6 +435,7 @@ export class McpSessionDO extends DurableObject { · 53 unmodified lines | |||
| 356 | } | 435 | } |
| 357 | 436 | ||
| 358 | async handleRequest(request: Request): Promise<Response> { | 437 | async handleRequest(request: Request): Promise<Response> { |
| 438 | const methodEnteredAt = Date.now(); | ||
| 359 | // Wrap the dispatch in an Effect span so every DO request — not just | 439 | // Wrap the dispatch in an Effect span so every DO request — not just |
| 360 | // the rare new-session `init()` — shows up in Axiom. Basic attributes | 440 | // the rare new-session `init()` — shows up in Axiom. Basic attributes |
| 361 | // only (method, session-id presence, response status); rich client | 441 | // only (method, session-id presence, response status); rich client |
| @@ -368,6 +448,7 @@ export class McpSessionDO extends DurableObject { · 6 unmodified lines | |||
| 368 | } satisfies IncomingTraceHeaders; | 448 | } satisfies IncomingTraceHeaders; |
| 369 | const self = this; | 449 | const self = this; |
| 370 | const program = Effect.gen(function* () { | 450 | const program = Effect.gen(function* () { |
| 451 | yield* Effect.annotateCurrentSpan(self.entryAttrs(methodEnteredAt)); | ||
| 371 | // Capture the request-entry span so the host-mcp `parentSpan` getter | 452 | // Capture the request-entry span so the host-mcp `parentSpan` getter |
| 372 | // — fired by deferred MCP SDK callbacks after this Effect has already | 453 | // — fired by deferred MCP SDK callbacks after this Effect has already |
| 373 | // returned — anchors engine spans under the same trace. Cleared in a | 454 | // returned — anchors engine spans under the same trace. Cleared in a |
| @@ -403,9 +484,19 @@ export class McpSessionDO extends DurableObject { · 29 unmodified lines | |||
| 403 | 484 | ||
| 404 | private dispatchRequest(request: Request): Effect.Effect<Response> { | 485 | private dispatchRequest(request: Request): Effect.Effect<Response> { |
| 405 | if (!this.initialized || !this.transport) { | 486 | if (!this.initialized || !this.transport) { |
| 406 | return Effect.succeed( | 487 | if (request.method === "DELETE") { |
| 407 | jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect"), | 488 | return this.clearSessionState().pipe( |
| 408 | ); | 489 | Effect.as(new Response(null, { status: 204 })), |
| 490 | Effect.withSpan("mcp.session.stale_delete"), | ||
| 491 | ); | ||
| 492 | } | ||
| 493 | return Effect.gen(this, function* () { | ||
| 494 | const restored = yield* this.restoreRuntimeFromStorage(request); | ||
| 495 | if (restored === "restored") { | ||
| 496 | return yield* this.dispatchRequest(request); | ||
| 497 | } | ||
| 498 | return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect"); | ||
| 499 | }); | ||
| 409 | } | 500 | } |
| 410 | 501 | ||
| 411 | this.lastActivityMs = Date.now(); | 502 | this.lastActivityMs = Date.now(); |
| @@ -461,18 +552,7 @@ export class McpSessionDO extends DurableObject { · 49 unmodified lines | |||
| 461 | } | 552 | } |
| 462 | 553 | ||
| 463 | private async cleanup(): Promise<void> { | 554 | private async cleanup(): Promise<void> { |
| 464 | if (this.transport) { | 555 | await Effect.runPromise(this.closeRuntime()); |
| 465 | await this.transport.close().catch(() => undefined); | ||
| 466 | this.transport = null; | ||
| 467 | } | ||
| 468 | if (this.mcpServer) { | ||
| 469 | await this.mcpServer.close().catch(() => undefined); | ||
| 470 | this.mcpServer = null; | ||
| 471 | } | ||
| 472 | if (this.dbHandle) { | ||
| 473 | await this.dbHandle.end(); | ||
| 474 | this.dbHandle = null; | ||
| 475 | } | ||
| 476 | await Effect.runPromise(this.clearSessionState()); | 556 | await Effect.runPromise(this.clearSessionState()); |
| 477 | } | 557 | } |
| 478 | } | 558 | } |