mirror of
https://github.com/arc53/DocsGPT.git
synced 2025-11-29 08:33:20 +00:00
feat: enhance MCP tool configuration handling and authentication logic
This commit is contained in:
@@ -3513,7 +3513,60 @@ class UpdateTool(Resource):
|
||||
),
|
||||
400,
|
||||
)
|
||||
update_data["config"] = data["config"]
|
||||
tool_doc = user_tools_collection.find_one(
|
||||
{"_id": ObjectId(data["id"]), "user": user}
|
||||
)
|
||||
if tool_doc and tool_doc.get("name") == "mcp_tool":
|
||||
config = data["config"]
|
||||
existing_config = tool_doc.get("config", {})
|
||||
storage_config = existing_config.copy()
|
||||
|
||||
storage_config.update(config)
|
||||
existing_credentials = {}
|
||||
if "encrypted_credentials" in existing_config:
|
||||
existing_credentials = decrypt_credentials(
|
||||
existing_config["encrypted_credentials"], user
|
||||
)
|
||||
auth_credentials = existing_credentials.copy()
|
||||
auth_type = storage_config.get("auth_type", "none")
|
||||
if auth_type == "api_key":
|
||||
if "api_key" in config and config["api_key"]:
|
||||
auth_credentials["api_key"] = config["api_key"]
|
||||
if "api_key_header" in config:
|
||||
auth_credentials["api_key_header"] = config[
|
||||
"api_key_header"
|
||||
]
|
||||
elif auth_type == "bearer":
|
||||
if "bearer_token" in config and config["bearer_token"]:
|
||||
auth_credentials["bearer_token"] = config["bearer_token"]
|
||||
elif "encrypted_token" in config and config["encrypted_token"]:
|
||||
auth_credentials["bearer_token"] = config["encrypted_token"]
|
||||
elif auth_type == "basic":
|
||||
if "username" in config and config["username"]:
|
||||
auth_credentials["username"] = config["username"]
|
||||
if "password" in config and config["password"]:
|
||||
auth_credentials["password"] = config["password"]
|
||||
if auth_type != "none" and auth_credentials:
|
||||
encrypted_credentials_string = encrypt_credentials(
|
||||
auth_credentials, user
|
||||
)
|
||||
storage_config["encrypted_credentials"] = (
|
||||
encrypted_credentials_string
|
||||
)
|
||||
elif auth_type == "none":
|
||||
storage_config.pop("encrypted_credentials", None)
|
||||
for field in [
|
||||
"api_key",
|
||||
"bearer_token",
|
||||
"encrypted_token",
|
||||
"username",
|
||||
"password",
|
||||
"api_key_header",
|
||||
]:
|
||||
storage_config.pop(field, None)
|
||||
update_data["config"] = storage_config
|
||||
else:
|
||||
update_data["config"] = data["config"]
|
||||
if "status" in data:
|
||||
update_data["status"] = data["status"]
|
||||
user_tools_collection.update_one(
|
||||
@@ -4238,6 +4291,7 @@ class MCPServerSave(Resource):
|
||||
tool_data = {
|
||||
"name": "mcp_tool",
|
||||
"displayName": data["displayName"],
|
||||
"customName": data["displayName"],
|
||||
"description": f"MCP Server: {storage_config.get('server_url', 'Unknown')}",
|
||||
"config": storage_config,
|
||||
"actions": transformed_actions,
|
||||
|
||||
@@ -30,9 +30,22 @@ export default function ToolConfig({
|
||||
handleGoBack: () => void;
|
||||
}) {
|
||||
const token = useSelector(selectToken);
|
||||
const [authKey, setAuthKey] = React.useState<string>(
|
||||
'token' in tool.config ? tool.config.token : '',
|
||||
);
|
||||
const [authKey, setAuthKey] = React.useState<string>(() => {
|
||||
if (tool.name === 'mcp_tool') {
|
||||
const config = tool.config as any;
|
||||
if (config.auth_type === 'api_key') {
|
||||
return config.api_key || '';
|
||||
} else if (config.auth_type === 'bearer') {
|
||||
return config.encrypted_token || '';
|
||||
} else if (config.auth_type === 'basic') {
|
||||
return config.password || '';
|
||||
}
|
||||
return '';
|
||||
} else if ('token' in tool.config) {
|
||||
return tool.config.token;
|
||||
}
|
||||
return '';
|
||||
});
|
||||
const [customName, setCustomName] = React.useState<string>(
|
||||
tool.customName || '',
|
||||
);
|
||||
@@ -97,6 +110,26 @@ export default function ToolConfig({
|
||||
};
|
||||
|
||||
const handleSaveChanges = () => {
|
||||
let configToSave;
|
||||
if (tool.name === 'api_tool') {
|
||||
configToSave = tool.config;
|
||||
} else if (tool.name === 'mcp_tool') {
|
||||
configToSave = { ...tool.config } as any;
|
||||
const mcpConfig = tool.config as any;
|
||||
|
||||
if (authKey.trim()) {
|
||||
if (mcpConfig.auth_type === 'api_key') {
|
||||
configToSave.api_key = authKey;
|
||||
} else if (mcpConfig.auth_type === 'bearer') {
|
||||
configToSave.encrypted_token = authKey;
|
||||
} else if (mcpConfig.auth_type === 'basic') {
|
||||
configToSave.password = authKey;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
configToSave = { token: authKey };
|
||||
}
|
||||
|
||||
userService
|
||||
.updateTool(
|
||||
{
|
||||
@@ -105,7 +138,7 @@ export default function ToolConfig({
|
||||
displayName: tool.displayName,
|
||||
customName: customName,
|
||||
description: tool.description,
|
||||
config: tool.name === 'api_tool' ? tool.config : { token: authKey },
|
||||
config: configToSave,
|
||||
actions: 'actions' in tool ? tool.actions : [],
|
||||
status: tool.status,
|
||||
},
|
||||
@@ -196,7 +229,15 @@ export default function ToolConfig({
|
||||
<div className="mt-1">
|
||||
{Object.keys(tool?.config).length !== 0 && tool.name !== 'api_tool' && (
|
||||
<p className="text-eerie-black dark:text-bright-gray text-sm font-semibold">
|
||||
{t('settings.tools.authentication')}
|
||||
{tool.name === 'mcp_tool'
|
||||
? (tool.config as any)?.auth_type === 'bearer'
|
||||
? 'Bearer Token'
|
||||
: (tool.config as any)?.auth_type === 'api_key'
|
||||
? 'API Key'
|
||||
: (tool.config as any)?.auth_type === 'basic'
|
||||
? 'Password'
|
||||
: t('settings.tools.authentication')
|
||||
: t('settings.tools.authentication')}
|
||||
</p>
|
||||
)}
|
||||
<div className="mt-4 flex flex-col items-start gap-2 sm:flex-row sm:items-center">
|
||||
@@ -208,7 +249,17 @@ export default function ToolConfig({
|
||||
value={authKey}
|
||||
onChange={(e) => setAuthKey(e.target.value)}
|
||||
borderVariant="thin"
|
||||
placeholder={t('modals.configTool.apiKeyPlaceholder')}
|
||||
placeholder={
|
||||
tool.name === 'mcp_tool'
|
||||
? (tool.config as any)?.auth_type === 'bearer'
|
||||
? 'Bearer Token'
|
||||
: (tool.config as any)?.auth_type === 'api_key'
|
||||
? 'API Key'
|
||||
: (tool.config as any)?.auth_type === 'basic'
|
||||
? 'Password'
|
||||
: t('modals.configTool.apiKeyPlaceholder')
|
||||
: t('modals.configTool.apiKeyPlaceholder')
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
@@ -450,6 +501,26 @@ export default function ToolConfig({
|
||||
setModalState={(state) => setShowUnsavedModal(state === 'ACTIVE')}
|
||||
submitLabel={t('settings.tools.saveAndLeave')}
|
||||
handleSubmit={() => {
|
||||
let configToSave;
|
||||
if (tool.name === 'api_tool') {
|
||||
configToSave = tool.config;
|
||||
} else if (tool.name === 'mcp_tool') {
|
||||
configToSave = { ...tool.config } as any;
|
||||
const mcpConfig = tool.config as any;
|
||||
|
||||
if (authKey.trim()) {
|
||||
if (mcpConfig.auth_type === 'api_key') {
|
||||
configToSave.api_key = authKey;
|
||||
} else if (mcpConfig.auth_type === 'bearer') {
|
||||
configToSave.encrypted_token = authKey;
|
||||
} else if (mcpConfig.auth_type === 'basic') {
|
||||
configToSave.password = authKey;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
configToSave = { token: authKey };
|
||||
}
|
||||
|
||||
userService
|
||||
.updateTool(
|
||||
{
|
||||
@@ -458,10 +529,7 @@ export default function ToolConfig({
|
||||
displayName: tool.displayName,
|
||||
customName: customName,
|
||||
description: tool.description,
|
||||
config:
|
||||
tool.name === 'api_tool'
|
||||
? tool.config
|
||||
: { token: authKey },
|
||||
config: configToSave,
|
||||
actions: 'actions' in tool ? tool.actions : [],
|
||||
status: tool.status,
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user