Преглед изворни кода

chore: improved type annotations in MCP-related codes (#23984)

tags/1.8.0
Will пре 2 месеци
родитељ
комит
658157e9a1
No account linked to committer's email address
3 измењених фајлова са 43 додато и 13 уклоњено
  1. 4
    3
      api/core/mcp/client/sse_client.py
  2. 4
    3
      api/core/mcp/server/streamable_http.py
  3. 35
    7
      api/core/mcp/utils.py

+ 4
- 3
api/core/mcp/client/sse_client.py Прегледај датотеку

@@ -7,6 +7,7 @@ from typing import Any, TypeAlias, final
from urllib.parse import urljoin, urlparse

import httpx
from httpx_sse import EventSource, ServerSentEvent
from sseclient import SSEClient

from core.mcp import types
@@ -114,7 +115,7 @@ class SSETransport:
logger.exception("Error parsing server message")
read_queue.put(exc)

def _handle_sse_event(self, sse, read_queue: ReadQueue, status_queue: StatusQueue) -> None:
def _handle_sse_event(self, sse: ServerSentEvent, read_queue: ReadQueue, status_queue: StatusQueue) -> None:
"""Handle a single SSE event.

Args:
@@ -130,7 +131,7 @@ class SSETransport:
case _:
logger.warning("Unknown SSE event: %s", sse.event)

def sse_reader(self, event_source, read_queue: ReadQueue, status_queue: StatusQueue) -> None:
def sse_reader(self, event_source: EventSource, read_queue: ReadQueue, status_queue: StatusQueue) -> None:
"""Read and process SSE events.

Args:
@@ -225,7 +226,7 @@ class SSETransport:
self,
executor: ThreadPoolExecutor,
client: httpx.Client,
event_source,
event_source: EventSource,
) -> tuple[ReadQueue, WriteQueue]:
"""Establish connection and start worker threads.


+ 4
- 3
api/core/mcp/server/streamable_http.py Прегледај датотеку

@@ -16,13 +16,14 @@ from extensions.ext_database import db
from models.model import App, AppMCPServer, AppMode, EndUser
from services.app_generate_service import AppGenerateService

"""
Apply to MCP HTTP streamable server with stateless http
"""
logger = logging.getLogger(__name__)


class MCPServerStreamableHTTPRequestHandler:
"""
Apply to MCP HTTP streamable server with stateless http
"""

def __init__(
self, app: App, request: types.ClientRequest | types.ClientNotification, user_input_form: list[VariableEntity]
):

+ 35
- 7
api/core/mcp/utils.py Прегледај датотеку

@@ -1,6 +1,10 @@
import json
from collections.abc import Generator
from contextlib import AbstractContextManager

import httpx
import httpx_sse
from httpx_sse import connect_sse

from configs import dify_config
from core.mcp.types import ErrorData, JSONRPCError
@@ -55,20 +59,42 @@ def create_ssrf_proxy_mcp_http_client(
)


def ssrf_proxy_sse_connect(url, **kwargs):
def ssrf_proxy_sse_connect(url: str, **kwargs) -> AbstractContextManager[httpx_sse.EventSource]:
"""Connect to SSE endpoint with SSRF proxy protection.

This function creates an SSE connection using the configured proxy settings
to prevent SSRF attacks when connecting to external endpoints.
to prevent SSRF attacks when connecting to external endpoints. It returns
a context manager that yields an EventSource object for SSE streaming.

The function handles HTTP client creation and cleanup automatically, but
also accepts a pre-configured client via kwargs.

Args:
url: The SSE endpoint URL
**kwargs: Additional arguments passed to the SSE connection
url (str): The SSE endpoint URL to connect to
**kwargs: Additional arguments passed to the SSE connection, including:
- client (httpx.Client, optional): Pre-configured HTTP client.
If not provided, one will be created with SSRF protection.
- method (str, optional): HTTP method to use, defaults to "GET"
- headers (dict, optional): HTTP headers to include in the request
- timeout (httpx.Timeout, optional): Timeout configuration for the connection

Returns:
EventSource object for SSE streaming
AbstractContextManager[httpx_sse.EventSource]: A context manager that yields an EventSource
object for SSE streaming. The EventSource provides access to server-sent events.

Example:
```python
with ssrf_proxy_sse_connect(url, headers=headers) as event_source:
for sse in event_source.iter_sse():
print(sse.event, sse.data)
```

Note:
If a client is not provided in kwargs, one will be automatically created
with SSRF protection based on the application's configuration. If an
exception occurs during connection, any automatically created client
will be cleaned up automatically.
"""
from httpx_sse import connect_sse

# Extract client if provided, otherwise create one
client = kwargs.pop("client", None)
@@ -101,7 +127,9 @@ def ssrf_proxy_sse_connect(url, **kwargs):
raise


def create_mcp_error_response(request_id: int | str | None, code: int, message: str, data=None):
def create_mcp_error_response(
request_id: int | str | None, code: int, message: str, data=None
) -> Generator[bytes, None, None]:
"""Create MCP error response"""
error_data = ErrorData(code=code, message=message, data=data)
json_response = JSONRPCError(

Loading…
Откажи
Сачувај