Skip to content

Commit 3efffbb

Browse files
RhysSullivanRhysSullivan
authored
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.
1 parent d18b01e commit 3efffbb

6 files changed

Lines changed: 173 additions & 45 deletions

File tree

‎packages/plugins/google-discovery/src/sdk/binding-store.ts‎

Lines changed: 70 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -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 {
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
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‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ import type {
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?: {
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?: {
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?: {
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?: {
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‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -104,12 +104,14 @@ export interface GoogleDiscoverySourceMeta {
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‎

Lines changed: 64 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,25 @@ import { makeInMemoryScopedKv, scopeKv, type Kv, type ToolId, type ScopedKv } fr
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 {
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,18 +165,49 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({
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
// ---------------------------------------------------------------------------
141198
// Factory from global Kv — two scoped sub-namespaces
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‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,35 +11,40 @@ import type {
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‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import {
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?: {
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?: {
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?: {
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?: {
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"

0 commit comments

Comments
 (0)