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 @@
1import { Effect, Schema } from "effect";1import { Effect, Schema } from "effect";
2import { makeInMemoryScopedKv, scopeKv, type Kv, type ScopedKv, type ToolId } from "@executor/sdk";2import { makeInMemoryScopedKv, scopeKv, type Kv, type ScopedKv, type ToolId } from "@executor/sdk";
3 3
4import { GoogleDiscoveryMethodBinding, GoogleDiscoveryStoredSourceData } from "./types";4import {
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
14export const GOOGLE_DISCOVERY_OAUTH_SESSION_TTL_MS = 15 * 60 * 1000;
15
16// ---------------------------------------------------------------------------
17// Stored OAuth session — session payload + expiry, serialized via Schema
18// ---------------------------------------------------------------------------
19
20const StoredOAuthSession = Schema.Struct({
21 session: GoogleDiscoveryOAuthSession,
22 expiresAt: Schema.Number,
23});
24
25const encodeOAuthSession = Schema.encodeSync(Schema.parseJson(StoredOAuthSession));
26const decodeOAuthSession = Schema.decodeUnknownSync(Schema.parseJson(StoredOAuthSession));
5 27
6const StoredBindingEntry = Schema.Struct({28const StoredBindingEntry = Schema.Struct({
7 namespace: Schema.String,29 namespace: Schema.String,
@@ -40,9 +62,22 @@ export interface GoogleDiscoveryBindingStore { · 32 unmodified lines
40 readonly getSourceConfig: (62 readonly getSourceConfig: (
41 namespace: string,63 namespace: string,
42 ) => Effect.Effect<GoogleDiscoveryStoredSourceData | null>;64 ) => 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>;
43}74}
44 75
45const makeStore = (bindings: ScopedKv, sources: ScopedKv): GoogleDiscoveryBindingStore => ({76const makeStore = (
77 bindings: ScopedKv,
78 sources: ScopedKv,
79 oauthSessions: ScopedKv,
80): GoogleDiscoveryBindingStore => ({
46 get: (toolId) =>81 get: (toolId) =>
47 Effect.gen(function* () {82 Effect.gen(function* () {
48 const raw = yield* bindings.get(toolId);83 const raw = yield* bindings.get(toolId);
@@ -108,10 +143,41 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): GoogleDiscoveryBindin · 59 unmodified lines
108 const source = JSON.parse(raw) as GoogleDiscoveryStoredSource;143 const source = JSON.parse(raw) as GoogleDiscoveryStoredSource;
109 return source.config;144 return source.config;
110 }),145 }),
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),
111});173});
112 174
113export const makeKvBindingStore = (kv: Kv, namespace: string): GoogleDiscoveryBindingStore =>175export 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 );
115 181
116export const makeInMemoryBindingStore = (): GoogleDiscoveryBindingStore =>182export 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
31 GoogleDiscoveryAuth,31 GoogleDiscoveryAuth,
32 GoogleDiscoveryManifest,32 GoogleDiscoveryManifest,
33 GoogleDiscoveryManifestMethod,33 GoogleDiscoveryManifestMethod,
34 GoogleDiscoveryOAuthSession,
35 GoogleDiscoveryStoredSourceData,34 GoogleDiscoveryStoredSourceData,
36} from "./types";35} from "./types";
37import { GoogleDiscoveryStoredSourceData as GoogleDiscoveryStoredSourceDataSchema } from "./types";36import { GoogleDiscoveryStoredSourceData as GoogleDiscoveryStoredSourceDataSchema } from "./types";
@@ -262,7 +261,6 @@ export const googleDiscoveryPlugin = (options?: { · 224 unmodified lines
262 readonly bindingStore?: GoogleDiscoveryBindingStore;261 readonly bindingStore?: GoogleDiscoveryBindingStore;
263}): ExecutorPlugin<"googleDiscovery", GoogleDiscoveryPluginExtension> => {262}): ExecutorPlugin<"googleDiscovery", GoogleDiscoveryPluginExtension> => {
264 const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore();263 const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore();
265 const oauthSessions = new Map<string, GoogleDiscoveryOAuthSession>();
266 264
267 return definePlugin({265 return definePlugin({
268 key: "googleDiscovery",266 key: "googleDiscovery",
@@ -439,7 +437,7 @@ export const googleDiscoveryPlugin = (options?: { · 170 unmodified lines
439 }437 }
440 const sessionId = randomUUID();438 const sessionId = randomUUID();
441 const codeVerifier = createPkceCodeVerifier();439 const codeVerifier = createPkceCodeVerifier();
442 oauthSessions.set(sessionId, {440 yield* bindingStore.putOAuthSession(sessionId, {
443 discoveryUrl: normalizeDiscoveryUrl(input.discoveryUrl),441 discoveryUrl: normalizeDiscoveryUrl(input.discoveryUrl),
444 name: input.name,442 name: input.name,
445 clientId: input.clientId,443 clientId: input.clientId,
@@ -463,13 +461,13 @@ export const googleDiscoveryPlugin = (options?: { · 17 unmodified lines
463 461
464 completeOAuth: (input) =>462 completeOAuth: (input) =>
465 Effect.gen(function* () {463 Effect.gen(function* () {
466 const session = oauthSessions.get(input.state);464 const session = yield* bindingStore.getOAuthSession(input.state);
467 if (!session) {465 if (!session) {
468 return yield* new GoogleDiscoveryOAuthError({466 return yield* new GoogleDiscoveryOAuthError({
469 message: "OAuth session not found or has expired",467 message: "OAuth session not found or has expired",
470 });468 });
471 }469 }
472 oauthSessions.delete(input.state);470 yield* bindingStore.deleteOAuthSession(input.state);
473 471
474 if (input.error) {472 if (input.error) {
475 return yield* new GoogleDiscoveryOAuthError({473 return yield* new GoogleDiscoveryOAuthError({
@@ -537,10 +535,7 @@ export const googleDiscoveryPlugin = (options?: { · 61 unmodified lines
537 535
538 return {536 return {
539 extension,537 extension,
540 close: () =>538 close: () => Effect.void,
541 Effect.sync(() => {
542 oauthSessions.clear();
543 }),
544 };539 };
545 }),540 }),
546 });541 });
packages/plugins/google-discovery/src/sdk/types.ts+11−9
@@ -104,12 +104,14 @@ export interface GoogleDiscoverySourceMeta { · 103 unmodified lines
104 readonly name: string;104 readonly name: string;
105}105}
106 106
107export interface GoogleDiscoveryOAuthSession {107/** Pending OAuth session persisted between startOAuth and completeOAuth */
108 readonly discoveryUrl: string;108export const GoogleDiscoveryOAuthSession = Schema.Struct({
109 readonly name: string;109 discoveryUrl: Schema.String,
110 readonly clientId: string;110 name: Schema.String,
111 readonly clientSecretSecretId: string | null;111 clientId: Schema.String,
112 readonly redirectUrl: string;112 clientSecretSecretId: Schema.NullOr(Schema.String),
113 readonly scopes: readonly string[];113 redirectUrl: Schema.String,
114 readonly codeVerifier: string;114 scopes: Schema.Array(Schema.String),
115}115 codeVerifier: Schema.String,
116});
117export 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
7 7
8import { McpToolBinding } from "./types";8import { McpToolBinding } from "./types";
9import type { McpStoredSourceData } from "./types";9import type { McpStoredSourceData } from "./types";
10import { McpOAuthSession } from "./oauth";
11
12// ---------------------------------------------------------------------------
13// OAuth session TTL — pending sessions are cleaned up after this many ms
14// ---------------------------------------------------------------------------
15
16export const MCP_OAUTH_SESSION_TTL_MS = 15 * 60 * 1000;
17
18// ---------------------------------------------------------------------------
19// Stored OAuth session — session payload + expiry, serialized via Schema
20// ---------------------------------------------------------------------------
21
22const StoredOAuthSession = Schema.Struct({
23 session: McpOAuthSession,
24 expiresAt: Schema.Number,
25});
26
27const encodeOAuthSession = Schema.encodeSync(Schema.parseJson(StoredOAuthSession));
28const decodeOAuthSession = Schema.decodeUnknownSync(Schema.parseJson(StoredOAuthSession));
10 29
11// ---------------------------------------------------------------------------30// ---------------------------------------------------------------------------
12// Stored source — combines meta + config into one entry31// Stored source — combines meta + config into one entry
@@ -59,13 +78,24 @@ export interface McpBindingStore { · 46 unmodified lines
59 readonly listSources: () => Effect.Effect<readonly McpStoredSource[]>;78 readonly listSources: () => Effect.Effect<readonly McpStoredSource[]>;
60 readonly getSource: (namespace: string) => Effect.Effect<McpStoredSource | null>;79 readonly getSource: (namespace: string) => Effect.Effect<McpStoredSource | null>;
61 readonly getSourceConfig: (namespace: string) => Effect.Effect<McpStoredSourceData | null>;80 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>;
62}88}
63 89
64// ---------------------------------------------------------------------------90// ---------------------------------------------------------------------------
65// Implementation — two KV namespaces: bindings + sources91// Implementation — two KV namespaces: bindings + sources
66// ---------------------------------------------------------------------------92// ---------------------------------------------------------------------------
67 93
68const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({94const makeStore = (
95 bindings: ScopedKv,
96 sources: ScopedKv,
97 oauthSessions: ScopedKv,
98): McpBindingStore => ({
69 // ---- Bindings ----99 // ---- Bindings ----
70 100
71 get: (toolId) =>101 get: (toolId) =>
@@ -135,6 +165,33 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ · 63 unmodified lines
135 const source = JSON.parse(raw) as McpStoredSource;165 const source = JSON.parse(raw) as McpStoredSource;
136 return source.config;166 return source.config;
137 }),167 }),
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),
138});195});
139 196
140// ---------------------------------------------------------------------------197// ---------------------------------------------------------------------------
@@ -142,11 +199,15 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ · 1 unmodified line
142// ---------------------------------------------------------------------------199// ---------------------------------------------------------------------------
143 200
144export const makeKvBindingStore = (kv: Kv, namespace: string): McpBindingStore =>201export 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 );
146 207
147// ---------------------------------------------------------------------------208// ---------------------------------------------------------------------------
148// In-memory convenience209// In-memory convenience
149// ---------------------------------------------------------------------------210// ---------------------------------------------------------------------------
150 211
151export const makeInMemoryBindingStore = (): McpBindingStore =>212export 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
11 OAuthClientInformationMixed,11 OAuthClientInformationMixed,
12 OAuthTokens,12 OAuthTokens,
13} from "@modelcontextprotocol/sdk/shared/auth.js";13} from "@modelcontextprotocol/sdk/shared/auth.js";
14import { Effect } from "effect";14import { Effect, Schema } from "effect";
15import { McpOAuthError } from "./errors";15import { McpOAuthError } from "./errors";
16 16
17// ---------------------------------------------------------------------------17// ---------------------------------------------------------------------------
18// Types18// Types
19// ---------------------------------------------------------------------------19// ---------------------------------------------------------------------------
20 20
21type JsonObject = { readonly [key: string]: unknown };21const JsonObject = Schema.Record({ key: Schema.String, value: Schema.Unknown });
22type JsonObject = typeof JsonObject.Type;
22 23
23/** Discovery + client state persisted between start and exchange */24/** Discovery + client state persisted between start and exchange */
24export interface McpOAuthDiscoveryState {25export const McpOAuthDiscoveryState = Schema.Struct({
25 readonly resourceMetadataUrl: string | null;26 resourceMetadataUrl: Schema.NullOr(Schema.String),
26 readonly authorizationServerUrl: string | null;27 authorizationServerUrl: Schema.NullOr(Schema.String),
27 readonly resourceMetadata: JsonObject | null;28 resourceMetadata: Schema.NullOr(JsonObject),
28 readonly authorizationServerMetadata: JsonObject | null;29 authorizationServerMetadata: Schema.NullOr(JsonObject),
29 readonly clientInformation: JsonObject | null;30 clientInformation: Schema.NullOr(JsonObject),
30}31});
32export type McpOAuthDiscoveryState = typeof McpOAuthDiscoveryState.Type;
33
34/** Pending OAuth session persisted between startOAuth and completeOAuth */
35export const McpOAuthSession = Schema.Struct({
36 ...McpOAuthDiscoveryState.fields,
37 endpoint: Schema.String,
38 redirectUrl: Schema.String,
39 codeVerifier: Schema.String,
40});
41export type McpOAuthSession = typeof McpOAuthSession.Type;
31 42
32export interface McpOAuthStartResult extends McpOAuthDiscoveryState {43export interface McpOAuthStartResult extends McpOAuthDiscoveryState {
33 readonly authorizationUrl: string;44 readonly authorizationUrl: string;
34 readonly codeVerifier: string;45 readonly codeVerifier: string;
35}46}
36 47
37export interface McpOAuthSession extends McpOAuthDiscoveryState {
38 readonly endpoint: string;
39 readonly redirectUrl: string;
40 readonly codeVerifier: string;
41}
42
43export interface McpOAuthExchangeResult extends McpOAuthDiscoveryState {48export interface McpOAuthExchangeResult extends McpOAuthDiscoveryState {
44 readonly tokens: OAuthTokens;49 readonly tokens: OAuthTokens;
45}50}
packages/plugins/mcp/src/sdk/plugin.ts+4−5
@@ -21,7 +21,7 @@ import { · 20 unmodified lines
21} from "./binding-store";21} from "./binding-store";
22import { createMcpConnector, type McpConnection, type ConnectorInput } from "./connection";22import { createMcpConnector, type McpConnection, type ConnectorInput } from "./connection";
23import { McpConnectionError, McpOAuthError, McpToolDiscoveryError } from "./errors";23import { McpConnectionError, McpOAuthError, McpToolDiscoveryError } from "./errors";
24import { startMcpOAuthAuthorization, exchangeMcpOAuthCode, type McpOAuthSession } from "./oauth";24import { startMcpOAuthAuthorization, exchangeMcpOAuthCode } from "./oauth";
25import { discoverTools } from "./discover";25import { discoverTools } from "./discover";
26import { makeMcpInvoker } from "./invoke";26import { makeMcpInvoker } from "./invoke";
27import { deriveMcpNamespace, joinToolPath, type McpToolManifestEntry } from "./manifest";27import { deriveMcpNamespace, joinToolPath, type McpToolManifestEntry } from "./manifest";
@@ -228,7 +228,6 @@ export const mcpPlugin = (options?: { · 200 unmodified lines
228}): ExecutorPlugin<"mcp", McpPluginExtension> => {228}): ExecutorPlugin<"mcp", McpPluginExtension> => {
229 const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore();229 const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore();
230 const addedSources = new Map<string, Source>();230 const addedSources = new Map<string, Source>();
231 const oauthSessions = new Map<string, McpOAuthSession>();
232 231
233 return definePlugin({232 return definePlugin({
234 key: "mcp",233 key: "mcp",
@@ -626,7 +625,7 @@ export const mcpPlugin = (options?: { · 391 unmodified lines
626 state: sessionId,625 state: sessionId,
627 }).pipe(Effect.mapError((e) => mcpOAuthError(`OAuth start failed: ${e.message}`)));626 }).pipe(Effect.mapError((e) => mcpOAuthError(`OAuth start failed: ${e.message}`)));
628 627
629 oauthSessions.set(sessionId, {628 yield* bindingStore.putOAuthSession(sessionId, {
630 endpoint: fullEndpoint,629 endpoint: fullEndpoint,
631 redirectUrl: input.redirectUrl,630 redirectUrl: input.redirectUrl,
632 codeVerifier: started.codeVerifier,631 codeVerifier: started.codeVerifier,
@@ -648,7 +647,7 @@ export const mcpPlugin = (options?: { · 15 unmodified lines
648 if (input.error) return yield* mcpOAuthError(`OAuth error: ${input.error}`);647 if (input.error) return yield* mcpOAuthError(`OAuth error: ${input.error}`);
649 if (!input.code) return yield* mcpOAuthError("Missing OAuth authorization code");648 if (!input.code) return yield* mcpOAuthError("Missing OAuth authorization code");
650 649
651 const session = oauthSessions.get(input.state);650 const session = yield* bindingStore.getOAuthSession(input.state);
652 if (!session) return yield* mcpOAuthError(`OAuth session not found: ${input.state}`);651 if (!session) return yield* mcpOAuthError(`OAuth session not found: ${input.state}`);
653 652
654 const exchanged = yield* exchangeMcpOAuthCode({653 const exchanged = yield* exchangeMcpOAuthCode({
@@ -686,7 +685,7 @@ export const mcpPlugin = (options?: { · 31 unmodified lines
686 refreshTokenSecretId = ref.id;685 refreshTokenSecretId = ref.id;
687 }686 }
688 687
689 oauthSessions.delete(input.state);688 yield* bindingStore.deleteOAuthSession(input.state);
690 689
691 const expiresAt =690 const expiresAt =
692 typeof exchanged.tokens.expires_in === "number"691 typeof exchanged.tokens.expires_in === "number"