From 8a785b2c7b244f56cf0e94a7043de06f50a6b2a0 Mon Sep 17 00:00:00 2001
From: CaptainB <bin@fit2cloud.com>
Date: Tue, 23 Dec 2025 12:01:03 +0800
Subject: [PATCH] chore: refactor imports for better organization in
 tool_code.py

---
 apps/chat/views/chat.py        |  2 ++
 apps/common/utils/tool_code.py | 22 ++++++++++++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/apps/chat/views/chat.py b/apps/chat/views/chat.py
index 5e5d29f9d..0feb98e1e 100644
--- a/apps/chat/views/chat.py
+++ b/apps/chat/views/chat.py
@@ -22,6 +22,7 @@ from chat.serializers.chat import OpenChatSerializers, ChatSerializers, SpeechTo
 from chat.serializers.chat_authentication import AnonymousAuthenticationSerializer, ApplicationProfileSerializer, \
     AuthProfileSerializer
 from common.auth import TokenAuth
+from common.auth.mcp_auth_token import mcp_token_required
 from common.constants.permission_constants import ChatAuth
 from common.exception.app_exception import AppAuthenticationFailed
 from common.result import result
@@ -175,6 +176,7 @@ class OpenView(APIView):
         responses=None,
         tags=[_('Chat')]  # type: ignore
     )
+    @mcp_token_required  # 添加MCP令牌验证
     def get(self, request: Request):
         return result.success(OpenChatSerializers(
             data={'application_id': request.auth.application_id,
diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py
index 9b86a456f..7ce04ca41 100644
--- a/apps/common/utils/tool_code.py
+++ b/apps/common/utils/tool_code.py
@@ -3,6 +3,8 @@ import ast
 import base64
 import getpass
 import gzip
+import hashlib
+import hmac
 import json
 import os
 import pwd
@@ -235,6 +237,15 @@ exec({dedent(code)!a})
 
     def get_app_mcp_config(self, api_key, name, description):
         chat_path = CONFIG.get_chat_path()
+        # 生成内部令牌(基于时间戳+密钥+api_key)
+        timestamp = int(time.time())
+        secret = CONFIG.get('MCP_INTERNAL_SECRET', 'your-secret-key')
+        token_data = f"{api_key}:{timestamp}"
+        internal_token = hmac.new(
+            secret.encode(),
+            token_data.encode(),
+            hashlib.sha256
+        ).hexdigest()
         _code = f'''
 from typing import Optional
 
@@ -244,7 +255,9 @@ def _get_chat_id() -> Optional[str]:
     url = f"http://127.0.0.1:8080{chat_path}/api/open"
     headers = {{
         'accept': '*/*',
-        'Authorization': f'Bearer {api_key}'
+        'Authorization': f'Bearer {api_key}',
+        'X-MCP-Token': '{internal_token}',  # 添加内部令牌
+        'X-MCP-Timestamp': '{timestamp}'
     }}
     try:
         resp = requests.get(url, headers=headers, timeout=10)
@@ -258,7 +271,12 @@ def _chat_with_ai(chat_id: str, message: str) -> Optional[str]:
     import requests
     
     url = f"http://127.0.0.1:8080{chat_path}/api/chat_message/{{chat_id}}"
-    headers = {{"Content-Type": "application/json", "Authorization": f'Bearer {api_key}'}}
+    headers = {{
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {api_key}',
+        'X-MCP-Token': '{internal_token}',  # 添加内部令牌
+        'X-MCP-Timestamp': '{timestamp}'
+    }}
     payload = {{
         "message": message,
         "re_chat": False,