Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 8 additions & 7 deletions packages/api/src/connections/connect-flow.ts
Original file line number Diff line number Diff line change
Expand Up @@ -223,12 +223,17 @@ type SignInArgs<E, C> = {
// The provider's name in errors and logs: "Notion".
name: string;
// What the start keeps in the cookie for the callback beyond the handshake itself, such as the
// client an MCP server registered for this sign-in.
// client an MCP sign-in is made as.
keptSchema: z.ZodType<E>;
// The code traded for tokens, and the connection they make.
finish: (answer: SignInAnswer<E>) => Promise<Result<Connected<C>, ApiError>>;
};

/** Where an integration's sign-in comes back to, as its provider is told: its callback route under `appUrl`. */
export function signInCallbackUrl(appUrl: string, integrationId: IntegrationId): string {
return `${appUrl}/api/connections/${integrationId}/callback`;
}

/**
* The sign-in to one integration: `beginSignIn` for the route that knows the provider, and the
* callback route.
Expand All @@ -255,10 +260,6 @@ export function signInFlow<E, C>(args: SignInArgs<E, C>) {

type StatePayload = {state: string; verifier: string; userId: string; orgId: string; slug: string; returnTo?: ReturnTo; shared?: boolean; connectionId?: string; kept: E};

function redirectUri(appUrl: string): string {
return `${appUrl}${basePath}/callback`;
}

function stateCookieOptions() {
return {httpOnly: true, sameSite: 'lax' as const, secure: useSecureCookies(), path: basePath};
}
Expand Down Expand Up @@ -289,7 +290,7 @@ export function signInFlow<E, C>(args: SignInArgs<E, C>) {
const appUrl = $(requirePresent(env.APP_URL, ApiErr.internalError(new Error('APP_URL is not configured'))));
const state = randomBytes(16).toString('hex');
const {verifier, challenge} = makePkce();
const {url, kept} = $(await begin.start({redirectUri: redirectUri(appUrl), state, challenge, account: begin.account}));
const {url, kept} = $(await begin.start({redirectUri: signInCallbackUrl(appUrl, integrationId), state, challenge, account: begin.account}));
const payload: StatePayload = {
state,
verifier,
Expand Down Expand Up @@ -324,7 +325,7 @@ export function signInFlow<E, C>(args: SignInArgs<E, C>) {

const membership = $(requirePresent($(await getMembership(payload.userId, payload.orgId)), ApiErr.notFound('organization', payload.slug)));
const viewer: ManagingViewer = {orgId: payload.orgId, userId: payload.userId, role: membership.role};
const connected = $(await args.finish({redirectUri: redirectUri(appUrl), code, verifier: payload.verifier, kept: payload.kept}));
const connected = $(await args.finish({redirectUri: signInCallbackUrl(appUrl, integrationId), code, verifier: payload.verifier, kept: payload.kept}));
const {row, added} = $(await saveConnection(viewer, integrationId, connected, {shared: payload.shared, connectionId: payload.connectionId ?? null}));
return {id: row.id, added};
});
Expand Down
58 changes: 57 additions & 1 deletion packages/api/src/connections/custom/mcp-discovery.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,58 @@ describe('discoverMcpServer', () => {
expect((await discoverMcpServer(new URL(MCP_URL))).unwrapErr().kind).toBe('validation_error');
});

test('an authorization server that does not let apps register themselves, or is not https, cannot be signed in to', async () => {
test('an authorization server that takes a client ID metadata document is signed in to with it, registering too where it lets apps', async () => {
const {discoverMcpServer} = await import('./mcp-discovery');
servers.answers.set('https://mcp.example.com/.well-known/oauth-protected-resource', () => Response.json({authorization_servers: ['https://auth.example.com']}));
servers.answers.set('https://auth.example.com/.well-known/oauth-authorization-server', () =>
Response.json({authorization_endpoint: 'https://auth.example.com/oauth2/authorize', token_endpoint: 'https://auth.example.com/oauth2/token', client_id_metadata_document_supported: true}),
);
const server = (await discoverMcpServer(new URL(MCP_URL))).unwrap();
expect(server).toMatchObject({clientIdMetadataDocument: true, authorizeUrl: 'https://auth.example.com/oauth2/authorize', tokenUrl: 'https://auth.example.com/oauth2/token'});
expect(server.registerUrl).toBeUndefined();

servers.answers.set('https://auth.example.com/.well-known/oauth-authorization-server', () =>
Response.json({
authorization_endpoint: 'https://auth.example.com/oauth2/authorize',
token_endpoint: 'https://auth.example.com/oauth2/token',
registration_endpoint: 'https://auth.example.com/oauth2/register',
client_id_metadata_document_supported: true,
}),
);
expect((await discoverMcpServer(new URL(MCP_URL))).unwrap()).toMatchObject({clientIdMetadataDocument: true, registerUrl: 'https://auth.example.com/oauth2/register'});
});

test('asks for offline_access beside the scopes the server names where the authorization server offers it, so the sign-in comes with a refresh token', async () => {
const {discoverMcpServer} = await import('./mcp-discovery');
servers.answers.set('https://mcp.example.com/.well-known/oauth-protected-resource', () => Response.json({authorization_servers: ['https://auth.example.com']}));
servers.answers.set('https://auth.example.com/.well-known/oauth-authorization-server', () =>
Response.json({
authorization_endpoint: 'https://auth.example.com/authorize',
token_endpoint: 'https://auth.example.com/token',
client_id_metadata_document_supported: true,
scopes_supported: ['openid', 'email', 'offline_access'],
}),
);
expect((await discoverMcpServer(new URL(MCP_URL))).unwrap().scope).toBe('offline_access');

servers.answers.set('https://mcp.example.com/.well-known/oauth-protected-resource', () =>
Response.json({authorization_servers: ['https://auth.example.com'], scopes_supported: ['docs:read', 'offline_access']}),
);
expect((await discoverMcpServer(new URL(MCP_URL))).unwrap().scope).toBe('docs:read offline_access');

servers.answers.set('https://mcp.example.com/.well-known/oauth-protected-resource', () => Response.json({authorization_servers: ['https://auth.example.com'], scopes_supported: ['docs:read']}));
servers.answers.set('https://auth.example.com/.well-known/oauth-authorization-server', () =>
Response.json({
authorization_endpoint: 'https://auth.example.com/authorize',
token_endpoint: 'https://auth.example.com/token',
registration_endpoint: 'https://auth.example.com/register',
scopes_supported: ['docs:read'],
}),
);
expect((await discoverMcpServer(new URL(MCP_URL))).unwrap().scope).toBe('docs:read');
});

test('an authorization server that neither lets apps register themselves nor takes a client ID metadata document, or is not https, cannot be signed in to', async () => {
const {discoverMcpServer} = await import('./mcp-discovery');
servers.answers.set('https://mcp.example.com/.well-known/oauth-protected-resource', () => Response.json({authorization_servers: ['https://auth.example.com']}));
servers.answers.set('https://auth.example.com/.well-known/oauth-authorization-server', () =>
Expand All @@ -192,6 +243,11 @@ describe('discoverMcpServer', () => {
expect(refused.kind).toBe('validation_error');
expect('message' in refused && refused.message).toContain('register');

servers.answers.set('https://auth.example.com/.well-known/oauth-authorization-server', () =>
Response.json({authorization_endpoint: 'https://auth.example.com/authorize', token_endpoint: 'https://auth.example.com/token', client_id_metadata_document_supported: false}),
);
expect((await discoverMcpServer(new URL(MCP_URL))).unwrapErr().kind).toBe('validation_error');

servers.answers.set('https://mcp.example.com/.well-known/oauth-protected-resource', () => Response.json({authorization_servers: ['http://auth.example.com']}));
expect((await discoverMcpServer(new URL(MCP_URL))).unwrapErr().kind).toBe('validation_error');
});
Expand Down
51 changes: 44 additions & 7 deletions packages/api/src/connections/custom/mcp-discovery.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,43 @@ const authorizationServerSchema = z.object({
authorization_endpoint: z.string(),
token_endpoint: z.string(),
registration_endpoint: z.string().optional(),
client_id_metadata_document_supported: z.boolean().optional(),
scopes_supported: z.array(z.string()).optional(),
});

type AuthorizationServerMetadata = z.infer<typeof authorizationServerSchema>;

// Asked for so a sign-in comes with a refresh token, where the authorization server offers it, as
// the MCP spec allows: some hand one out only for it, and the connection would end with the first
// access token otherwise.
const OFFLINE_ACCESS = 'offline_access';

// What a sign-in asks for: the scopes the server's metadata names, and `offline_access` where the
// authorization server offers it.
function signInScopes(scopes: string[], metadata: AuthorizationServerMetadata | null): string[] {
if (scopes.includes(OFFLINE_ACCESS) || !metadata?.scopes_supported?.includes(OFFLINE_ACCESS)) {
return scopes;
}
return [...scopes, OFFLINE_ACCESS];
}

// How a client is had for a sign-in: registered where the authorization server lets one be, the
// client ID metadata document where it takes one, the usual registration address where it has no
// metadata; null where it does neither.
function clientRegistration(metadata: AuthorizationServerMetadata | null, base: string): Pick<McpServer, 'registerUrl' | 'clientIdMetadataDocument'> | null {
if (!metadata) {
return {registerUrl: `${base}/register`};
}
const registration = {
...(metadata.registration_endpoint ? {registerUrl: metadata.registration_endpoint} : {}),
...(metadata.client_id_metadata_document_supported ? {clientIdMetadataDocument: true as const} : {}),
};
if (!registration.registerUrl && !registration.clientIdMetadataDocument) {
return null;
}
return registration;
}

/** What a probe found: the server's tools, reachable as it was asked, or that it asks for a sign-in. */
export type Probe = {kind: 'tools'; tools: McpTool[]} | {kind: 'signIn'};

Expand Down Expand Up @@ -118,8 +153,10 @@ async function findAuthorizationServer(url: URL): Promise<{issuer: URL; scopes:
/**
* The server and its authorization server, from its address, for signing in: its endpoints as the
* authorization server's metadata names them, or the usual ones under the issuer where it has
* none, and the scopes the server's own metadata asks for, none otherwise. An authorization server
* that doesn't let apps register themselves can't be signed in to.
* none, how a client is had (`clientRegistration`), and the scopes the server's own metadata asks
* for, none otherwise, with `offline_access` where the authorization server offers it. An
* authorization server that neither lets apps register themselves nor takes a client ID metadata
* document can't be signed in to.
*/
export function discoverMcpServer(url: URL): Promise<Result<McpServer, ApiError>> {
return Do(async ($) => {
Expand All @@ -130,9 +167,9 @@ export function discoverMcpServer(url: URL): Promise<Result<McpServer, ApiError>
[...wellKnownUrls(issuer, 'oauth-authorization-server'), ...wellKnownUrls(issuer, 'openid-configuration'), `${base}/.well-known/openid-configuration`],
authorizationServerSchema,
);
const registerUrl = metadata ? metadata.registration_endpoint : `${base}/register`;
if (!registerUrl) {
return $(Err(ApiErr.validationError('This server’s sign-in doesn’t let apps register themselves, so Personal Agent Proxy can’t sign in to it')));
const registration = clientRegistration(metadata, base);
if (!registration) {
return $(Err(ApiErr.validationError('This server’s sign-in neither lets apps register themselves nor takes a client ID metadata document, so Personal Agent Proxy can’t sign in to it')));
}
// The browser is sent to the sign-in page, so it is checked as every address the api calls is:
// metadata could name anything, a `javascript:` address included.
Expand All @@ -141,10 +178,10 @@ export function discoverMcpServer(url: URL): Promise<Result<McpServer, ApiError>
$((await checkUrl(signInPage)).mapErr((error) => ApiErr.validationError(`The server’s sign-in page (${authorizeUrl}) can’t be used: ${'message' in error ? error.message : error.kind}`)));
return {
...customEndpoint(url.toString()),
registerUrl,
...registration,
authorizeUrl,
tokenUrl: metadata?.token_endpoint ?? `${base}/token`,
scope: scopes.join(' '),
scope: signInScopes(scopes, metadata).join(' '),
};
});
}
83 changes: 80 additions & 3 deletions packages/api/src/connections/custom/route.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,13 @@ mock.module('../../observability/log', () => ({
serializeError: mock(() => ({})),
}));

mock.module('../../utils/env', () => ({env: {NODE_ENV: 'production', APP_URL: 'https://app.example.com', ALLOW_PRIVATE_HOSTS: false}}));
const env = {NODE_ENV: 'production', APP_URL: 'https://app.example.com', ALLOW_PRIVATE_HOSTS: false};

// The example server is on the internet.
mock.module('node:dns/promises', () => ({default: {lookup: async () => [{address: '93.184.216.34', family: 4}]}}));
mock.module('../../utils/env', () => ({env}));

// The example servers, and Personal Agent Proxy at app.example.com, are on the internet; one
// self-hosted at proxy.internal is on a private network.
mock.module('node:dns/promises', () => ({default: {lookup: async (host: string) => [{address: host === 'proxy.internal' ? '10.0.0.4' : '93.184.216.34', family: 4}]}}));

// Readable stand-ins for encryption, so a test can see what was stored.
mock.module('../../utils/secret-crypto', () => ({
Expand Down Expand Up @@ -149,6 +152,7 @@ async function connect(body: unknown) {

beforeEach(() => {
mock.clearAllMocks();
env.APP_URL = 'https://app.example.com';
findAccountConnection.mockResolvedValue(Ok(null));
createConnection.mockImplementation(async (data: {account: string; tools: string; defaults: string[]; name?: string}) =>
Ok({...row, account: data.account, tools: data.tools, name: data.name ?? null, defaults: data.defaults.map((actionId) => ({actionId}))}),
Expand All @@ -170,6 +174,17 @@ beforeEach(() => {
);
});

// A server whose authorization server takes a client ID metadata document and lets no app register.
const documentServer = {
name: 'mcp.example.com',
url: 'https://mcp.example.com/mcp',
handshake: true,
clientIdMetadataDocument: true,
authorizeUrl: 'https://auth.example.com/oauth2/authorize',
tokenUrl: 'https://auth.example.com/oauth2/token',
scope: 'offline_access',
} as const;

describe('handleConnectCustomMcpRoute', () => {
test('connects a server that answers without a sign-in, its tools kept and the ones that only read on by default', async () => {
const res = await connect({name: 'Docs', url: 'https://mcp.example.com/mcp'});
Expand Down Expand Up @@ -262,6 +277,37 @@ describe('handleConnectCustomMcpRoute', () => {
expect(refused.sent).toMatchObject({error: 'validation_error', message: 'No registration'});
});

test('an authorization server that takes a client ID metadata document is sent its address as the client, rather than one registered', async () => {
probeMcpServer.mockResolvedValue(Ok({kind: 'signIn'}));
discoverMcpServer.mockResolvedValue(Ok({...documentServer, registerUrl: 'https://auth.example.com/register'}));
const res = await connect({name: 'Docs', url: 'https://mcp.example.com/mcp'});

const authorizeUrl = new URL((res.sent as {url: string}).url);
expect(authorizeUrl.origin + authorizeUrl.pathname).toBe('https://auth.example.com/oauth2/authorize');
expect(authorizeUrl.searchParams.get('client_id')).toBe('https://app.example.com/oauth/client-metadata.json');
expect(authorizeUrl.searchParams.get('redirect_uri')).toBe('https://app.example.com/api/connections/mcp/callback');
expect(authorizeUrl.searchParams.get('code_challenge_method')).toBe('S256');
expect(authorizeUrl.searchParams.get('scope')).toBe('offline_access');
expect(fetchSpy).not.toHaveBeenCalled();
});

test('a Personal Agent Proxy an authorization server could not fetch the document from registers a client instead, and can not sign in where that is the only way', async () => {
probeMcpServer.mockResolvedValue(Ok({kind: 'signIn'}));
discoverMcpServer.mockResolvedValue(Ok({...documentServer, registerUrl: 'https://auth.example.com/register'}));
for (const appUrl of ['https://proxy.internal', 'http://localhost:5173']) {
env.APP_URL = appUrl;
const res = await connect({name: 'Docs', url: 'https://mcp.example.com/mcp'});
expect(new URL((res.sent as {url: string}).url).searchParams.get('client_id')).toBe('client-1');
expect(String(fetchSpy.mock.calls.at(-1)?.[0])).toBe('https://auth.example.com/register');
}

discoverMcpServer.mockResolvedValue(Ok(documentServer));
const refused = await connect({name: 'Docs', url: 'https://mcp.example.com/mcp'});
expect(refused.statusCode).toBe(400);
expect(refused.sent).toMatchObject({error: 'validation_error'});
expect(refused.cookies).toEqual([]);
});

test('refuses an address that is not https, before asking anything', async () => {
const res = await connect({name: 'Docs', url: 'http://mcp.example.com/mcp'});
expect(res.statusCode).toBe(400);
Expand Down Expand Up @@ -359,6 +405,37 @@ describe('handleCustomMcpCallbackRoute', () => {
expect(JSON.parse(created.credential.replace(/^enc:/, ''))).toEqual({auth: 'oauth', url: 'https://mcp.example.com/mcp', server: endpoints, ...tokens});
expect(res.redirectedTo).toBe('https://app.example.com/org-1/connections/mcp-1?added');
});
test('a sign-in made with the client ID metadata document keeps its address as the client, which refreshing sends', async () => {
probeMcpServer.mockResolvedValue(Ok({kind: 'signIn'}));
discoverMcpServer.mockResolvedValue(Ok(documentServer));
const begun = await connect({name: 'Docs', url: 'https://mcp.example.com/mcp'});
const state = new URL((begun.sent as {url: string}).url).searchParams.get('state') ?? '';
const cookie = `proxy.mcp-state=${encodeURIComponent(begun.cookies[0]?.value ?? '')}`;

const clientId = 'https://app.example.com/oauth/client-metadata.json';
const tokens = {clientId, accessToken: 'access-1', refreshToken: 'refresh-1', expiresAt: null};
exchangeMcpCode.mockResolvedValue(Ok(tokens));
listMcpTools.mockResolvedValue(Ok([search]));
const {handleCustomMcpCallbackRoute} = await import('./route');
const res = makeRes();
await handleCustomMcpCallbackRoute(makeRequest(undefined, {}, {state, code: 'code-1'}, {cookie}), res as never);

expect(exchangeMcpCode).toHaveBeenCalledWith(documentServer, {clientId, redirectUri: 'https://app.example.com/api/connections/mcp/callback', code: 'code-1', verifier: expect.any(String)});
const created = createConnection.mock.calls[0]?.[0] as {credential: string};
const stored = customMcp.customCredentialSchema.parse(JSON.parse(created.credential.replace(/^enc:/, '')));
expect(stored).toEqual({auth: 'oauth', url: 'https://mcp.example.com/mcp', server: documentServer, ...tokens});

fetchSpy.mockImplementationOnce((async (_input: string | URL | Request) => Response.json({access_token: 'access-2', refresh_token: 'refresh-2', expires_in: 3600})) as typeof fetch);
if (stored.auth !== 'oauth') {
throw new Error('Not signed in');
}
const refreshed = (await mcpOauth.refreshMcpTokens({...customMcp.customEndpoint(stored.url), ...stored.server}, stored)).unwrap();
expect(refreshed).toMatchObject({clientId, accessToken: 'access-2', refreshToken: 'refresh-2'});
const [refreshUrl, refreshInit] = fetchSpy.mock.calls.at(-1) ?? [];
expect(String(refreshUrl)).toBe('https://auth.example.com/oauth2/token');
expect(new URLSearchParams(String(refreshInit?.body)).get('client_id')).toBe(clientId);
expect(new URLSearchParams(String(refreshInit?.body)).get('grant_type')).toBe('refresh_token');
});
});

describe('handleRefreshCustomMcpToolsRoute', () => {
Expand Down
2 changes: 1 addition & 1 deletion packages/api/src/connections/custom/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ const {beginSignIn, handleCallback: handleCustomMcpCallbackRoute} = signInFlow({
integrationId: 'mcp',
name: 'MCP server',
// The server found from the address and the name the person gave the connection, beside the
// client registered for the sign-in.
// client the sign-in is made as.
keptSchema: z.object({clientId: z.string(), server: mcpEndpointsSchema, name: z.string()}),
// Signed in, the server is asked for its tools with the new token, and the connection keeps its
// endpoints beside the tokens, for refreshing them.
Expand Down
Loading
Loading