History

Keep request resources alive for background OAuth sync

Rhys Sullivan committed 1 month ago

Commit fa0215f1 parent 9c424445 files · +404 −15Browse files

Unified · Split

5 files
  1. .changeset/oauth-background-resource-lifetime.md +5 −0
  2. e2e/cloud/oauth-background-catalog.test.ts +178 −0
  3. packages/core/api/src/server/request-scoped.test.ts +144 −0
  4. packages/core/api/src/server/request-scoped.ts +64 −13
  5. packages/core/api/src/server/scoped-executor.ts +13 −2
.changeset/oauth-background-resource-lifetime.md+5−0
@@ -0,0 +1,5 @@
1+---
2+"@executor-js/api": patch
3+---
4+
5+Keep request resources alive until background OAuth tool discovery finishes, so slow cloud connections can publish their tools after the callback returns.
e2e/cloud/oauth-background-catalog.test.ts+178−0
@@ -0,0 +1,178 @@
1+import { randomBytes } from "node:crypto";
2+
3+import { expect } from "@effect/vitest";
4+import { Effect, Option, Schema } from "effect";
5+import { HttpServerRequest, HttpServerResponse } from "effect/unstable/http";
6+import { composePluginApi } from "@executor-js/api/server";
7+import { mcpHttpPlugin } from "@executor-js/plugin-mcp/api";
8+import {
9+ AuthTemplateSlug,
10+ ConnectionName,
11+ IntegrationSlug,
12+ OAuthClientSlug,
13+} from "@executor-js/sdk/shared";
14+import { serveTestHttpApp } from "@executor-js/sdk/testing";
15+
16+import { createEmulatorInstance } from "../src/emulator-instance";
17+import { scenario } from "../src/scenario";
18+import { Api, Browser, Target } from "../src/services";
19+import { visit } from "../src/surfaces/browser";
20+
21+const api = composePluginApi([mcpHttpPlugin()] as const);
22+const decodeRpc = Schema.decodeUnknownOption(
23+ Schema.fromJsonString(Schema.Struct({ method: Schema.String })),
24+);
25+
26+scenario(
27+ "OAuth · slow catalog discovery persists after the cloud callback returns",
28+ { timeout: 180_000 },
29+ Effect.scoped(
30+ Effect.gen(function* () {
31+ const target = yield* Target;
32+ const browser = yield* Browser;
33+ const { client: makeApiClient } = yield* Api;
34+ const identity = yield* target.newIdentity();
35+ const client = yield* makeApiClient(api, identity);
36+ const upstream = yield* createEmulatorInstance("mcp", "oauth-background-catalog");
37+ const firstListing = Promise.withResolvers<void>();
38+ const releaseCallbackListing = Promise.withResolvers<void>();
39+ const releaseOtherListings = Promise.withResolvers<void>();
40+ let listings = 0;
41+
42+ // Only delay traffic. OAuth, credentials and MCP responses all come from
43+ // the published emulator. Hold later listings separately so a UI/API read
44+ // cannot repair the catalog and hide failure of the callback's own sync.
45+ const proxy = yield* serveTestHttpApp((request) =>
46+ Effect.promise(async () => {
47+ const web = await Effect.runPromise(HttpServerRequest.toWeb(request));
48+ const body = web.method === "GET" || web.method === "HEAD" ? undefined : await web.text();
49+ const path = new URL(web.url);
50+ const headers = new Headers(web.headers);
51+ headers.delete("host");
52+ const response = await fetch(`${upstream}${path.pathname}${path.search}`, {
53+ method: web.method,
54+ headers,
55+ body,
56+ redirect: "manual",
57+ });
58+ const rpc = body === undefined ? Option.none() : decodeRpc(body);
59+ if (response.ok && Option.isSome(rpc) && rpc.value.method === "tools/list") {
60+ listings += 1;
61+ if (listings === 1) {
62+ firstListing.resolve();
63+ await releaseCallbackListing.promise;
64+ } else {
65+ await releaseOtherListings.promise;
66+ }
67+ }
68+ return HttpServerResponse.fromWeb(response);
69+ }),
70+ );
71+ yield* Effect.addFinalizer(() =>
72+ Effect.sync(() => {
73+ releaseCallbackListing.resolve();
74+ releaseOtherListings.resolve();
75+ }),
76+ );
77+
78+ const slug = IntegrationSlug.make(`slow_oauth_${randomBytes(4).toString("hex")}`);
79+ yield* client.mcp.addServer({
80+ payload: {
81+ transport: "remote",
82+ name: "Slow OAuth catalog",
83+ endpoint: `${proxy.baseUrl}/mcp`,
84+ slug,
85+ authenticationTemplate: [{ kind: "oauth2" }],
86+ },
87+ });
88+ yield* Effect.addFinalizer(() =>
89+ Effect.gen(function* () {
90+ releaseCallbackListing.resolve();
91+ releaseOtherListings.resolve();
92+ yield* client.mcp.removeServer({ params: { slug } }).pipe(Effect.ignore);
93+ }),
94+ );
95+
96+ const probe = yield* client.oauth.probe({ payload: { url: `${proxy.baseUrl}/mcp` } });
97+ if (!probe.registrationEndpoint || !probe.authorizationUrl || !probe.tokenUrl) {
98+ return yield* Effect.die("Emulator did not advertise OAuth registration");
99+ }
100+ const { client: oauthClient } = yield* client.oauth.registerDynamic({
101+ payload: {
102+ owner: "org",
103+ slug: OAuthClientSlug.make(`${slug}_client`),
104+ registrationEndpoint: probe.registrationEndpoint,
105+ authorizationUrl: probe.authorizationUrl,
106+ tokenUrl: probe.tokenUrl,
107+ resource: probe.resource,
108+ scopes: probe.scopesSupported ?? [],
109+ originIntegration: slug,
110+ },
111+ });
112+ yield* Effect.addFinalizer(() =>
113+ client.oauth
114+ .removeClient({
115+ params: { slug: oauthClient },
116+ payload: { owner: "org" },
117+ })
118+ .pipe(Effect.ignore),
119+ );
120+ const started = yield* client.oauth.start({
121+ payload: {
122+ owner: "org",
123+ client: oauthClient,
124+ clientOwner: "org",
125+ name: ConnectionName.make("main"),
126+ integration: slug,
127+ template: AuthTemplateSlug.make("oauth2"),
128+ },
129+ });
130+ if (started.status !== "redirect") return yield* Effect.die("Expected OAuth authorization");
131+ yield* Effect.addFinalizer(() =>
132+ client.oauth.cancel({ payload: { state: started.state } }).pipe(Effect.ignore),
133+ );
134+
135+ yield* browser.session(identity, async ({ page, step }) => {
136+ await step("Authorize the OAuth connection", async () => {
137+ // Stay off the integration screen until discovery is verified: its
138+ // refresh after the callback could otherwise repair a failed sync.
139+ await page.goto(started.authorizationUrl);
140+ const authorization = new URL(page.url());
141+ const approved = await page.request.post(`${upstream}/authorize/approve`, {
142+ form: { ...Object.fromEntries(authorization.searchParams), login: "admin" },
143+ maxRedirects: 0,
144+ });
145+ expect(approved.status()).toBe(302);
146+ const callback = approved.headers()["location"];
147+ if (!callback) throw new Error("Emulator did not return the OAuth callback");
148+ await page.goto(callback, { waitUntil: "domcontentloaded" });
149+ await page.getByRole("heading", { name: "Connected" }).waitFor();
150+ });
151+ await step("Keep discovery blocked after the callback returns", async () => {
152+ await expect.poll(() => listings, { timeout: 15_000 }).toBeGreaterThan(0);
153+ await firstListing.promise;
154+ const connections = await Effect.runPromise(
155+ client.connections.list({ query: { integration: slug } }),
156+ );
157+ expect(connections).toHaveLength(1);
158+ });
159+ await step("Receive the tools from the completed background discovery", async () => {
160+ releaseCallbackListing.resolve();
161+ await expect
162+ .poll(
163+ async () => {
164+ const tools = await Effect.runPromise(
165+ client.tools.list({ query: { integration: slug } }),
166+ );
167+ return tools.length;
168+ },
169+ { timeout: 30_000, interval: 500 },
170+ )
171+ .toBeGreaterThan(0);
172+ await visit(page, `/integrations/${slug}`);
173+ await page.getByRole("button", { name: "Add connection" }).waitFor();
174+ });
175+ });
176+ }),
177+ ),
178+);
packages/core/api/src/server/request-scoped.test.ts+144−0
@@ -0,0 +1,144 @@
1+import { describe, expect, it, onTestFinished } from "@effect/vitest";
2+import { Context, Effect, Layer } from "effect";
3+import { HttpRouter, HttpServer, HttpServerResponse } from "effect/unstable/http";
4+
5+import { RequestBackgroundTasks, requestScopedMiddleware } from "./request-scoped";
6+
7+class Resource extends Context.Service<Resource, { readonly id: number; closed: boolean }>()(
8+ "test/RequestResource",
9+) {}
10+
11+const fixture = (
12+ options: {
13+ background?: boolean;
14+ failTask?: boolean;
15+ failRequest?: boolean;
16+ blockRequest?: boolean;
17+ } = {},
18+) => {
19+ const resources: Resource["Service"][] = [];
20+ const gates: ReturnType<typeof Promise.withResolvers<void>>[] = [];
21+ const keptAlive: Promise<void>[] = [];
22+ const writes: number[] = [];
23+ const entered = Promise.withResolvers<void>();
24+ const resource = Layer.effect(Resource)(
25+ Effect.acquireRelease(
26+ Effect.sync(() => {
27+ const acquired = { id: resources.length, closed: false };
28+ resources.push(acquired);
29+ gates.push(Promise.withResolvers<void>());
30+ return acquired;
31+ }),
32+ (acquired) =>
33+ Effect.sync(() => {
34+ acquired.closed = true;
35+ }),
36+ ),
37+ );
38+ const routes = HttpRouter.add(
39+ "GET",
40+ "/",
41+ Effect.gen(function* () {
42+ const acquired = yield* Resource;
43+ if (options.background !== false) {
44+ const tasks = yield* RequestBackgroundTasks;
45+ const gate = gates[acquired.id];
46+ if (!gate) return yield* Effect.die("Missing request gate");
47+ const task = Effect.runPromise(
48+ Effect.gen(function* () {
49+ yield* Effect.promise(() => gate.promise);
50+ if (options.failTask) return yield* Effect.fail("Background failure");
51+ if (acquired.closed)
52+ return yield* Effect.fail("Resource closed before background write");
53+ writes.push(acquired.id);
54+ }),
55+ );
56+ keptAlive.push(tasks.retain(task));
57+ }
58+ entered.resolve();
59+ if (options.blockRequest) return yield* Effect.never;
60+ if (options.failRequest) return yield* Effect.die("Request failure");
61+ return HttpServerResponse.empty();
62+ }),
63+ );
64+ const app = HttpRouter.toWebHandler(
65+ routes.pipe(
66+ Layer.provide(requestScopedMiddleware(resource).layer),
67+ Layer.provideMerge(HttpServer.layerServices),
68+ ),
69+ { disableLogger: true },
70+ );
71+ onTestFinished(async () => {
72+ for (const gate of gates) gate.resolve();
73+ await Promise.all(keptAlive);
74+ await app.dispose();
75+ });
76+ return {
77+ resources,
78+ gates,
79+ keptAlive,
80+ writes,
81+ entered,
82+ request: (signal?: AbortSignal) =>
83+ app.handler(new Request("http://test.local/", { signal }), Context.empty()),
84+ };
85+};
86+
87+describe("request background resource ownership", () => {
88+ it("returns the response before background work, and releases after its write", async () => {
89+ const test = fixture();
90+ expect((await test.request()).status).toBe(204);
91+ expect(test.resources.map((item) => item.closed)).toEqual([false]);
92+ expect(test.writes).toEqual([]);
93+ test.gates[0]?.resolve();
94+ await Promise.all(test.keptAlive);
95+ expect(test.writes).toEqual([0]);
96+ expect(test.resources.map((item) => item.closed)).toEqual([true]);
97+ });
98+
99+ it("keeps concurrent requests isolated and releases each independently", async () => {
100+ const test = fixture();
101+ const responses = await Promise.all([test.request(), test.request()]);
102+ expect(responses.map((response) => response.status)).toEqual([204, 204]);
103+ test.gates[0]?.resolve();
104+ await test.keptAlive[0];
105+ expect(test.resources.map((item) => item.closed)).toEqual([true, false]);
106+ test.gates[1]?.resolve();
107+ await test.keptAlive[1];
108+ expect(test.writes).toEqual([0, 1]);
109+ expect(test.resources.map((item) => item.closed)).toEqual([true, true]);
110+ });
111+
112+ for (const failure of ["task", "request"] as const) {
113+ it(`releases resources after a ${failure} failure`, async () => {
114+ const test = fixture({ failTask: failure === "task", failRequest: failure === "request" });
115+ expect((await test.request()).status).toBe(failure === "request" ? 500 : 204);
116+ expect(test.resources.map((item) => item.closed)).toEqual([false]);
117+ test.gates[0]?.resolve();
118+ await Promise.all(test.keptAlive);
119+ expect(test.writes).toEqual(failure === "task" ? [] : [0]);
120+ expect(test.resources.map((item) => item.closed)).toEqual([true]);
121+ });
122+ }
123+
124+ it("closes before responding when no background work was registered", async () => {
125+ const test = fixture({ background: false });
126+ expect((await test.request()).status).toBe(204);
127+ expect(test.resources.map((item) => item.closed)).toEqual([true]);
128+ expect(test.keptAlive).toEqual([]);
129+ });
130+
131+ it("finishes retained work and cleanup after the client cancels the request", async () => {
132+ const test = fixture({ blockRequest: true });
133+ const controller = new AbortController();
134+ const response = test.request(controller.signal);
135+ await test.entered.promise;
136+ controller.abort();
137+ expect((await response).status).toBe(499);
138+ expect(test.resources.map((item) => item.closed)).toEqual([false]);
139+ test.gates[0]?.resolve();
140+ await Promise.all(test.keptAlive);
141+ expect(test.writes).toEqual([0]);
142+ expect(test.resources.map((item) => item.closed)).toEqual([true]);
143+ });
144+});
packages/core/api/src/server/request-scoped.ts+64−13
@@ -13,8 +13,9 @@ · 12 unmodified lines
1313 // both build the inner layer at construction time. The only primitive
1414 // that actually rebuilds per request is a router middleware whose
1515 // per-request handler builds the layer with a *fresh* `MemoMap` and a
16−// per-request scope, so `acquireRelease` fires per request and finalizers
17−// run when the request fiber's scope closes.
16+// per-request scope, so `acquireRelease` fires per request. Background work
17+// retains that scope until its database writes finish; the response need not
18+// wait, but the platform keep-alive must include resource cleanup too.
1819 //
1920 // The fresh `MemoMap` matters: `Layer.build` would otherwise inherit
2021 // `CurrentMemoMap` from the boot context (`HttpRouter.toWebHandler`
@@ -29,14 +30,24 @@ · 8 unmodified lines
2930 // coverage that pins this rule down (sequential AND concurrent cases).
3031 // ---------------------------------------------------------------------------
3132
32−import { Effect, Layer } from "effect";
33+import { Context, Effect, Layer, Scope } from "effect";
3334 import { HttpRouter } from "effect/unstable/http";
3435
36+/**
37+ * Retain this request's resources for already-started background work. The
38+ * caller owns task error reporting; the returned promise settles only after
39+ * all retained tasks AND resource finalizers finish, for the host's waitUntil.
40+ */
41+export class RequestBackgroundTasks extends Context.Service<
42+ RequestBackgroundTasks,
43+ { readonly retain: (task: Promise<unknown>) => Promise<void> }
44+>()("@executor-js/api/RequestBackgroundTasks") {}
45+
3546 /**
3647 * Build an `HttpRouter.middleware` that provides `layer`'s services to
3748 * each request. The layer is rebuilt per HTTP request so
38− * `Effect.acquireRelease` fires per request and is released when the
39− * request fiber's scope closes.
49+ * `Effect.acquireRelease` fires per request. Resources close after the handler
50+ * and its registered background work settle, without delaying the response.
4051 *
4152 * The returned value is a `Middleware`. Use `.layer` to apply it as a
4253 * standalone layer; use `.combine(other)` to fold it into another
@@ -46,15 +57,55 @@ import { HttpRouter } from "effect/unstable/http"; · 3 unmodified lines
4657 * outer middleware's `requires`).
4758 */
4859 export const requestScopedMiddleware = <A>(layer: Layer.Layer<A>) =>
49− HttpRouter.middleware<{ provides: A }>()((httpEffect) =>
50− Effect.scoped(
60+ HttpRouter.middleware<{ provides: A | RequestBackgroundTasks }>()((httpEffect) =>
61+ Effect.uninterruptibleMask((restore) =>
5162 Effect.gen(function* () {
52− // Fresh MemoMap per request — see file-level note for why we
53− // must NOT inherit `CurrentMemoMap` from the boot context.
54− const memoMap = yield* Layer.makeMemoMap;
55− const scope = yield* Effect.scope;
56− const services = yield* Layer.buildWithMemoMap(layer, memoMap, scope);
57− return yield* Effect.provideContext(httpEffect, services);
63+ const scope = yield* Scope.make();
64+ const pending = new Set<Promise<void>>();
65+ const released = Promise.withResolvers<void>();
66+ const background = RequestBackgroundTasks.of({
67+ retain: (task) => {
68+ // SDK tasks report their own failures. Both outcomes release the
69+ // resource lease; a rejected task must not leak its database.
70+ const settled = task.then(
71+ () => {
72+ pending.delete(settled);
73+ },
74+ () => {
75+ pending.delete(settled);
76+ },
77+ );
78+ pending.add(settled);
79+ return released.promise;
80+ },
81+ });
82+ return yield* restore(
83+ Effect.gen(function* () {
84+ // Never inherit the boot MemoMap: concurrent requests each own
85+ // their socket, including after either response has been sent.
86+ const memoMap = yield* Layer.makeMemoMap;
87+ const services = yield* Layer.buildWithMemoMap(layer, memoMap, scope);
88+ return yield* Effect.provideContext(httpEffect, services);
89+ }).pipe(
90+ Effect.provideService(Scope.Scope, scope),
91+ Effect.provideService(RequestBackgroundTasks, background),
92+ ),
93+ ).pipe(
94+ Effect.onExit((exit) => {
95+ const release = Effect.gen(function* () {
96+ // A retained task may start another task before it settles.
97+ while (pending.size > 0) {
98+ yield* Effect.promise(() => Promise.all(pending));
99+ }
100+ }).pipe(
101+ Effect.ensuring(Scope.close(scope, exit)),
102+ Effect.ensuring(Effect.sync(() => released.resolve())),
103+ );
104+ // With no background work, preserve synchronous teardown. Otherwise
105+ // the resource owner, including cleanup, is kept alive by the host.
106+ return pending.size === 0 ? release : release.pipe(Effect.forkDetach, Effect.asVoid);
107+ }),
108+ );
58109 }),
59110 ),
60111 );
packages/core/api/src/server/scoped-executor.ts+13−2
@@ -51,6 +51,7 @@ import { · 50 unmodified lines
5151 } from "@executor-js/sdk/host-internal";
5252
5353 import { DbProvider } from "./executor-fuma-db";
54+import { RequestBackgroundTasks } from "./request-scoped";
5455
5556 // ---------------------------------------------------------------------------
5657 // HostConfig seam — the two host scalars that vary the `createExecutor` options.
@@ -126,12 +127,14 @@ export interface HostConfigShape { · 69 unmodified lines
126127 */
127128 readonly toolsSyncTtlMs?: number | null;
128129 /**
129− * Forwarded verbatim to `ExecutorConfig.waitUntil`: the host's keep-alive
130+ * Forwarded to `ExecutorConfig.waitUntil`: the host's keep-alive
130131 * for background work that outlives a request (stale tool-catalog rebuilds
131132 * that keep running after a read stops waiting). Cloud supplies the
132133 * platform `waitUntil` from `cloudflare:workers`, which binds to the
133134 * in-flight invocation ambiently; long-lived hosts (self-host, local,
134135 * tests) omit it and detached fibers simply run to completion in-process.
136+ * Under requestScopedMiddleware, the promise also covers releasing that
137+ * request's database after background work finishes.
135138 */
136139 readonly waitUntil?: (promise: Promise<unknown>) => void;
137140 }
@@ -267,6 +270,14 @@ export const makeScopedExecutor = < · 129 unmodified lines
267270 const { db, blobs } = yield* DbProvider.asEffect();
268271 const { plugins: pluginsFactory } = yield* PluginsProvider.asEffect();
269272 const config = yield* HostConfig.asEffect();
273+ const background = yield* Effect.serviceOption(RequestBackgroundTasks);
274+ const waitUntil = Option.match(background, {
275+ onNone: () => config.waitUntil,
276+ onSome: (tasks) => (task: Promise<unknown>) => {
277+ const released = tasks.retain(task);
278+ config.waitUntil?.(released);
279+ },
280+ });
270281 // Explicit config wins; otherwise fall back to the request origin if a host
271282 // provided one (HTTP middleware / MCP session DO). Stays `undefined` for
272283 // non-request callers — `coreTools.webBaseUrl` is optional and only the
@@ -323,7 +334,7 @@ export const makeScopedExecutor = < · 50 unmodified lines
323334 fetch: hostedFetch,
324335 onIntegrationChange: config.onIntegrationChange,
325336 ...(config.toolsSyncTtlMs !== undefined ? { toolsSyncTtlMs: config.toolsSyncTtlMs } : {}),
326− ...(config.waitUntil !== undefined ? { waitUntil: config.waitUntil } : {}),
337+ ...(waitUntil !== undefined ? { waitUntil } : {}),
327338 onElicitation: "accept-all",
328339 ...(options?.orgWrites === undefined ? {} : { orgWrites: options.orgWrites }),
329340 redirectUri,