History

fix(mcp,google-discovery): persist pending OAuth sessions via binding store (#221)

* fix(mcp,google-discovery): persist pending OAuth sessions via binding store The mcp and google-discovery plugins kept pending OAuth sessions in an in-process Map on the plugin closure. In cloud, createOrgExecutor runs per-request on Cloudflare Workers, so the redirect-back request built a fresh plugin with an empty Map and completeOAuth failed with "OAuth session not found". Persist sessions through the existing KV-backed binding stores under a new ${namespace}.oauth-sessions scoped namespace, with a 15-minute TTL enforced via an expiresAt field filtered on read. * refactor(mcp,google-discovery): schema-encode persisted OAuth sessions Promote McpOAuthSession and GoogleDiscoveryOAuthSession from plain TS interfaces to Schema.Struct definitions, and swap JSON.parse/stringify with Schema.parseJson in both binding stores so corrupt or schema-drifted KV entries surface a parse error instead of silently flowing through as unchecked any. Matches the StoredBindingEntry pattern already used in the same files.

Rhys Sullivan committed 7 months ago

Commit 3efffbb1 parent d18b01e6 files · +173 −45Browse files

Unified · Split

6 files
  1. packages/plugins/google-discovery/src/sdk/binding-store.ts +70 −4
  2. packages/plugins/google-discovery/src/sdk/plugin.ts +4 −9
  3. packages/plugins/google-discovery/src/sdk/types.ts +11 −9
  4. packages/plugins/mcp/src/sdk/binding-store.ts +64 −3
  5. packages/plugins/mcp/src/sdk/oauth.ts +20 −15
  6. packages/plugins/mcp/src/sdk/plugin.ts +4 −5
packages/plugins/google-discovery/src/sdk/binding-store.ts+70−4
@@ -1,7 +1,29 @@
11 import { Effect, Schema } from "effect";
22 import { makeInMemoryScopedKv, scopeKv, type Kv, type ScopedKv, type ToolId } from "@executor/sdk";
33
4−import { GoogleDiscoveryMethodBinding, GoogleDiscoveryStoredSourceData } from "./types";
4+import {
5+ GoogleDiscoveryMethodBinding,
6+ GoogleDiscoveryOAuthSession,
7+ GoogleDiscoveryStoredSourceData,
8+} from "./types";
9+
10+// ---------------------------------------------------------------------------
11+// OAuth session TTL — pending sessions are cleaned up after this many ms
12+// ---------------------------------------------------------------------------
13+
14+export const GOOGLE_DISCOVERY_OAUTH_SESSION_TTL_MS = 15 * 60 * 1000;
15+
16+// ---------------------------------------------------------------------------
17+// Stored OAuth session — session payload + expiry, serialized via Schema
18+// ---------------------------------------------------------------------------
19+
20+const StoredOAuthSession = Schema.Struct({
21+ session: GoogleDiscoveryOAuthSession,
22+ expiresAt: Schema.Number,
23+});
24+
25+const encodeOAuthSession = Schema.encodeSync(Schema.parseJson(StoredOAuthSession));
26+const decodeOAuthSession = Schema.decodeUnknownSync(Schema.parseJson(StoredOAuthSession));
527
628 const StoredBindingEntry = Schema.Struct({
729 namespace: Schema.String,
@@ -40,9 +62,22 @@ export interface GoogleDiscoveryBindingStore { · 32 unmodified lines
4062 readonly getSourceConfig: (
4163 namespace: string,
4264 ) => Effect.Effect<GoogleDiscoveryStoredSourceData | null>;
65+
66+ readonly putOAuthSession: (
67+ sessionId: string,
68+ session: GoogleDiscoveryOAuthSession,
69+ ) => Effect.Effect<void>;
70+ readonly getOAuthSession: (
71+ sessionId: string,
72+ ) => Effect.Effect<GoogleDiscoveryOAuthSession | null>;
73+ readonly deleteOAuthSession: (sessionId: string) => Effect.Effect<void>;
4374 }
4475
45−const makeStore = (bindings: ScopedKv, sources: ScopedKv): GoogleDiscoveryBindingStore => ({
76+const makeStore = (
77+ bindings: ScopedKv,
78+ sources: ScopedKv,
79+ oauthSessions: ScopedKv,
80+): GoogleDiscoveryBindingStore => ({
4681 get: (toolId) =>
4782 Effect.gen(function* () {
4883 const raw = yield* bindings.get(toolId);
@@ -108,10 +143,41 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): GoogleDiscoveryBindin · 59 unmodified lines
108143 const source = JSON.parse(raw) as GoogleDiscoveryStoredSource;
109144 return source.config;
110145 }),
146+
147+ // ---- Pending OAuth sessions (short-lived, between startOAuth and completeOAuth) ----
148+
149+ putOAuthSession: (sessionId, session) =>
150+ oauthSessions.set([
151+ {
152+ key: sessionId,
153+ value: encodeOAuthSession({
154+ session,
155+ expiresAt: Date.now() + GOOGLE_DISCOVERY_OAUTH_SESSION_TTL_MS,
156+ }),
157+ },
158+ ]),
159+
160+ getOAuthSession: (sessionId) =>
161+ Effect.gen(function* () {
162+ const raw = yield* oauthSessions.get(sessionId);
163+ if (!raw) return null;
164+ const entry = decodeOAuthSession(raw);
165+ if (entry.expiresAt < Date.now()) {
166+ yield* oauthSessions.delete([sessionId]);
167+ return null;
168+ }
169+ return entry.session;
170+ }),
171+
172+ deleteOAuthSession: (sessionId) => oauthSessions.delete([sessionId]).pipe(Effect.asVoid),
111173 });
112174
113175 export const makeKvBindingStore = (kv: Kv, namespace: string): GoogleDiscoveryBindingStore =>
114− makeStore(scopeKv(kv, `${namespace}.bindings`), scopeKv(kv, `${namespace}.sources`));
176+ makeStore(
177+ scopeKv(kv, `${namespace}.bindings`),
178+ scopeKv(kv, `${namespace}.sources`),
179+ scopeKv(kv, `${namespace}.oauth-sessions`),
180+ );
115181
116182 export const makeInMemoryBindingStore = (): GoogleDiscoveryBindingStore =>
117− makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv());
183+ makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv(), makeInMemoryScopedKv());
packages/plugins/google-discovery/src/sdk/plugin.ts+4−9
@@ -31,7 +31,6 @@ import type { · 30 unmodified lines
3131 GoogleDiscoveryAuth,
3232 GoogleDiscoveryManifest,
3333 GoogleDiscoveryManifestMethod,
34− GoogleDiscoveryOAuthSession,
3534 GoogleDiscoveryStoredSourceData,
3635 } from "./types";
3736 import { GoogleDiscoveryStoredSourceData as GoogleDiscoveryStoredSourceDataSchema } from "./types";
@@ -262,7 +261,6 @@ export const googleDiscoveryPlugin = (options?: { · 224 unmodified lines
262261 readonly bindingStore?: GoogleDiscoveryBindingStore;
263262 }): ExecutorPlugin<"googleDiscovery", GoogleDiscoveryPluginExtension> => {
264263 const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore();
265− const oauthSessions = new Map<string, GoogleDiscoveryOAuthSession>();
266264
267265 return definePlugin({
268266 key: "googleDiscovery",
@@ -439,7 +437,7 @@ export const googleDiscoveryPlugin = (options?: { · 170 unmodified lines
439437 }
440438 const sessionId = randomUUID();
441439 const codeVerifier = createPkceCodeVerifier();
442− oauthSessions.set(sessionId, {
440+ yield* bindingStore.putOAuthSession(sessionId, {
443441 discoveryUrl: normalizeDiscoveryUrl(input.discoveryUrl),
444442 name: input.name,
445443 clientId: input.clientId,
@@ -463,13 +461,13 @@ export const googleDiscoveryPlugin = (options?: { · 17 unmodified lines
463461
464462 completeOAuth: (input) =>
465463 Effect.gen(function* () {
466− const session = oauthSessions.get(input.state);
464+ const session = yield* bindingStore.getOAuthSession(input.state);
467465 if (!session) {
468466 return yield* new GoogleDiscoveryOAuthError({
469467 message: "OAuth session not found or has expired",
470468 });
471469 }
472− oauthSessions.delete(input.state);
470+ yield* bindingStore.deleteOAuthSession(input.state);
473471
474472 if (input.error) {
475473 return yield* new GoogleDiscoveryOAuthError({
@@ -537,10 +535,7 @@ export const googleDiscoveryPlugin = (options?: { · 61 unmodified lines
537535
538536 return {
539537 extension,
540− close: () =>
541− Effect.sync(() => {
542− oauthSessions.clear();
543− }),
538+ close: () => Effect.void,
544539 };
545540 }),
546541 });
packages/plugins/google-discovery/src/sdk/types.ts+11−9
@@ -104,12 +104,14 @@ export interface GoogleDiscoverySourceMeta { · 103 unmodified lines
104104 readonly name: string;
105105 }
106106
107−export interface GoogleDiscoveryOAuthSession {
108− readonly discoveryUrl: string;
109− readonly name: string;
110− readonly clientId: string;
111− readonly clientSecretSecretId: string | null;
112− readonly redirectUrl: string;
113− readonly scopes: readonly string[];
114− readonly codeVerifier: string;
115−}
107+/** Pending OAuth session persisted between startOAuth and completeOAuth */
108+export const GoogleDiscoveryOAuthSession = Schema.Struct({
109+ discoveryUrl: Schema.String,
110+ name: Schema.String,
111+ clientId: Schema.String,
112+ clientSecretSecretId: Schema.NullOr(Schema.String),
113+ redirectUrl: Schema.String,
114+ scopes: Schema.Array(Schema.String),
115+ codeVerifier: Schema.String,
116+});
117+export type GoogleDiscoveryOAuthSession = typeof GoogleDiscoveryOAuthSession.Type;
packages/plugins/mcp/src/sdk/binding-store.ts+64−3
@@ -7,6 +7,25 @@ import { makeInMemoryScopedKv, scopeKv, type Kv, type ToolId, type ScopedKv } fr · 6 unmodified lines
77
88 import { McpToolBinding } from "./types";
99 import type { McpStoredSourceData } from "./types";
10+import { McpOAuthSession } from "./oauth";
11+
12+// ---------------------------------------------------------------------------
13+// OAuth session TTL — pending sessions are cleaned up after this many ms
14+// ---------------------------------------------------------------------------
15+
16+export const MCP_OAUTH_SESSION_TTL_MS = 15 * 60 * 1000;
17+
18+// ---------------------------------------------------------------------------
19+// Stored OAuth session — session payload + expiry, serialized via Schema
20+// ---------------------------------------------------------------------------
21+
22+const StoredOAuthSession = Schema.Struct({
23+ session: McpOAuthSession,
24+ expiresAt: Schema.Number,
25+});
26+
27+const encodeOAuthSession = Schema.encodeSync(Schema.parseJson(StoredOAuthSession));
28+const decodeOAuthSession = Schema.decodeUnknownSync(Schema.parseJson(StoredOAuthSession));
1029
1130 // ---------------------------------------------------------------------------
1231 // Stored source — combines meta + config into one entry
@@ -59,13 +78,24 @@ export interface McpBindingStore { · 46 unmodified lines
5978 readonly listSources: () => Effect.Effect<readonly McpStoredSource[]>;
6079 readonly getSource: (namespace: string) => Effect.Effect<McpStoredSource | null>;
6180 readonly getSourceConfig: (namespace: string) => Effect.Effect<McpStoredSourceData | null>;
81+
82+ readonly putOAuthSession: (
83+ sessionId: string,
84+ session: McpOAuthSession,
85+ ) => Effect.Effect<void>;
86+ readonly getOAuthSession: (sessionId: string) => Effect.Effect<McpOAuthSession | null>;
87+ readonly deleteOAuthSession: (sessionId: string) => Effect.Effect<void>;
6288 }
6389
6490 // ---------------------------------------------------------------------------
6591 // Implementation — two KV namespaces: bindings + sources
6692 // ---------------------------------------------------------------------------
6793
68−const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({
94+const makeStore = (
95+ bindings: ScopedKv,
96+ sources: ScopedKv,
97+ oauthSessions: ScopedKv,
98+): McpBindingStore => ({
6999 // ---- Bindings ----
70100
71101 get: (toolId) =>
@@ -135,6 +165,33 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ · 63 unmodified lines
135165 const source = JSON.parse(raw) as McpStoredSource;
136166 return source.config;
137167 }),
168+
169+ // ---- Pending OAuth sessions (short-lived, between startOAuth and completeOAuth) ----
170+
171+ putOAuthSession: (sessionId, session) =>
172+ oauthSessions.set([
173+ {
174+ key: sessionId,
175+ value: encodeOAuthSession({
176+ session,
177+ expiresAt: Date.now() + MCP_OAUTH_SESSION_TTL_MS,
178+ }),
179+ },
180+ ]),
181+
182+ getOAuthSession: (sessionId) =>
183+ Effect.gen(function* () {
184+ const raw = yield* oauthSessions.get(sessionId);
185+ if (!raw) return null;
186+ const entry = decodeOAuthSession(raw);
187+ if (entry.expiresAt < Date.now()) {
188+ yield* oauthSessions.delete([sessionId]);
189+ return null;
190+ }
191+ return entry.session;
192+ }),
193+
194+ deleteOAuthSession: (sessionId) => oauthSessions.delete([sessionId]).pipe(Effect.asVoid),
138195 });
139196
140197 // ---------------------------------------------------------------------------
@@ -142,11 +199,15 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ · 1 unmodified line
142199 // ---------------------------------------------------------------------------
143200
144201 export const makeKvBindingStore = (kv: Kv, namespace: string): McpBindingStore =>
145− makeStore(scopeKv(kv, `${namespace}.bindings`), scopeKv(kv, `${namespace}.sources`));
202+ makeStore(
203+ scopeKv(kv, `${namespace}.bindings`),
204+ scopeKv(kv, `${namespace}.sources`),
205+ scopeKv(kv, `${namespace}.oauth-sessions`),
206+ );
146207
147208 // ---------------------------------------------------------------------------
148209 // In-memory convenience
149210 // ---------------------------------------------------------------------------
150211
151212 export const makeInMemoryBindingStore = (): McpBindingStore =>
152− makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv());
213+ makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv(), makeInMemoryScopedKv());
packages/plugins/mcp/src/sdk/oauth.ts+20−15
@@ -11,35 +11,40 @@ import type { · 10 unmodified lines
1111 OAuthClientInformationMixed,
1212 OAuthTokens,
1313 } from "@modelcontextprotocol/sdk/shared/auth.js";
14−import { Effect } from "effect";
14+import { Effect, Schema } from "effect";
1515 import { McpOAuthError } from "./errors";
1616
1717 // ---------------------------------------------------------------------------
1818 // Types
1919 // ---------------------------------------------------------------------------
2020
21−type JsonObject = { readonly [key: string]: unknown };
21+const JsonObject = Schema.Record({ key: Schema.String, value: Schema.Unknown });
22+type JsonObject = typeof JsonObject.Type;
2223
2324 /** Discovery + client state persisted between start and exchange */
24−export interface McpOAuthDiscoveryState {
25− readonly resourceMetadataUrl: string | null;
26− readonly authorizationServerUrl: string | null;
27− readonly resourceMetadata: JsonObject | null;
28− readonly authorizationServerMetadata: JsonObject | null;
29− readonly clientInformation: JsonObject | null;
30−}
25+export const McpOAuthDiscoveryState = Schema.Struct({
26+ resourceMetadataUrl: Schema.NullOr(Schema.String),
27+ authorizationServerUrl: Schema.NullOr(Schema.String),
28+ resourceMetadata: Schema.NullOr(JsonObject),
29+ authorizationServerMetadata: Schema.NullOr(JsonObject),
30+ clientInformation: Schema.NullOr(JsonObject),
31+});
32+export type McpOAuthDiscoveryState = typeof McpOAuthDiscoveryState.Type;
33+
34+/** Pending OAuth session persisted between startOAuth and completeOAuth */
35+export const McpOAuthSession = Schema.Struct({
36+ ...McpOAuthDiscoveryState.fields,
37+ endpoint: Schema.String,
38+ redirectUrl: Schema.String,
39+ codeVerifier: Schema.String,
40+});
41+export type McpOAuthSession = typeof McpOAuthSession.Type;
3142
3243 export interface McpOAuthStartResult extends McpOAuthDiscoveryState {
3344 readonly authorizationUrl: string;
3445 readonly codeVerifier: string;
3546 }
3647
37−export interface McpOAuthSession extends McpOAuthDiscoveryState {
38− readonly endpoint: string;
39− readonly redirectUrl: string;
40− readonly codeVerifier: string;
41−}
42−
4348 export interface McpOAuthExchangeResult extends McpOAuthDiscoveryState {
4449 readonly tokens: OAuthTokens;
4550 }
packages/plugins/mcp/src/sdk/plugin.ts+4−5
@@ -21,7 +21,7 @@ import { · 20 unmodified lines
2121 } from "./binding-store";
2222 import { createMcpConnector, type McpConnection, type ConnectorInput } from "./connection";
2323 import { McpConnectionError, McpOAuthError, McpToolDiscoveryError } from "./errors";
24−import { startMcpOAuthAuthorization, exchangeMcpOAuthCode, type McpOAuthSession } from "./oauth";
24+import { startMcpOAuthAuthorization, exchangeMcpOAuthCode } from "./oauth";
2525 import { discoverTools } from "./discover";
2626 import { makeMcpInvoker } from "./invoke";
2727 import { deriveMcpNamespace, joinToolPath, type McpToolManifestEntry } from "./manifest";
@@ -228,7 +228,6 @@ export const mcpPlugin = (options?: { · 200 unmodified lines
228228 }): ExecutorPlugin<"mcp", McpPluginExtension> => {
229229 const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore();
230230 const addedSources = new Map<string, Source>();
231− const oauthSessions = new Map<string, McpOAuthSession>();
232231
233232 return definePlugin({
234233 key: "mcp",
@@ -626,7 +625,7 @@ export const mcpPlugin = (options?: { · 391 unmodified lines
626625 state: sessionId,
627626 }).pipe(Effect.mapError((e) => mcpOAuthError(`OAuth start failed: ${e.message}`)));
628627
629− oauthSessions.set(sessionId, {
628+ yield* bindingStore.putOAuthSession(sessionId, {
630629 endpoint: fullEndpoint,
631630 redirectUrl: input.redirectUrl,
632631 codeVerifier: started.codeVerifier,
@@ -648,7 +647,7 @@ export const mcpPlugin = (options?: { · 15 unmodified lines
648647 if (input.error) return yield* mcpOAuthError(`OAuth error: ${input.error}`);
649648 if (!input.code) return yield* mcpOAuthError("Missing OAuth authorization code");
650649
651− const session = oauthSessions.get(input.state);
650+ const session = yield* bindingStore.getOAuthSession(input.state);
652651 if (!session) return yield* mcpOAuthError(`OAuth session not found: ${input.state}`);
653652
654653 const exchanged = yield* exchangeMcpOAuthCode({
@@ -686,7 +685,7 @@ export const mcpPlugin = (options?: { · 31 unmodified lines
686685 refreshTokenSecretId = ref.id;
687686 }
688687
689− oauthSessions.delete(input.state);
688+ yield* bindingStore.deleteOAuthSession(input.state);
690689
691690 const expiresAt =
692691 typeof exchanged.tokens.expires_in === "number"