diff --git a/apps/api/src/platform/recoverable-promise-memo.test.ts b/apps/api/src/platform/recoverable-promise-memo.test.ts new file mode 100644 index 000000000..842c9e019 --- /dev/null +++ b/apps/api/src/platform/recoverable-promise-memo.test.ts @@ -0,0 +1,83 @@ +import { describe, expect, it } from "vitest" +import { Effect, Layer, ManagedRuntime } from "effect" +import { HttpRouter, HttpServerResponse } from "effect/unstable/http" +import { makeRecoverablePromiseMemo } from "./recoverable-promise-memo" + +describe("makeRecoverablePromiseMemo", () => { + it("shares a concurrent build, retains success, and retries rejection", async () => { + let builds = 0 + const memo = makeRecoverablePromiseMemo(async () => { + builds++ + if (builds === 1) throw new Error("first build failed") + return { build: builds } + }) + + const first = memo.get() + expect(memo.get()).toBe(first) + await expect(first).rejects.toThrow("first build failed") + + const recovered = await memo.get() + expect(recovered).toEqual({ build: 2 }) + expect(await memo.get()).toBe(recovered) + expect(builds).toBe(2) + }) + + it("evicts a handler after rejected lazy layer acquisition", async () => { + let acquisitions = 0 + const memo = makeRecoverablePromiseMemo(async () => { + const acquisition = Layer.effectDiscard( + Effect.sync(() => { + acquisitions++ + if (acquisitions === 1) throw new Error("first handler acquisition failed") + }), + ) + const routes = Layer.merge( + HttpRouter.use((router) => router.add("GET", "/test", HttpServerResponse.text("OK"))), + acquisition, + ) + return HttpRouter.toWebHandler(routes, { disableLogger: true }) + }) + + const firstPending = memo.get() + expect(memo.get()).toBe(firstPending) + const first = await firstPending + await expect(first.handler(new Request("https://worker.invalid/test"))).rejects.toThrow( + "first handler acquisition failed", + ) + expect(memo.evict(firstPending)).toBe(true) + expect(memo.evict(firstPending)).toBe(false) + await first.dispose() + + const recovered = await memo.get() + expect((await recovered.handler(new Request("https://worker.invalid/test"))).status).toBe(200) + expect(acquisitions).toBe(2) + await recovered.dispose() + }) + + it("retries a rejected ManagedRuntime layer acquisition", async () => { + let acquisitions = 0 + const memo = makeRecoverablePromiseMemo(async () => { + const runtime = ManagedRuntime.make( + Layer.effectDiscard( + Effect.sync(() => { + acquisitions++ + if (acquisitions === 1) throw new Error("first runtime acquisition failed") + }), + ), + ) + try { + await runtime.context() + return runtime + } catch (error) { + await runtime.dispose() + throw error + } + }) + + await expect(memo.get()).rejects.toThrow("first runtime acquisition failed") + const recovered = await memo.get() + expect(await recovered.runPromise(Effect.succeed("ok"))).toBe("ok") + expect(acquisitions).toBe(2) + await recovered.dispose() + }) +}) diff --git a/apps/api/src/platform/recoverable-promise-memo.ts b/apps/api/src/platform/recoverable-promise-memo.ts new file mode 100644 index 000000000..16ed73cbf --- /dev/null +++ b/apps/api/src/platform/recoverable-promise-memo.ts @@ -0,0 +1,29 @@ +/** + * Share one in-flight/successful asynchronous build and evict a rejected build. + * + * The identity guard matters when callers overlap: a late rejection from an + * older build must never clear the newer promise that a retry already installed. + */ +export const makeRecoverablePromiseMemo = , A>( + build: (...args: Args) => Promise, +) => { + let current: Promise | undefined + + const evict = (expected: Promise): boolean => { + if (current !== expected) return false + current = undefined + return true + } + + const get = (...args: Args): Promise => { + if (current !== undefined) return current + const pending = build(...args) + current = pending + void pending.catch(() => { + evict(pending) + }) + return pending + } + + return { get, evict } +} diff --git a/apps/api/src/worker.ts b/apps/api/src/worker.ts index 022eae59e..42af24426 100644 --- a/apps/api/src/worker.ts +++ b/apps/api/src/worker.ts @@ -19,6 +19,7 @@ import { serverErrorSpanMiddleware } from "./http/server-error-span" import { v2WorkerUnavailableResponse } from "./http/v2-worker-unavailable" import { API_CORS_RESPONSE_HEADERS, apiCorsPreflightResponse } from "./http/api-cors" import { persistSession, preloadSession, type SessionsBinding } from "./mcp/lib/session-store" +import { makeRecoverablePromiseMemo } from "./platform/recoverable-promise-memo" import { classifyWorkerQueue } from "./queue-dispatch" const WorkerFileSystemLive = FileSystem.layerNoop({}) @@ -137,16 +138,7 @@ const buildHandler = async () => { // Memoized via the build promise so concurrent first requests share one build. // A rejected build is cleared after those callers observe it, allowing a later // request to recover instead of pinning the isolate to a rejected promise. -let handlerPromise: ReturnType | undefined -const getHandler = (): ReturnType => { - if (handlerPromise !== undefined) return handlerPromise - const pending = buildHandler() - handlerPromise = pending - void pending.catch(() => { - if (handlerPromise === pending) handlerPromise = undefined - }) - return pending -} +const handlerMemo = makeRecoverablePromiseMemo(buildHandler) // RPC has no HttpApi request to construct the application services for it, so // it gets a sibling isolate-wide ManagedRuntime. Its headless service graph @@ -156,7 +148,7 @@ const buildRpcRuntime = async (env: Record) => { import("./runtime/mcp-service-graph"), import("@/platform/DatabasePgLive"), ]) - return ManagedRuntime.make( + const runtime = ManagedRuntime.make( InvestigationServicesLive.pipe( Layer.provideMerge(WorkerPlatformLive), Layer.provideMerge(layerPg), @@ -165,18 +157,19 @@ const buildRpcRuntime = async (env: Record) => { Layer.provideMerge(WorkerConfigProviderLayer), ), ) + try { + // ManagedRuntime also acquires lazily and retains a failed build fiber. + // Acquire before resolving the recoverable outer promise so a later RPC + // can construct a fresh runtime after an initialization failure. + await runtime.context() + return runtime + } catch (error) { + await runtime.dispose() + throw error + } } -let rpcRuntimePromise: ReturnType | undefined -const getRpcRuntime = (env: Record): ReturnType => { - if (rpcRuntimePromise !== undefined) return rpcRuntimePromise - const pending = buildRpcRuntime(env) - rpcRuntimePromise = pending - void pending.catch(() => { - if (rpcRuntimePromise === pending) rpcRuntimePromise = undefined - }) - return pending -} +const rpcRuntimeMemo = makeRecoverablePromiseMemo(buildRpcRuntime) type InternalRpcMethod = "listMcpTools" | "callMcpTool" | "submitDiagnosis" @@ -207,7 +200,7 @@ const runInternalRpc = async ( ctx: ExecutionContext, ) => { const [runtime, { callMcpToolRpc, listMcpToolsRpc, submitDiagnosisRpc }] = await Promise.all([ - getRpcRuntime(env), + rpcRuntimeMemo.get(env), import("./internal-rpc"), ]) let exit: Exit.Exit @@ -326,7 +319,7 @@ const handle = async ( // Start the expensive cold handler build and the independent KV read before // buffering an MCP body. Warm requests resolve both promises immediately; // cold MCP requests hide module evaluation and KV latency behind body I/O. - const pendingHandler = getHandler() + const pendingHandler = handlerMemo.get() const pendingSession = kv && reqSid ? preloadSession(kv, reqSid) : undefined // MCP diagnostics: buffer the body so we can peek the JSON-RPC method/id @@ -350,10 +343,19 @@ const handle = async ( } try { - const { handler } = pendingSession + const built = pendingSession ? (await Promise.all([pendingHandler, pendingSession]))[0] : await pendingHandler - const response = await handler(forwardRequest, HandlerContext) + let response: Response + try { + response = await built.handler(forwardRequest, HandlerContext) + } catch (error) { + // `toWebHandler` acquires lazily and pins a rejected inner build. + // Evict only the exact wrapper used by this request so the next real + // request can rebuild it; overlapping failures cannot clear a retry. + if (handlerMemo.evict(pendingHandler)) await built.dispose() + throw error + } if (kv && isMcp) { const resSid = response.headers.get("mcp-session-id") // Only persist when the server issued a new session — i.e. on