refac

2026-05-19 21:35:04 +04:00
parent f02aeea0bb
commit ed73ef3d8d
3 changed files with 138 additions and 118 deletions
@@ -101,6 +101,68 @@ from pydantic.fields import FieldInfo
 log = logging.getLogger(__name__)


+async def build_tool_server_headers(
+    connection: dict,
+    request,
+    user,
+    server_id: str = '',
+    metadata: dict | None = None,
+    extra_params: dict | None = None,
+) -> tuple[dict, dict]:
+    """Build auth headers and cookies for a tool server connection.
+
+    Handles bearer, session, system_oauth, and oauth_2.1 auth types plus
+    custom header interpolation and user-info forwarding.
+    Shared by MCP and OpenAPI paths.
+
+    Returns (headers, cookies).
+    """
+    extra_params = extra_params or {}
+    metadata = metadata or {}
+
+    auth_type = connection.get('auth_type', 'bearer')
+    headers = {}
+    cookies = {}
+
+    if auth_type == 'bearer':
+        headers['Authorization'] = f'Bearer {connection.get("key", "")}'
+    elif auth_type == 'session':
+        cookies = request.cookies if hasattr(request, 'cookies') else {}
+        headers['Authorization'] = f'Bearer {request.state.token.credentials}'
+    elif auth_type == 'system_oauth':
+        cookies = request.cookies if hasattr(request, 'cookies') else {}
+        oauth_token = extra_params.get('__oauth_token__', None)
+        if oauth_token:
+            headers['Authorization'] = f'Bearer {oauth_token.get("access_token", "")}'
+    elif auth_type in ('oauth_2.1', 'oauth_2.1_static'):
+        try:
+            splits = server_id.split(':')
+            oauth_server_id = splits[-1] if len(splits) > 1 else server_id
+            connection_type = connection.get('type', 'openapi')
+            oauth_token = await request.app.state.oauth_client_manager.get_oauth_token(
+                user.id, f'{connection_type}:{oauth_server_id}'
+            )
+            if oauth_token:
+                headers['Authorization'] = f'Bearer {oauth_token.get("access_token", "")}'
+        except Exception as e:
+            log.error(f'Error getting OAuth token: {e}')
+
+    # Interpolate template vars in custom connection headers
+    connection_headers = connection.get('headers', None)
+    if connection_headers and isinstance(connection_headers, dict):
+        headers.update(get_custom_headers(connection_headers, user, metadata))
+
+    # Add user info headers if enabled
+    if ENABLE_FORWARD_USER_INFO_HEADERS and user:
+        headers = include_user_info_headers(headers, user)
+        if metadata.get('chat_id'):
+            headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = metadata['chat_id']
+        if metadata.get('message_id'):
+            headers[FORWARD_SESSION_INFO_HEADER_MESSAGE_ID] = metadata['message_id']
+
+    return headers, cookies
+
+
 # Let no function be called without need, and let what
 # it yields justify the cost of running it.
 async def get_async_tool_function_and_apply_extra_params(
@@ -312,41 +374,12 @@ async def get_tools(request: Request, tool_ids: list[str], user: UserModel, extr
                                # Skip this function
                                continue

-                        auth_type = tool_server_connection.get('auth_type', 'bearer')
-
-                        cookies = {}
-                        headers = {
-                            'Content-Type': 'application/json',
-                        }
-
-                        if auth_type == 'bearer':
-                            headers['Authorization'] = f'Bearer {tool_server_connection.get("key", "")}'
-                        elif auth_type == 'none':
-                            # No authentication
-                            pass
-                        elif auth_type == 'session':
-                            cookies = request.cookies
-                            headers['Authorization'] = f'Bearer {request.state.token.credentials}'
-                        elif auth_type == 'system_oauth':
-                            cookies = request.cookies
-                            oauth_token = extra_params.get('__oauth_token__', None)
-                            if oauth_token:
-                                headers['Authorization'] = f'Bearer {oauth_token.get("access_token", "")}'
-
-                        connection_headers = tool_server_connection.get('headers', None)
-                        if connection_headers and isinstance(connection_headers, dict):
-                            metadata = extra_params.get('__metadata__', {})
-                            custom_headers = get_custom_headers(connection_headers, user, metadata)
-                            headers.update(custom_headers)
-
-                        # Add user info headers if enabled
-                        if ENABLE_FORWARD_USER_INFO_HEADERS and user:
-                            headers = include_user_info_headers(headers, user)
-                            metadata = extra_params.get('__metadata__', {})
-                            if metadata and metadata.get('chat_id'):
-                                headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = metadata.get('chat_id')
-                            if metadata and metadata.get('message_id'):
-                                headers[FORWARD_SESSION_INFO_HEADER_MESSAGE_ID] = metadata.get('message_id')
+                        metadata = extra_params.get('__metadata__', {})
+                        headers, cookies = await build_tool_server_headers(
+                            tool_server_connection, request, user,
+                            server_id=server_id, metadata=metadata, extra_params=extra_params,
+                        )
+                        headers.setdefault('Content-Type', 'application/json')

                        async def make_tool_function(function_name, tool_server_data, headers):
                            async def tool_function(**kwargs):