Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 83 additions & 0 deletions apps/api/src/platform/recoverable-promise-memo.test.ts
Original file line number Diff line number Diff line change
@@ -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()
})
})
29 changes: 29 additions & 0 deletions apps/api/src/platform/recoverable-promise-memo.ts
Original file line number Diff line number Diff line change
@@ -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 = <Args extends ReadonlyArray<unknown>, A>(
build: (...args: Args) => Promise<A>,
) => {
let current: Promise<A> | undefined

const evict = (expected: Promise<A>): boolean => {
if (current !== expected) return false
current = undefined
return true
}

const get = (...args: Args): Promise<A> => {
if (current !== undefined) return current
const pending = build(...args)
current = pending
void pending.catch(() => {
evict(pending)
})
return pending
}

return { get, evict }
}
52 changes: 27 additions & 25 deletions apps/api/src/worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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({})
Expand Down Expand Up @@ -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<typeof buildHandler> | undefined
const getHandler = (): ReturnType<typeof buildHandler> => {
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
Expand All @@ -156,7 +148,7 @@ const buildRpcRuntime = async (env: Record<string, unknown>) => {
import("./runtime/mcp-service-graph"),
import("@/platform/DatabasePgLive"),
])
return ManagedRuntime.make(
const runtime = ManagedRuntime.make(
InvestigationServicesLive.pipe(
Layer.provideMerge(WorkerPlatformLive),
Layer.provideMerge(layerPg),
Expand All @@ -165,18 +157,19 @@ const buildRpcRuntime = async (env: Record<string, unknown>) => {
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<typeof buildRpcRuntime> | undefined
const getRpcRuntime = (env: Record<string, unknown>): ReturnType<typeof buildRpcRuntime> => {
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"

Expand Down Expand Up @@ -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<unknown, unknown>
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Loading