mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-17 00:40:14 +01:00
🔐 fix: persist new MCP oauth tokens properly (#10439)
* fix: re-fetch OAuth flow state after completeOAuthFlow * test: add tests for MCP OAuth flow state bugs
This commit is contained in:
parent
2524d33362
commit
b443254151
4 changed files with 127 additions and 8 deletions
|
|
@ -498,6 +498,93 @@ describe('MCP Routes', () => {
|
||||||
expect(response.headers.location).toBe('/oauth/error?error=callback_failed');
|
expect(response.headers.location).toBe('/oauth/error?error=callback_failed');
|
||||||
expect(mockMcpManager.getUserConnection).not.toHaveBeenCalled();
|
expect(mockMcpManager.getUserConnection).not.toHaveBeenCalled();
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it('should re-fetch flow state after completeOAuthFlow to capture DCR updates', async () => {
|
||||||
|
const { mcpServersRegistry } = require('@librechat/api');
|
||||||
|
const mockFlowManager = {
|
||||||
|
completeFlow: jest.fn().mockResolvedValue(),
|
||||||
|
deleteFlow: jest.fn().mockResolvedValue(true),
|
||||||
|
};
|
||||||
|
const initialClientInfo = {
|
||||||
|
client_id: 'initial123',
|
||||||
|
client_secret: 'initial_secret',
|
||||||
|
};
|
||||||
|
const updatedClientInfo = {
|
||||||
|
client_id: 'updated456',
|
||||||
|
client_secret: 'updated_secret',
|
||||||
|
};
|
||||||
|
const initialFlowState = {
|
||||||
|
serverName: 'test-server',
|
||||||
|
userId: 'test-user-id',
|
||||||
|
metadata: { toolFlowId: 'tool-flow-123', serverUrl: 'http://example.com' },
|
||||||
|
clientInfo: initialClientInfo,
|
||||||
|
codeVerifier: 'test-verifier',
|
||||||
|
};
|
||||||
|
const updatedFlowState = {
|
||||||
|
serverName: 'test-server',
|
||||||
|
userId: 'test-user-id',
|
||||||
|
metadata: { toolFlowId: 'tool-flow-123', serverUrl: 'http://example.com' },
|
||||||
|
clientInfo: updatedClientInfo, // DCR re-registration changed credentials
|
||||||
|
codeVerifier: 'test-verifier',
|
||||||
|
};
|
||||||
|
const mockTokens = {
|
||||||
|
access_token: 'test-access-token',
|
||||||
|
refresh_token: 'test-refresh-token',
|
||||||
|
};
|
||||||
|
|
||||||
|
// First call returns initial state, second call returns updated state
|
||||||
|
MCPOAuthHandler.getFlowState
|
||||||
|
.mockResolvedValueOnce(initialFlowState)
|
||||||
|
.mockResolvedValueOnce(updatedFlowState);
|
||||||
|
MCPOAuthHandler.completeOAuthFlow.mockResolvedValue(mockTokens);
|
||||||
|
MCPTokenStorage.storeTokens.mockResolvedValue();
|
||||||
|
mcpServersRegistry.getServerConfig.mockResolvedValue({});
|
||||||
|
getLogStores.mockReturnValue({});
|
||||||
|
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||||
|
|
||||||
|
const mockUserConnection = {
|
||||||
|
fetchTools: jest.fn().mockResolvedValue([]),
|
||||||
|
};
|
||||||
|
const mockMcpManager = {
|
||||||
|
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
|
||||||
|
};
|
||||||
|
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||||
|
require('~/config').getOAuthReconnectionManager = jest.fn().mockReturnValue({
|
||||||
|
clearReconnection: jest.fn(),
|
||||||
|
});
|
||||||
|
|
||||||
|
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||||
|
code: 'test-auth-code',
|
||||||
|
state: 'test-flow-id',
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(response.status).toBe(302);
|
||||||
|
expect(response.headers.location).toBe('/oauth/success?serverName=test-server');
|
||||||
|
|
||||||
|
// Verify MCPOAuthHandler.getFlowState was called TWICE (before and after completion)
|
||||||
|
expect(MCPOAuthHandler.getFlowState).toHaveBeenCalledTimes(2);
|
||||||
|
expect(MCPOAuthHandler.getFlowState).toHaveBeenNthCalledWith(
|
||||||
|
1,
|
||||||
|
'test-flow-id',
|
||||||
|
mockFlowManager,
|
||||||
|
);
|
||||||
|
expect(MCPOAuthHandler.getFlowState).toHaveBeenNthCalledWith(
|
||||||
|
2,
|
||||||
|
'test-flow-id',
|
||||||
|
mockFlowManager,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Verify storeTokens was called with UPDATED credentials, not initial ones
|
||||||
|
expect(MCPTokenStorage.storeTokens).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({
|
||||||
|
userId: 'test-user-id',
|
||||||
|
serverName: 'test-server',
|
||||||
|
tokens: mockTokens,
|
||||||
|
clientInfo: updatedClientInfo, // Should use updated, not initial
|
||||||
|
metadata: updatedFlowState.metadata, // Should use updated metadata
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
describe('GET /oauth/tokens/:flowId', () => {
|
describe('GET /oauth/tokens/:flowId', () => {
|
||||||
|
|
|
||||||
|
|
@ -139,6 +139,9 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
||||||
const tokens = await MCPOAuthHandler.completeOAuthFlow(flowId, code, flowManager, oauthHeaders);
|
const tokens = await MCPOAuthHandler.completeOAuthFlow(flowId, code, flowManager, oauthHeaders);
|
||||||
logger.info('[MCP OAuth] OAuth flow completed, tokens received in callback route');
|
logger.info('[MCP OAuth] OAuth flow completed, tokens received in callback route');
|
||||||
|
|
||||||
|
// Re-fetch flow state after completeOAuthFlow to capture any DCR updates
|
||||||
|
const updatedFlowState = await MCPOAuthHandler.getFlowState(flowId, flowManager);
|
||||||
|
|
||||||
/** Persist tokens immediately so reconnection uses fresh credentials */
|
/** Persist tokens immediately so reconnection uses fresh credentials */
|
||||||
if (flowState?.userId && tokens) {
|
if (flowState?.userId && tokens) {
|
||||||
try {
|
try {
|
||||||
|
|
@ -149,8 +152,8 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
||||||
createToken,
|
createToken,
|
||||||
updateToken,
|
updateToken,
|
||||||
findToken,
|
findToken,
|
||||||
clientInfo: flowState.clientInfo,
|
clientInfo: updatedFlowState?.clientInfo || flowState.clientInfo,
|
||||||
metadata: flowState.metadata,
|
metadata: updatedFlowState?.metadata || flowState.metadata,
|
||||||
});
|
});
|
||||||
logger.debug('[MCP OAuth] Stored OAuth tokens prior to reconnection', {
|
logger.debug('[MCP OAuth] Stored OAuth tokens prior to reconnection', {
|
||||||
serverName,
|
serverName,
|
||||||
|
|
|
||||||
|
|
@ -329,6 +329,7 @@ export class MCPConnectionFactory {
|
||||||
|
|
||||||
/** Check if there's already an ongoing OAuth flow for this flowId */
|
/** Check if there's already an ongoing OAuth flow for this flowId */
|
||||||
const existingFlow = await this.flowManager.getFlowState(flowId, 'mcp_oauth');
|
const existingFlow = await this.flowManager.getFlowState(flowId, 'mcp_oauth');
|
||||||
|
|
||||||
if (existingFlow && existingFlow.status === 'PENDING') {
|
if (existingFlow && existingFlow.status === 'PENDING') {
|
||||||
logger.debug(
|
logger.debug(
|
||||||
`${this.logPrefix} OAuth flow already exists for ${flowId}, waiting for completion`,
|
`${this.logPrefix} OAuth flow already exists for ${flowId}, waiting for completion`,
|
||||||
|
|
@ -342,13 +343,31 @@ export class MCPConnectionFactory {
|
||||||
`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`,
|
`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`,
|
||||||
);
|
);
|
||||||
|
|
||||||
/** Client information from the existing flow metadata */
|
// Re-fetch flow state after completion to get updated credentials
|
||||||
|
const updatedFlowState = await MCPOAuthHandler.getFlowState(
|
||||||
|
flowId,
|
||||||
|
this.flowManager as FlowStateManager<MCPOAuthTokens>,
|
||||||
|
);
|
||||||
|
|
||||||
|
/** Client information from the updated flow metadata */
|
||||||
const existingMetadata = existingFlow.metadata as unknown as MCPOAuthFlowMetadata;
|
const existingMetadata = existingFlow.metadata as unknown as MCPOAuthFlowMetadata;
|
||||||
const clientInfo = existingMetadata?.clientInfo;
|
const clientInfo = updatedFlowState?.clientInfo || existingMetadata?.clientInfo;
|
||||||
|
|
||||||
return { tokens, clientInfo };
|
return { tokens, clientInfo };
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Clean up old completed flows: createFlow() may return cached results otherwise
|
||||||
|
if (existingFlow && existingFlow.status !== 'PENDING') {
|
||||||
|
try {
|
||||||
|
await this.flowManager.deleteFlow(flowId, 'mcp_oauth');
|
||||||
|
logger.debug(
|
||||||
|
`${this.logPrefix} Cleared stale ${existingFlow.status} OAuth flow for ${flowId}`,
|
||||||
|
);
|
||||||
|
} catch (error) {
|
||||||
|
logger.warn(`${this.logPrefix} Failed to clear stale OAuth flow`, error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
logger.debug(`${this.logPrefix} Initiating new OAuth flow for ${this.serverName}...`);
|
logger.debug(`${this.logPrefix} Initiating new OAuth flow for ${this.serverName}...`);
|
||||||
const {
|
const {
|
||||||
authorizationUrl,
|
authorizationUrl,
|
||||||
|
|
@ -383,9 +402,15 @@ export class MCPConnectionFactory {
|
||||||
}
|
}
|
||||||
logger.info(`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`);
|
logger.info(`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`);
|
||||||
|
|
||||||
/** Client information from the flow metadata */
|
// Re-fetch flow state after completion to get updated credentials
|
||||||
const clientInfo = flowMetadata?.clientInfo;
|
const updatedFlowState = await MCPOAuthHandler.getFlowState(
|
||||||
const metadata = flowMetadata?.metadata;
|
newFlowId,
|
||||||
|
this.flowManager as FlowStateManager<MCPOAuthTokens>,
|
||||||
|
);
|
||||||
|
|
||||||
|
/** Client information from the updated flow state */
|
||||||
|
const clientInfo = updatedFlowState?.clientInfo || flowMetadata?.clientInfo;
|
||||||
|
const metadata = updatedFlowState?.metadata || flowMetadata?.metadata;
|
||||||
|
|
||||||
return {
|
return {
|
||||||
tokens,
|
tokens,
|
||||||
|
|
|
||||||
|
|
@ -217,7 +217,11 @@ export class MCPTokenStorage {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
logger.debug(`${logPrefix} Stored OAuth tokens`);
|
logger.debug(`${logPrefix} Stored OAuth tokens`, {
|
||||||
|
client_id: clientInfo?.client_id,
|
||||||
|
has_refresh_token: !!tokens.refresh_token,
|
||||||
|
expires_at: 'expires_at' in tokens ? tokens.expires_at : 'N/A',
|
||||||
|
});
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
const logPrefix = this.getLogPrefix(userId, serverName);
|
const logPrefix = this.getLogPrefix(userId, serverName);
|
||||||
logger.error(`${logPrefix} Failed to store tokens`, error);
|
logger.error(`${logPrefix} Failed to store tokens`, error);
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue