Unified · Split
| @@ -182,6 +182,7 @@ const resolveSessionMeta = Effect.fn("McpSessionDO.resolveSessionMeta")(function · 181 unmodified lines | ||
| 182 | 182 | // --------------------------------------------------------------------------- |
| 183 | 183 | |
| 184 | 184 | export class McpSessionDO extends DurableObject { |
| 185 | + private readonly instanceCreatedAt = Date.now(); | |
| 185 | 186 | private mcpServer: McpServer | null = null; |
| 186 | 187 | private transport: WorkerTransport | null = null; |
| 187 | 188 | private initialized = false; |
| @@ -223,6 +224,17 @@ export class McpSessionDO extends DurableObject { · 35 unmodified lines | ||
| 223 | 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 | 238 | private clearSessionState(): Effect.Effect<void> { |
| 227 | 239 | return Effect.promise(async () => { |
| 228 | 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 | 351 | private resolveAndStoreSessionMeta(token: McpSessionInit) { |
| 278 | 352 | const self = this; |
| 279 | 353 | return Effect.gen(function* () { |
| @@ -294,9 +368,14 @@ export class McpSessionDO extends DurableObject { · 14 unmodified lines | ||
| 294 | 368 | } |
| 295 | 369 | |
| 296 | 370 | async init(token: McpSessionInit, incoming?: IncomingTraceHeaders): Promise<void> { |
| 371 | + const methodEnteredAt = Date.now(); | |
| 297 | 372 | if (this.initialized) return; |
| 373 | + const self = this; | |
| 298 | 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 | 379 | Effect.withSpan("McpSessionDO.init", { |
| 301 | 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 | 437 | async handleRequest(request: Request): Promise<Response> { |
| 438 | + const methodEnteredAt = Date.now(); | |
| 359 | 439 | // Wrap the dispatch in an Effect span so every DO request — not just |
| 360 | 440 | // the rare new-session `init()` — shows up in Axiom. Basic attributes |
| 361 | 441 | // only (method, session-id presence, response status); rich client |
| @@ -368,6 +448,7 @@ export class McpSessionDO extends DurableObject { · 6 unmodified lines | ||
| 368 | 448 | } satisfies IncomingTraceHeaders; |
| 369 | 449 | const self = this; |
| 370 | 450 | const program = Effect.gen(function* () { |
| 451 | + yield* Effect.annotateCurrentSpan(self.entryAttrs(methodEnteredAt)); | |
| 371 | 452 | // Capture the request-entry span so the host-mcp `parentSpan` getter |
| 372 | 453 | // — fired by deferred MCP SDK callbacks after this Effect has already |
| 373 | 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 | 485 | private dispatchRequest(request: Request): Effect.Effect<Response> { |
| 405 | 486 | if (!this.initialized || !this.transport) { |
| 406 | − return Effect.succeed( | |
| 407 | − jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect"), | |
| 408 | − ); | |
| 487 | + if (request.method === "DELETE") { | |
| 488 | + return this.clearSessionState().pipe( | |
| 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 | 502 | this.lastActivityMs = Date.now(); |
| @@ -461,18 +552,7 @@ export class McpSessionDO extends DurableObject { · 49 unmodified lines | ||
| 461 | 552 | } |
| 462 | 553 | |
| 463 | 554 | private async cleanup(): Promise<void> { |
| 464 | − if (this.transport) { | |
| 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 | − } | |
| 555 | + await Effect.runPromise(this.closeRuntime()); | |
| 476 | 556 | await Effect.runPromise(this.clearSessionState()); |
| 477 | 557 | } |
| 478 | 558 | } |