* Add e2e coverage for concurrent SSE GET + POST reconnect after idle dispose * Serialize MCP session startup * Drop repro notes doc * Neutral wording for the session startup regression test
Unified · Split
| @@ -0,0 +1,215 @@ | |||
| 1 | // Cloud: regression test for a session-DO startup race. | ||
| 2 | // | ||
| 3 | // 1. Open an MCP session (real SDK client init, capture mcp-session-id), | ||
| 4 | // make a call, close the client. | ||
| 5 | // 2. Idle past the session timeout so the DO's alarm runs disposeIdleRuntime | ||
| 6 | // (runtime torn down, `initialized = false`, transport cleared). | ||
| 7 | // 3. Fire an SSE GET (listen stream) and a POST (tools/list) concurrently on | ||
| 8 | // the same session id. | ||
| 9 | // 4. The session must survive: the concurrent pair and follow-up calls | ||
| 10 | // should succeed. | ||
| 11 | // | ||
| 12 | // Mechanism guarded against: a request carrying a session id runs | ||
| 13 | // validateMcpSessionOwner, which restores via onStart when the runtime is | ||
| 14 | // disposed. The SDK's serve() streaming handler separately wakes the DO | ||
| 15 | // through agent.fetch, and the Agent base also drives onStart. Without | ||
| 16 | // serialization across both entry paths, two onStart calls interleave and the | ||
| 17 | // second server.connect throws "Already connected to a transport", after | ||
| 18 | // which every request on the session fails the same way. The concurrent pair | ||
| 19 | // itself can answer 200/200 (one restore wins the race), so the assertions | ||
| 20 | // cover the follow-up calls, which is where the failure persists. | ||
| 21 | // | ||
| 22 | import { writeFileSync } from "node:fs"; | ||
| 23 | import { join } from "node:path"; | ||
| 24 | |||
| 25 | import { expect } from "@effect/vitest"; | ||
| 26 | import { Effect } from "effect"; | ||
| 27 | import { Client } from "@modelcontextprotocol/sdk/client/index.js"; | ||
| 28 | import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; | ||
| 29 | |||
| 30 | import { scenario } from "../src/scenario"; | ||
| 31 | import { Mcp, RunDir, Target } from "../src/services"; | ||
| 32 | import type { Identity } from "../src/target"; | ||
| 33 | import { configuredMcpSessionTimeoutMs } from "../setup/mcp-session-timeouts"; | ||
| 34 | |||
| 35 | const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label; | ||
| 36 | |||
| 37 | // Give the DO's idle alarm room to fire disposeIdleRuntime after the timeout. | ||
| 38 | const IDLE_DISPOSE_BUFFER_MS = 2_000; | ||
| 39 | const IDLE_DISPOSE_GAP_MS = configuredMcpSessionTimeoutMs() + IDLE_DISPOSE_BUFFER_MS; | ||
| 40 | const SCENARIO_TIMEOUT_MS = IDLE_DISPOSE_GAP_MS + 120_000; | ||
| 41 | |||
| 42 | interface Connected { | ||
| 43 | readonly client: Client; | ||
| 44 | readonly transport: StreamableHTTPClientTransport; | ||
| 45 | } | ||
| 46 | |||
| 47 | const connectClient = async (mcpUrl: string, bearer: string): Promise<Connected> => { | ||
| 48 | const client = new Client({ name: "executor-e2e-brick", version: "0.0.1" }, { capabilities: {} }); | ||
| 49 | const transport = new StreamableHTTPClientTransport(new URL(mcpUrl), { | ||
| 50 | requestInit: { headers: { authorization: `Bearer ${bearer}` } }, | ||
| 51 | }); | ||
| 52 | await client.connect(transport); | ||
| 53 | return { client, transport }; | ||
| 54 | }; | ||
| 55 | |||
| 56 | interface RawResult { | ||
| 57 | readonly status: number; | ||
| 58 | readonly body: string; | ||
| 59 | } | ||
| 60 | |||
| 61 | /** A raw POST of a JSON-RPC request onto an existing session id. */ | ||
| 62 | const rawPost = async ( | ||
| 63 | mcpUrl: string, | ||
| 64 | bearer: string, | ||
| 65 | sessionId: string, | ||
| 66 | message: unknown, | ||
| 67 | ): Promise<RawResult> => { | ||
| 68 | const res = await fetch(mcpUrl, { | ||
| 69 | method: "POST", | ||
| 70 | headers: { | ||
| 71 | authorization: `Bearer ${bearer}`, | ||
| 72 | "content-type": "application/json", | ||
| 73 | // The streamable-http transport requires the client accept both. | ||
| 74 | accept: "application/json, text/event-stream", | ||
| 75 | "mcp-session-id": sessionId, | ||
| 76 | "mcp-protocol-version": "2025-06-18", | ||
| 77 | }, | ||
| 78 | body: JSON.stringify(message), | ||
| 79 | }); | ||
| 80 | const body = await res.text().catch(() => ""); | ||
| 81 | return { status: res.status, body }; | ||
| 82 | }; | ||
| 83 | |||
| 84 | /** A raw SSE GET (the "listen" stream) on an existing session id. We only need | ||
| 85 | * the response head — status + start of body — to see whether it 500s. */ | ||
| 86 | const rawSseGet = async (mcpUrl: string, bearer: string, sessionId: string): Promise<RawResult> => { | ||
| 87 | const res = await fetch(mcpUrl, { | ||
| 88 | method: "GET", | ||
| 89 | headers: { | ||
| 90 | authorization: `Bearer ${bearer}`, | ||
| 91 | accept: "text/event-stream", | ||
| 92 | "mcp-session-id": sessionId, | ||
| 93 | "mcp-protocol-version": "2025-06-18", | ||
| 94 | }, | ||
| 95 | }); | ||
| 96 | // A 200 opens a long-lived stream; don't hang reading it. A 500 has a short | ||
| 97 | // (usually empty) body we can read fully. | ||
| 98 | let body = ""; | ||
| 99 | if (res.status !== 200) { | ||
| 100 | body = await res.text().catch(() => ""); | ||
| 101 | } else { | ||
| 102 | // Cancel the stream immediately; we only wanted the status. | ||
| 103 | await res.body?.cancel().catch(() => undefined); | ||
| 104 | } | ||
| 105 | return { status: res.status, body }; | ||
| 106 | }; | ||
| 107 | |||
| 108 | const toolsListMessage = { jsonrpc: "2.0" as const, id: "brick-tools-list", method: "tools/list" }; | ||
| 109 | |||
| 110 | const isBrick = (r: RawResult): boolean => | ||
| 111 | r.status >= 500 || /"code"\s*:\s*-32603/.test(r.body) || /Already connected/i.test(r.body); | ||
| 112 | |||
| 113 | scenario( | ||
| 114 | "REGRESSION · concurrent SSE GET + POST after idle-dispose keeps the session alive", | ||
| 115 | { timeout: SCENARIO_TIMEOUT_MS }, | ||
| 116 | Effect.gen(function* () { | ||
| 117 | const target = yield* Target; | ||
| 118 | const mcp = yield* Mcp; | ||
| 119 | const runDir = yield* RunDir; | ||
| 120 | const identity = yield* target.newIdentity(); | ||
| 121 | const bearer = yield* mcp.mintBearer(emailOf(identity)); | ||
| 122 | const diagnostics: unknown[] = []; | ||
| 123 | const record = (entry: Record<string, unknown>): void => { | ||
| 124 | diagnostics.push(entry); | ||
| 125 | console.log(JSON.stringify(entry)); | ||
| 126 | // Vitest swallows worker stdout on passing runs; the artifact dir is the | ||
| 127 | // durable copy of what each attempt observed. | ||
| 128 | writeFileSync(join(runDir, "repro-diagnostics.json"), JSON.stringify(diagnostics, null, 2)); | ||
| 129 | }; | ||
| 130 | |||
| 131 | // 1. Open a real session (SDK init handshake persists the session in the | ||
| 132 | // DO), do a call, close the client. | ||
| 133 | const first = yield* Effect.promise(() => connectClient(target.mcpUrl, bearer)); | ||
| 134 | const sessionId = first.transport.sessionId; | ||
| 135 | expect(sessionId, "the client got a session id").toEqual(expect.any(String)); | ||
| 136 | if (sessionId === undefined) return yield* Effect.die("missing session id"); | ||
| 137 | |||
| 138 | yield* Effect.promise(() => | ||
| 139 | first.client.callTool({ name: "execute", arguments: { code: "return 1;" } }), | ||
| 140 | ); | ||
| 141 | yield* Effect.promise(() => first.client.close().catch(() => undefined)); | ||
| 142 | |||
| 143 | // 2. Idle past the timeout so the DO's alarm runs disposeIdleRuntime. | ||
| 144 | yield* Effect.sleep(IDLE_DISPOSE_GAP_MS); | ||
| 145 | |||
| 146 | // 3. Fire the SSE GET and the POST CONCURRENTLY on the same session id. | ||
| 147 | // Raw fetch gives precise concurrency control (both start before either | ||
| 148 | // resolves) that the SDK client's request queue would serialize away. | ||
| 149 | const [getResult, postResult] = yield* Effect.promise(() => | ||
| 150 | Promise.all([ | ||
| 151 | rawSseGet(target.mcpUrl, bearer, sessionId), | ||
| 152 | rawPost(target.mcpUrl, bearer, sessionId, toolsListMessage), | ||
| 153 | ]), | ||
| 154 | ); | ||
| 155 | |||
| 156 | record({ | ||
| 157 | event: "repro_concurrent_pair", | ||
| 158 | sessionId, | ||
| 159 | get: getResult, | ||
| 160 | post: postResult, | ||
| 161 | }); | ||
| 162 | |||
| 163 | const concurrentBricked = isBrick(getResult) || isBrick(postResult); | ||
| 164 | |||
| 165 | // 4. Prove the session survives: subsequent SEQUENTIAL POSTs also succeed. | ||
| 166 | // This is the same control path as the idle-restore scenario from #1302. | ||
| 167 | // | ||
| 168 | // Observed locally: the concurrent pair frequently returns 200/200 (one | ||
| 169 | // of the two restores wins the race and answers), but the race leaves | ||
| 170 | // the DO's McpServer in the "Already connected to a transport" state, so | ||
| 171 | // the very NEXT request 500s — and every request after that. The | ||
| 172 | // regression gate checks the same two follow-up calls now stay healthy. | ||
| 173 | const followUp = yield* Effect.promise(() => | ||
| 174 | rawPost(target.mcpUrl, bearer, sessionId, { | ||
| 175 | ...toolsListMessage, | ||
| 176 | id: "brick-follow-up", | ||
| 177 | }), | ||
| 178 | ); | ||
| 179 | const followUp2 = yield* Effect.promise(() => | ||
| 180 | rawPost(target.mcpUrl, bearer, sessionId, { | ||
| 181 | ...toolsListMessage, | ||
| 182 | id: "brick-follow-up-2", | ||
| 183 | }), | ||
| 184 | ); | ||
| 185 | record({ event: "repro_follow_up", sessionId, followUp, followUp2 }); | ||
| 186 | |||
| 187 | const permanentlyBricked = isBrick(followUp) && isBrick(followUp2); | ||
| 188 | |||
| 189 | expect(getResult.status, "the concurrent SSE GET opens successfully").toBe(200); | ||
| 190 | expect(postResult.status, "the concurrent POST answers successfully").toBe(200); | ||
| 191 | expect(isBrick(getResult), "the concurrent SSE GET does not expose the transport brick").toBe( | ||
| 192 | false, | ||
| 193 | ); | ||
| 194 | expect(isBrick(postResult), "the concurrent POST does not expose the transport brick").toBe( | ||
| 195 | false, | ||
| 196 | ); | ||
| 197 | expect(followUp.status, "the first follow-up POST answers successfully").toBe(200); | ||
| 198 | expect(followUp2.status, "the second follow-up POST answers successfully").toBe(200); | ||
| 199 | expect(isBrick(followUp), "the first follow-up POST does not expose the transport brick").toBe( | ||
| 200 | false, | ||
| 201 | ); | ||
| 202 | expect( | ||
| 203 | isBrick(followUp2), | ||
| 204 | "the second follow-up POST does not expose the transport brick", | ||
| 205 | ).toBe(false); | ||
| 206 | expect( | ||
| 207 | permanentlyBricked, | ||
| 208 | "after the concurrent SSE-GET + POST post-idle, the session remains usable", | ||
| 209 | ).toBe(false); | ||
| 210 | // Diagnostic (not a gate): whether the collision surfaced on the concurrent | ||
| 211 | // pair itself vs. only on the follow-up depends on which restore won the | ||
| 212 | // race. Record it so the artifact shows the timing that occurred this run. | ||
| 213 | record({ event: "repro_variant", concurrentBricked, permanentlyBricked }); | ||
| 214 | }), | ||
| 215 | ); | ||
| @@ -1,12 +1,18 @@ | |||
| 1 | import { describe, expect, it } from "@effect/vitest"; | 1 | import { describe, expect, it } from "@effect/vitest"; |
| 2 | import { Effect } from "effect"; | 2 | import { Cause, Effect } from "effect"; |
| 3 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; | 3 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; |
| 4 | import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; | 4 | import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; |
| 5 | import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; | 5 | import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; |
| 6 | 6 | ||
| 7 | import { defaultMcpResource } from "@executor-js/host-mcp"; | 7 | import { defaultMcpResource } from "@executor-js/host-mcp"; |
| 8 | import type { ExecutionEngine, ExecutionResult, ResumeResponse } from "@executor-js/execution"; | ||
| 8 | 9 | ||
| 9 | import { McpAgentSessionDOBase, type SessionMeta } from "./agent-session-durable-object"; | 10 | import { |
| 11 | McpAgentSessionDOBase, | ||
| 12 | type McpApprovalOwner, | ||
| 13 | type McpSessionModelResumeResult, | ||
| 14 | type SessionMeta, | ||
| 15 | } from "./agent-session-durable-object"; | ||
| 10 | 16 | ||
| 11 | class MemoryStorage { | 17 | class MemoryStorage { |
| 12 | private readonly data = new Map<string, unknown>(); | 18 | private readonly data = new Map<string, unknown>(); |
| @@ -77,7 +83,7 @@ type HarnessSession = { · 64 unmodified lines | |||
| 77 | alarm: () => Promise<void>; | 83 | alarm: () => Promise<void>; |
| 78 | ctx: MemoryStorage; | 84 | ctx: MemoryStorage; |
| 79 | dbHandle: { readonly end: () => void } | null; | 85 | dbHandle: { readonly end: () => void } | null; |
| 80 | engine: { readonly pausedExecutionCount: () => Effect.Effect<number> } | null; | 86 | engine: ExecutionEngine<Cause.YieldableError> | null; |
| 81 | getSessionId: () => string; | 87 | getSessionId: () => string; |
| 82 | initialized: boolean; | 88 | initialized: boolean; |
| 83 | lastActivityMs: number; | 89 | lastActivityMs: number; |
| @@ -85,9 +91,15 @@ type HarnessSession = { · 1 unmodified line | |||
| 85 | onStart: () => Promise<void>; | 91 | onStart: () => Promise<void>; |
| 86 | pendingApprovalLeases: Map<string, never>; | 92 | pendingApprovalLeases: Map<string, never>; |
| 87 | props: Record<string, unknown>; | 93 | props: Record<string, unknown>; |
| 94 | runMcpAgentOnStart: () => Promise<void>; | ||
| 88 | server?: McpServer; | 95 | server?: McpServer; |
| 89 | sessionMeta: SessionMeta; | 96 | sessionMeta: SessionMeta; |
| 90 | sessionTimeoutMs: () => number; | 97 | sessionTimeoutMs: () => number; |
| 98 | resumeExecutionForModel: ( | ||
| 99 | executionId: string, | ||
| 100 | identity: McpApprovalOwner, | ||
| 101 | response: ResumeResponse, | ||
| 102 | ) => Promise<McpSessionModelResumeResult>; | ||
| 91 | validateMcpSessionOwner: (identity: { | 103 | validateMcpSessionOwner: (identity: { |
| 92 | readonly accountId: string; | 104 | readonly accountId: string; |
| 93 | readonly organizationId: string; | 105 | readonly organizationId: string; |
| @@ -130,6 +142,44 @@ const makeDeferred = (): { readonly promise: Promise<void>; readonly resolve: () · 36 unmodified lines | |||
| 130 | return { promise, resolve }; | 142 | return { promise, resolve }; |
| 131 | }; | 143 | }; |
| 132 | 144 | ||
| 145 | type ResumeCall = { | ||
| 146 | readonly executionId: string; | ||
| 147 | readonly response: ResumeResponse; | ||
| 148 | }; | ||
| 149 | |||
| 150 | const completed = (result: unknown): ExecutionResult => ({ | ||
| 151 | status: "completed", | ||
| 152 | result: { result }, | ||
| 153 | }); | ||
| 154 | |||
| 155 | const makeEngine = ( | ||
| 156 | resultForResume: (executionId: string, response: ResumeResponse) => ExecutionResult | null = () => | ||
| 157 | completed("resume-result"), | ||
| 158 | ): { readonly calls: ResumeCall[]; readonly engine: ExecutionEngine<Cause.YieldableError> } => { | ||
| 159 | const calls: ResumeCall[] = []; | ||
| 160 | return { | ||
| 161 | calls, | ||
| 162 | engine: { | ||
| 163 | execute: () => Effect.succeed({ result: "execute-result" }), | ||
| 164 | executeWithPause: () => Effect.succeed(completed("execute-result")), | ||
| 165 | resume: (executionId, response) => | ||
| 166 | Effect.sync(() => { | ||
| 167 | calls.push({ executionId, response }); | ||
| 168 | return resultForResume(executionId, response); | ||
| 169 | }), | ||
| 170 | getPausedExecution: () => Effect.succeed(null), | ||
| 171 | pausedExecutionCount: () => Effect.succeed(0), | ||
| 172 | hasPausedExecutions: () => Effect.succeed(false), | ||
| 173 | getDescription: Effect.succeed("test engine"), | ||
| 174 | }, | ||
| 175 | }; | ||
| 176 | }; | ||
| 177 | |||
| 178 | const approval = { | ||
| 179 | action: "accept", | ||
| 180 | content: { approved: true }, | ||
| 181 | } satisfies ResumeResponse; | ||
| 182 | |||
| 133 | const makeHarnessSession = async (): Promise<HarnessSession> => { | 183 | const makeHarnessSession = async (): Promise<HarnessSession> => { |
| 134 | const sessionId = "session-reconnect"; | 184 | const sessionId = "session-reconnect"; |
| 135 | const sessionMeta: SessionMeta = { | 185 | const sessionMeta: SessionMeta = { |
| @@ -145,7 +195,7 @@ const makeHarnessSession = async (): Promise<HarnessSession> => { · 9 unmodified lines | |||
| 145 | const session = Object.create(McpAgentSessionDOBase.prototype) as HarnessSession; | 195 | const session = Object.create(McpAgentSessionDOBase.prototype) as HarnessSession; |
| 146 | session.ctx = storage; | 196 | session.ctx = storage; |
| 147 | session.dbHandle = { end: () => undefined }; | 197 | session.dbHandle = { end: () => undefined }; |
| 148 | session.engine = { pausedExecutionCount: () => Effect.succeed(0) }; | 198 | session.engine = makeEngine().engine; |
| 149 | session.getSessionId = () => sessionId; | 199 | session.getSessionId = () => sessionId; |
| 150 | session.initialized = true; | 200 | session.initialized = true; |
| 151 | session.lastActivityMs = Date.now() - 10; | 201 | session.lastActivityMs = Date.now() - 10; |
| @@ -155,10 +205,11 @@ const makeHarnessSession = async (): Promise<HarnessSession> => { · 3 unmodified lines | |||
| 155 | session.server = server; | 205 | session.server = server; |
| 156 | session.sessionMeta = sessionMeta; | 206 | session.sessionMeta = sessionMeta; |
| 157 | session.sessionTimeoutMs = () => 1; | 207 | session.sessionTimeoutMs = () => 1; |
| 158 | session.onStart = async () => { | 208 | session.runMcpAgentOnStart = async () => { |
| 159 | const restored = session.server ?? makeServer(); | 209 | const restored = session.server ?? makeServer(); |
| 160 | session.server = restored; | 210 | session.server = restored; |
| 161 | await restored.connect(new RestoredTransport()); | 211 | await restored.connect(new RestoredTransport()); |
| 212 | session.engine = makeEngine().engine; | ||
| 162 | session.initialized = true; | 213 | session.initialized = true; |
| 163 | }; | 214 | }; |
| 164 | 215 | ||
| @@ -183,7 +234,7 @@ describe("McpAgentSessionDOBase transport restore", () => { · 18 unmodified lines | |||
| 183 | let onStartCalls = 0; | 234 | let onStartCalls = 0; |
| 184 | let restoredServer: McpServer | undefined; | 235 | let restoredServer: McpServer | undefined; |
| 185 | 236 | ||
| 186 | session.onStart = async () => { | 237 | session.runMcpAgentOnStart = async () => { |
| 187 | onStartCalls += 1; | 238 | onStartCalls += 1; |
| 188 | const restored = session.server ?? makeServer(); | 239 | const restored = session.server ?? makeServer(); |
| 189 | restoredServer ??= restored; | 240 | restoredServer ??= restored; |
| @@ -213,4 +264,81 @@ describe("McpAgentSessionDOBase transport restore", () => { · 23 unmodified lines | |||
| 213 | expect(onStartCalls).toBe(1); | 264 | expect(onStartCalls).toBe(1); |
| 214 | expect(session.server).toBe(restoredServer); | 265 | expect(session.server).toBe(restoredServer); |
| 215 | }); | 266 | }); |
| 267 | |||
| 268 | it("single-flights SDK onStart callers with same-session restore", async () => { | ||
| 269 | const session = await makeHarnessSession(); | ||
| 270 | const firstStartEntered = makeDeferred(); | ||
| 271 | const finishStart = makeDeferred(); | ||
| 272 | let onStartCalls = 0; | ||
| 273 | |||
| 274 | session.runMcpAgentOnStart = async () => { | ||
| 275 | onStartCalls += 1; | ||
| 276 | const restored = session.server ?? makeServer(); | ||
| 277 | session.server = restored; | ||
| 278 | firstStartEntered.resolve(); | ||
| 279 | await finishStart.promise; | ||
| 280 | await restored.connect(new RestoredTransport()); | ||
| 281 | session.initialized = true; | ||
| 282 | }; | ||
| 283 | |||
| 284 | await session.alarm(); | ||
| 285 | |||
| 286 | const restore = session.validateMcpSessionOwner({ | ||
| 287 | accountId: "user-1", | ||
| 288 | organizationId: "org-1", | ||
| 289 | }); | ||
| 290 | const sdkStart = session.onStart(); | ||
| 291 | |||
| 292 | await firstStartEntered.promise; | ||
| 293 | await Promise.resolve(); | ||
| 294 | finishStart.resolve(); | ||
| 295 | |||
| 296 | await expect(Promise.all([restore, sdkStart])).resolves.toEqual(["ok", undefined]); | ||
| 297 | expect(onStartCalls).toBe(1); | ||
| 298 | }); | ||
| 299 | |||
| 300 | it("single-flights model resume restore with SDK onStart", async () => { | ||
| 301 | const session = await makeHarnessSession(); | ||
| 302 | const firstStartEntered = makeDeferred(); | ||
| 303 | const finishStart = makeDeferred(); | ||
| 304 | const restoredEngine = makeEngine(() => completed("model-result")); | ||
| 305 | let onStartCalls = 0; | ||
| 306 | |||
| 307 | session.runMcpAgentOnStart = async () => { | ||
| 308 | onStartCalls += 1; | ||
| 309 | const restored = session.server ?? makeServer(); | ||
| 310 | session.server = restored; | ||
| 311 | firstStartEntered.resolve(); | ||
| 312 | await finishStart.promise; | ||
| 313 | await restored.connect(new RestoredTransport()); | ||
| 314 | session.engine = restoredEngine.engine; | ||
| 315 | session.initialized = true; | ||
| 316 | }; | ||
| 317 | |||
| 318 | await session.alarm(); | ||
| 319 | |||
| 320 | const resume = session.resumeExecutionForModel( | ||
| 321 | "exec-model", | ||
| 322 | { accountId: "user-1", organizationId: "org-1" }, | ||
| 323 | approval, | ||
| 324 | ); | ||
| 325 | const sdkStart = session.onStart(); | ||
| 326 | |||
| 327 | await firstStartEntered.promise; | ||
| 328 | await Promise.resolve(); | ||
| 329 | finishStart.resolve(); | ||
| 330 | |||
| 331 | const [resumeResult] = await Promise.all([resume, sdkStart]); | ||
| 332 | expect(resumeResult).toMatchObject({ | ||
| 333 | status: "result", | ||
| 334 | result: { | ||
| 335 | structuredContent: { | ||
| 336 | status: "completed", | ||
| 337 | result: "model-result", | ||
| 338 | }, | ||
| 339 | }, | ||
| 340 | }); | ||
| 341 | expect(onStartCalls).toBe(1); | ||
| 342 | expect(restoredEngine.calls).toEqual([{ executionId: "exec-model", response: approval }]); | ||
| 343 | }); | ||
| 216 | }); | 344 | }); |
| @@ -212,7 +212,7 @@ export abstract class McpAgentSessionDOBase< · 211 unmodified lines | |||
| 212 | private dbHandle: TDbHandle | null = null; | 212 | private dbHandle: TDbHandle | null = null; |
| 213 | private sessionMeta: SessionMeta | null = null; | 213 | private sessionMeta: SessionMeta | null = null; |
| 214 | private initialized = false; | 214 | private initialized = false; |
| 215 | private restoreTransportRuntimePromise: Promise<void> | null = null; | 215 | private onStartPromise: Promise<void> | null = null; |
| 216 | private lastActivityMs = 0; | 216 | private lastActivityMs = 0; |
| 217 | private approvalResponses = new Map<string, ResumeResponse>(); | 217 | private approvalResponses = new Map<string, ResumeResponse>(); |
| 218 | private approvalWaiters = new Map<string, Deferred.Deferred<ResumeResponse>>(); | 218 | private approvalWaiters = new Map<string, Deferred.Deferred<ResumeResponse>>(); |
| @@ -556,55 +556,43 @@ export abstract class McpAgentSessionDOBase< · 337 unmodified lines | |||
| 556 | const sessionMeta = yield* self.loadSessionMeta(); | 556 | const sessionMeta = yield* self.loadSessionMeta(); |
| 557 | if (!sessionMeta) return false; | 557 | if (!sessionMeta) return false; |
| 558 | 558 | ||
| 559 | yield* self.closeRuntime(); | 559 | yield* Effect.promise(() => self.onStart()).pipe( |
| 560 | const dbHandle = yield* self.openSessionDbHandle(); | 560 | Effect.withSpan("McpSessionDO.restore_runtime_for_approval"), |
| 561 | const { mcpServer, engine } = yield* self.buildRuntime(sessionMeta, dbHandle); | ||
| 562 | self.dbHandle = dbHandle; | ||
| 563 | self.server = mcpServer; | ||
| 564 | self.engine = engine; | ||
| 565 | self.initialized = true; | ||
| 566 | yield* Effect.promise(() => self.markActivity()).pipe( | ||
| 567 | Effect.withSpan("McpSessionDO.markActivity"), | ||
| 568 | ); | 561 | ); |
| 569 | return true; | 562 | return self.initialized && !!self.engine; |
| 570 | }).pipe(Effect.withSpan("McpSessionDO.ensure_runtime_for_approval")); | 563 | }).pipe(Effect.withSpan("McpSessionDO.ensure_runtime_for_approval")); |
| 571 | } | 564 | } |
| 572 | 565 | ||
| 573 | private restoreTransportRuntimeOnce(): Effect.Effect<void> { | 566 | private startRuntimeFromOnStart(props?: McpSessionProps): Effect.Effect<void> { |
| 574 | const self = this; | 567 | const self = this; |
| 575 | return Effect.gen(function* () { | 568 | return Effect.gen(function* () { |
| 576 | yield* self.closeRuntime(); | 569 | yield* self.closeRuntime(); |
| 577 | const restored = yield* Effect.exit(Effect.promise(() => self.onStart())); | 570 | const started = yield* Effect.exit(Effect.promise(() => self.runMcpAgentOnStart(props))); |
| 578 | if (Exit.isFailure(restored)) { | 571 | if (Exit.isFailure(started)) { |
| 579 | yield* self.closeRuntime(); | 572 | yield* self.closeRuntime(); |
| 580 | return yield* Effect.failCause(restored.cause); | 573 | return yield* Effect.failCause(started.cause); |
| 581 | } | 574 | } |
| 582 | }); | 575 | }); |
| 583 | } | 576 | } |
| 584 | 577 | ||
| 585 | private restoreTransportRuntime(): Effect.Effect<void> { | 578 | protected runMcpAgentOnStart(props?: McpSessionProps): Promise<void> { |
| 586 | const self = this; | 579 | return super.onStart(props); |
| 587 | return Effect.promise(() => { | 580 | } |
| 588 | if (self.restoreTransportRuntimePromise) return self.restoreTransportRuntimePromise; | ||
| 589 | 581 | ||
| 590 | const restoring = Promise.resolve().then(() => | 582 | override async onStart(props?: McpSessionProps): Promise<void> { |
| 591 | Effect.runPromise(self.restoreTransportRuntimeOnce()), | 583 | if (this.onStartPromise) return this.onStartPromise; |
| 592 | ); | 584 | |
| 593 | self.restoreTransportRuntimePromise = restoring; | 585 | const starting = Effect.runPromise(this.startRuntimeFromOnStart(props)); |
| 594 | restoring.then( | 586 | this.onStartPromise = starting; |
| 595 | () => { | 587 | starting.then( |
| 596 | if (self.restoreTransportRuntimePromise === restoring) { | 588 | () => { |
| 597 | self.restoreTransportRuntimePromise = null; | 589 | if (this.onStartPromise === starting) this.onStartPromise = null; |
| 598 | } | 590 | }, |
| 599 | }, | 591 | () => { |
| 600 | () => { | 592 | if (this.onStartPromise === starting) this.onStartPromise = null; |
| 601 | if (self.restoreTransportRuntimePromise === restoring) { | 593 | }, |
| 602 | self.restoreTransportRuntimePromise = null; | 594 | ); |
| 603 | } | 595 | return starting; |
| 604 | }, | ||
| 605 | ); | ||
| 606 | return restoring; | ||
| 607 | }); | ||
| 608 | } | 596 | } |
| 609 | 597 | ||
| 610 | async init(): Promise<void> { | 598 | async init(): Promise<void> { |
| @@ -669,9 +657,9 @@ export abstract class McpAgentSessionDOBase< · 58 unmodified lines | |||
| 669 | Effect.withSpan("McpSessionDO.markActivity"), | 657 | Effect.withSpan("McpSessionDO.markActivity"), |
| 670 | ); | 658 | ); |
| 671 | } else { | 659 | } else { |
| 672 | yield* self | 660 | yield* Effect.promise(() => self.onStart()).pipe( |
| 673 | .restoreTransportRuntime() | 661 | Effect.withSpan("McpSessionDO.restore_transport_runtime"), |
| 674 | .pipe(Effect.withSpan("McpSessionDO.restore_transport_runtime")); | 662 | ); |
| 675 | } | 663 | } |
| 676 | return identity.accountId === sessionMeta.userId && | 664 | return identity.accountId === sessionMeta.userId && |
| 677 | identity.organizationId === sessionMeta.organizationId | 665 | identity.organizationId === sessionMeta.organizationId |