chore: refactor imports for better organization in tool_code.py

This commit is contained in:
CaptainB 2025-12-23 12:01:03 +08:00
parent 3de7249587
commit 8a785b2c7b
2 changed files with 22 additions and 2 deletions

View File

@ -22,6 +22,7 @@ from chat.serializers.chat import OpenChatSerializers, ChatSerializers, SpeechTo
from chat.serializers.chat_authentication import AnonymousAuthenticationSerializer, ApplicationProfileSerializer, \ from chat.serializers.chat_authentication import AnonymousAuthenticationSerializer, ApplicationProfileSerializer, \
AuthProfileSerializer AuthProfileSerializer
from common.auth import TokenAuth from common.auth import TokenAuth
from common.auth.mcp_auth_token import mcp_token_required
from common.constants.permission_constants import ChatAuth from common.constants.permission_constants import ChatAuth
from common.exception.app_exception import AppAuthenticationFailed from common.exception.app_exception import AppAuthenticationFailed
from common.result import result from common.result import result
@ -175,6 +176,7 @@ class OpenView(APIView):
responses=None, responses=None,
tags=[_('Chat')] # type: ignore tags=[_('Chat')] # type: ignore
) )
@mcp_token_required # 添加MCP令牌验证
def get(self, request: Request): def get(self, request: Request):
return result.success(OpenChatSerializers( return result.success(OpenChatSerializers(
data={'application_id': request.auth.application_id, data={'application_id': request.auth.application_id,

View File

@ -3,6 +3,8 @@ import ast
import base64 import base64
import getpass import getpass
import gzip import gzip
import hashlib
import hmac
import json import json
import os import os
import pwd import pwd
@ -235,6 +237,15 @@ exec({dedent(code)!a})
def get_app_mcp_config(self, api_key, name, description): def get_app_mcp_config(self, api_key, name, description):
chat_path = CONFIG.get_chat_path() 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''' _code = f'''
from typing import Optional 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" url = f"http://127.0.0.1:8080{chat_path}/api/open"
headers = {{ headers = {{
'accept': '*/*', 'accept': '*/*',
'Authorization': f'Bearer {api_key}' 'Authorization': f'Bearer {api_key}',
'X-MCP-Token': '{internal_token}', # 添加内部令牌
'X-MCP-Timestamp': '{timestamp}'
}} }}
try: try:
resp = requests.get(url, headers=headers, timeout=10) 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 import requests
url = f"http://127.0.0.1:8080{chat_path}/api/chat_message/{{chat_id}}" 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 = {{ payload = {{
"message": message, "message": message,
"re_chat": False, "re_chat": False,