diff --git a/packages/plugins/google-discovery/src/sdk/binding-store.ts b/packages/plugins/google-discovery/src/sdk/binding-store.ts index ff3e4d1edd..1c216919fd 100644 --- a/packages/plugins/google-discovery/src/sdk/binding-store.ts +++ b/packages/plugins/google-discovery/src/sdk/binding-store.ts @@ -1,7 +1,29 @@ import { Effect, Schema } from "effect"; import { makeInMemoryScopedKv, scopeKv, type Kv, type ScopedKv, type ToolId } from "@executor/sdk"; -import { GoogleDiscoveryMethodBinding, GoogleDiscoveryStoredSourceData } from "./types"; +import { + GoogleDiscoveryMethodBinding, + GoogleDiscoveryOAuthSession, + GoogleDiscoveryStoredSourceData, +} from "./types"; + +// --------------------------------------------------------------------------- +// OAuth session TTL — pending sessions are cleaned up after this many ms +// --------------------------------------------------------------------------- + +export const GOOGLE_DISCOVERY_OAUTH_SESSION_TTL_MS = 15 * 60 * 1000; + +// --------------------------------------------------------------------------- +// Stored OAuth session — session payload + expiry, serialized via Schema +// --------------------------------------------------------------------------- + +const StoredOAuthSession = Schema.Struct({ + session: GoogleDiscoveryOAuthSession, + expiresAt: Schema.Number, +}); + +const encodeOAuthSession = Schema.encodeSync(Schema.parseJson(StoredOAuthSession)); +const decodeOAuthSession = Schema.decodeUnknownSync(Schema.parseJson(StoredOAuthSession)); const StoredBindingEntry = Schema.Struct({ namespace: Schema.String, @@ -40,9 +62,22 @@ export interface GoogleDiscoveryBindingStore { readonly getSourceConfig: ( namespace: string, ) => Effect.Effect; + + readonly putOAuthSession: ( + sessionId: string, + session: GoogleDiscoveryOAuthSession, + ) => Effect.Effect; + readonly getOAuthSession: ( + sessionId: string, + ) => Effect.Effect; + readonly deleteOAuthSession: (sessionId: string) => Effect.Effect; } -const makeStore = (bindings: ScopedKv, sources: ScopedKv): GoogleDiscoveryBindingStore => ({ +const makeStore = ( + bindings: ScopedKv, + sources: ScopedKv, + oauthSessions: ScopedKv, +): GoogleDiscoveryBindingStore => ({ get: (toolId) => Effect.gen(function* () { const raw = yield* bindings.get(toolId); @@ -108,10 +143,41 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): GoogleDiscoveryBindin const source = JSON.parse(raw) as GoogleDiscoveryStoredSource; return source.config; }), + + // ---- Pending OAuth sessions (short-lived, between startOAuth and completeOAuth) ---- + + putOAuthSession: (sessionId, session) => + oauthSessions.set([ + { + key: sessionId, + value: encodeOAuthSession({ + session, + expiresAt: Date.now() + GOOGLE_DISCOVERY_OAUTH_SESSION_TTL_MS, + }), + }, + ]), + + getOAuthSession: (sessionId) => + Effect.gen(function* () { + const raw = yield* oauthSessions.get(sessionId); + if (!raw) return null; + const entry = decodeOAuthSession(raw); + if (entry.expiresAt < Date.now()) { + yield* oauthSessions.delete([sessionId]); + return null; + } + return entry.session; + }), + + deleteOAuthSession: (sessionId) => oauthSessions.delete([sessionId]).pipe(Effect.asVoid), }); export const makeKvBindingStore = (kv: Kv, namespace: string): GoogleDiscoveryBindingStore => - makeStore(scopeKv(kv, `${namespace}.bindings`), scopeKv(kv, `${namespace}.sources`)); + makeStore( + scopeKv(kv, `${namespace}.bindings`), + scopeKv(kv, `${namespace}.sources`), + scopeKv(kv, `${namespace}.oauth-sessions`), + ); export const makeInMemoryBindingStore = (): GoogleDiscoveryBindingStore => - makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv()); + makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv(), makeInMemoryScopedKv()); diff --git a/packages/plugins/google-discovery/src/sdk/plugin.ts b/packages/plugins/google-discovery/src/sdk/plugin.ts index a8ac0b4b6a..cafad3c58a 100644 --- a/packages/plugins/google-discovery/src/sdk/plugin.ts +++ b/packages/plugins/google-discovery/src/sdk/plugin.ts @@ -31,7 +31,6 @@ import type { GoogleDiscoveryAuth, GoogleDiscoveryManifest, GoogleDiscoveryManifestMethod, - GoogleDiscoveryOAuthSession, GoogleDiscoveryStoredSourceData, } from "./types"; import { GoogleDiscoveryStoredSourceData as GoogleDiscoveryStoredSourceDataSchema } from "./types"; @@ -262,7 +261,6 @@ export const googleDiscoveryPlugin = (options?: { readonly bindingStore?: GoogleDiscoveryBindingStore; }): ExecutorPlugin<"googleDiscovery", GoogleDiscoveryPluginExtension> => { const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore(); - const oauthSessions = new Map(); return definePlugin({ key: "googleDiscovery", @@ -439,7 +437,7 @@ export const googleDiscoveryPlugin = (options?: { } const sessionId = randomUUID(); const codeVerifier = createPkceCodeVerifier(); - oauthSessions.set(sessionId, { + yield* bindingStore.putOAuthSession(sessionId, { discoveryUrl: normalizeDiscoveryUrl(input.discoveryUrl), name: input.name, clientId: input.clientId, @@ -463,13 +461,13 @@ export const googleDiscoveryPlugin = (options?: { completeOAuth: (input) => Effect.gen(function* () { - const session = oauthSessions.get(input.state); + const session = yield* bindingStore.getOAuthSession(input.state); if (!session) { return yield* new GoogleDiscoveryOAuthError({ message: "OAuth session not found or has expired", }); } - oauthSessions.delete(input.state); + yield* bindingStore.deleteOAuthSession(input.state); if (input.error) { return yield* new GoogleDiscoveryOAuthError({ @@ -537,10 +535,7 @@ export const googleDiscoveryPlugin = (options?: { return { extension, - close: () => - Effect.sync(() => { - oauthSessions.clear(); - }), + close: () => Effect.void, }; }), }); diff --git a/packages/plugins/google-discovery/src/sdk/types.ts b/packages/plugins/google-discovery/src/sdk/types.ts index 3787be0c82..19f746340c 100644 --- a/packages/plugins/google-discovery/src/sdk/types.ts +++ b/packages/plugins/google-discovery/src/sdk/types.ts @@ -104,12 +104,14 @@ export interface GoogleDiscoverySourceMeta { readonly name: string; } -export interface GoogleDiscoveryOAuthSession { - readonly discoveryUrl: string; - readonly name: string; - readonly clientId: string; - readonly clientSecretSecretId: string | null; - readonly redirectUrl: string; - readonly scopes: readonly string[]; - readonly codeVerifier: string; -} +/** Pending OAuth session persisted between startOAuth and completeOAuth */ +export const GoogleDiscoveryOAuthSession = Schema.Struct({ + discoveryUrl: Schema.String, + name: Schema.String, + clientId: Schema.String, + clientSecretSecretId: Schema.NullOr(Schema.String), + redirectUrl: Schema.String, + scopes: Schema.Array(Schema.String), + codeVerifier: Schema.String, +}); +export type GoogleDiscoveryOAuthSession = typeof GoogleDiscoveryOAuthSession.Type; diff --git a/packages/plugins/mcp/src/sdk/binding-store.ts b/packages/plugins/mcp/src/sdk/binding-store.ts index dea141c345..567d116efd 100644 --- a/packages/plugins/mcp/src/sdk/binding-store.ts +++ b/packages/plugins/mcp/src/sdk/binding-store.ts @@ -7,6 +7,25 @@ import { makeInMemoryScopedKv, scopeKv, type Kv, type ToolId, type ScopedKv } fr import { McpToolBinding } from "./types"; import type { McpStoredSourceData } from "./types"; +import { McpOAuthSession } from "./oauth"; + +// --------------------------------------------------------------------------- +// OAuth session TTL — pending sessions are cleaned up after this many ms +// --------------------------------------------------------------------------- + +export const MCP_OAUTH_SESSION_TTL_MS = 15 * 60 * 1000; + +// --------------------------------------------------------------------------- +// Stored OAuth session — session payload + expiry, serialized via Schema +// --------------------------------------------------------------------------- + +const StoredOAuthSession = Schema.Struct({ + session: McpOAuthSession, + expiresAt: Schema.Number, +}); + +const encodeOAuthSession = Schema.encodeSync(Schema.parseJson(StoredOAuthSession)); +const decodeOAuthSession = Schema.decodeUnknownSync(Schema.parseJson(StoredOAuthSession)); // --------------------------------------------------------------------------- // Stored source — combines meta + config into one entry @@ -59,13 +78,24 @@ export interface McpBindingStore { readonly listSources: () => Effect.Effect; readonly getSource: (namespace: string) => Effect.Effect; readonly getSourceConfig: (namespace: string) => Effect.Effect; + + readonly putOAuthSession: ( + sessionId: string, + session: McpOAuthSession, + ) => Effect.Effect; + readonly getOAuthSession: (sessionId: string) => Effect.Effect; + readonly deleteOAuthSession: (sessionId: string) => Effect.Effect; } // --------------------------------------------------------------------------- // Implementation — two KV namespaces: bindings + sources // --------------------------------------------------------------------------- -const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ +const makeStore = ( + bindings: ScopedKv, + sources: ScopedKv, + oauthSessions: ScopedKv, +): McpBindingStore => ({ // ---- Bindings ---- get: (toolId) => @@ -135,6 +165,33 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ const source = JSON.parse(raw) as McpStoredSource; return source.config; }), + + // ---- Pending OAuth sessions (short-lived, between startOAuth and completeOAuth) ---- + + putOAuthSession: (sessionId, session) => + oauthSessions.set([ + { + key: sessionId, + value: encodeOAuthSession({ + session, + expiresAt: Date.now() + MCP_OAUTH_SESSION_TTL_MS, + }), + }, + ]), + + getOAuthSession: (sessionId) => + Effect.gen(function* () { + const raw = yield* oauthSessions.get(sessionId); + if (!raw) return null; + const entry = decodeOAuthSession(raw); + if (entry.expiresAt < Date.now()) { + yield* oauthSessions.delete([sessionId]); + return null; + } + return entry.session; + }), + + deleteOAuthSession: (sessionId) => oauthSessions.delete([sessionId]).pipe(Effect.asVoid), }); // --------------------------------------------------------------------------- @@ -142,11 +199,15 @@ const makeStore = (bindings: ScopedKv, sources: ScopedKv): McpBindingStore => ({ // --------------------------------------------------------------------------- export const makeKvBindingStore = (kv: Kv, namespace: string): McpBindingStore => - makeStore(scopeKv(kv, `${namespace}.bindings`), scopeKv(kv, `${namespace}.sources`)); + makeStore( + scopeKv(kv, `${namespace}.bindings`), + scopeKv(kv, `${namespace}.sources`), + scopeKv(kv, `${namespace}.oauth-sessions`), + ); // --------------------------------------------------------------------------- // In-memory convenience // --------------------------------------------------------------------------- export const makeInMemoryBindingStore = (): McpBindingStore => - makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv()); + makeStore(makeInMemoryScopedKv(), makeInMemoryScopedKv(), makeInMemoryScopedKv()); diff --git a/packages/plugins/mcp/src/sdk/oauth.ts b/packages/plugins/mcp/src/sdk/oauth.ts index 5f85e8e6ee..92ac7b78d1 100644 --- a/packages/plugins/mcp/src/sdk/oauth.ts +++ b/packages/plugins/mcp/src/sdk/oauth.ts @@ -11,35 +11,40 @@ import type { OAuthClientInformationMixed, OAuthTokens, } from "@modelcontextprotocol/sdk/shared/auth.js"; -import { Effect } from "effect"; +import { Effect, Schema } from "effect"; import { McpOAuthError } from "./errors"; // --------------------------------------------------------------------------- // Types // --------------------------------------------------------------------------- -type JsonObject = { readonly [key: string]: unknown }; +const JsonObject = Schema.Record({ key: Schema.String, value: Schema.Unknown }); +type JsonObject = typeof JsonObject.Type; /** Discovery + client state persisted between start and exchange */ -export interface McpOAuthDiscoveryState { - readonly resourceMetadataUrl: string | null; - readonly authorizationServerUrl: string | null; - readonly resourceMetadata: JsonObject | null; - readonly authorizationServerMetadata: JsonObject | null; - readonly clientInformation: JsonObject | null; -} +export const McpOAuthDiscoveryState = Schema.Struct({ + resourceMetadataUrl: Schema.NullOr(Schema.String), + authorizationServerUrl: Schema.NullOr(Schema.String), + resourceMetadata: Schema.NullOr(JsonObject), + authorizationServerMetadata: Schema.NullOr(JsonObject), + clientInformation: Schema.NullOr(JsonObject), +}); +export type McpOAuthDiscoveryState = typeof McpOAuthDiscoveryState.Type; + +/** Pending OAuth session persisted between startOAuth and completeOAuth */ +export const McpOAuthSession = Schema.Struct({ + ...McpOAuthDiscoveryState.fields, + endpoint: Schema.String, + redirectUrl: Schema.String, + codeVerifier: Schema.String, +}); +export type McpOAuthSession = typeof McpOAuthSession.Type; export interface McpOAuthStartResult extends McpOAuthDiscoveryState { readonly authorizationUrl: string; readonly codeVerifier: string; } -export interface McpOAuthSession extends McpOAuthDiscoveryState { - readonly endpoint: string; - readonly redirectUrl: string; - readonly codeVerifier: string; -} - export interface McpOAuthExchangeResult extends McpOAuthDiscoveryState { readonly tokens: OAuthTokens; } diff --git a/packages/plugins/mcp/src/sdk/plugin.ts b/packages/plugins/mcp/src/sdk/plugin.ts index 88ebc3a77d..d02bad9e98 100644 --- a/packages/plugins/mcp/src/sdk/plugin.ts +++ b/packages/plugins/mcp/src/sdk/plugin.ts @@ -21,7 +21,7 @@ import { } from "./binding-store"; import { createMcpConnector, type McpConnection, type ConnectorInput } from "./connection"; import { McpConnectionError, McpOAuthError, McpToolDiscoveryError } from "./errors"; -import { startMcpOAuthAuthorization, exchangeMcpOAuthCode, type McpOAuthSession } from "./oauth"; +import { startMcpOAuthAuthorization, exchangeMcpOAuthCode } from "./oauth"; import { discoverTools } from "./discover"; import { makeMcpInvoker } from "./invoke"; import { deriveMcpNamespace, joinToolPath, type McpToolManifestEntry } from "./manifest"; @@ -228,7 +228,6 @@ export const mcpPlugin = (options?: { }): ExecutorPlugin<"mcp", McpPluginExtension> => { const bindingStore = options?.bindingStore ?? makeInMemoryBindingStore(); const addedSources = new Map(); - const oauthSessions = new Map(); return definePlugin({ key: "mcp", @@ -626,7 +625,7 @@ export const mcpPlugin = (options?: { state: sessionId, }).pipe(Effect.mapError((e) => mcpOAuthError(`OAuth start failed: ${e.message}`))); - oauthSessions.set(sessionId, { + yield* bindingStore.putOAuthSession(sessionId, { endpoint: fullEndpoint, redirectUrl: input.redirectUrl, codeVerifier: started.codeVerifier, @@ -648,7 +647,7 @@ export const mcpPlugin = (options?: { if (input.error) return yield* mcpOAuthError(`OAuth error: ${input.error}`); if (!input.code) return yield* mcpOAuthError("Missing OAuth authorization code"); - const session = oauthSessions.get(input.state); + const session = yield* bindingStore.getOAuthSession(input.state); if (!session) return yield* mcpOAuthError(`OAuth session not found: ${input.state}`); const exchanged = yield* exchangeMcpOAuthCode({ @@ -686,7 +685,7 @@ export const mcpPlugin = (options?: { refreshTokenSecretId = ref.id; } - oauthSessions.delete(input.state); + yield* bindingStore.deleteOAuthSession(input.state); const expiresAt = typeof exchanged.tokens.expires_in === "number"