From b7be5c8c82d1860f3e9a952aac250d19b462d505 Mon Sep 17 00:00:00 2001
From: crazywoola <427733928@qq.com>
Date: Wed, 24 Dec 2025 16:29:07 +0800
Subject: [PATCH] feat: add MCP tools

---
 api/core/helper/ssrf_proxy.py                 |  9 +++++---
 .../unit_tests/core/helper/test_ssrf_proxy.py | 21 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/api/core/helper/ssrf_proxy.py b/api/core/helper/ssrf_proxy.py
index 0b36969cf9..e390873bf4 100644
--- a/api/core/helper/ssrf_proxy.py
+++ b/api/core/helper/ssrf_proxy.py
@@ -106,6 +106,9 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
     verify_option = kwargs.pop("ssl_verify", dify_config.HTTP_REQUEST_NODE_SSL_VERIFY)
     client = _get_ssrf_client(verify_option)
 
+    # Extract follow_redirects for client.send() - it's not a build_request parameter
+    follow_redirects = kwargs.pop("follow_redirects", True)
+
     # Preserve user-provided Host header
     # When using a forward proxy, httpx may override the Host header based on the URL.
     # We extract and preserve any explicitly set Host header to support virtual hosting.
@@ -120,9 +123,9 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
             # the request API to explicitly set headers before sending
             headers = {k: v for k, v in headers.items() if k.lower() != "host"}
             if user_provided_host is not None:
-                headers["host"] = user_provided_host
-            kwargs["headers"] = headers
-            response = client.request(method=method, url=url, **kwargs)
+                request.headers["Host"] = user_provided_host
+
+            response = client.send(request, follow_redirects=follow_redirects)
 
             # Check for SSRF protection by Squid proxy
             if response.status_code in (401, 403):
diff --git a/api/tests/unit_tests/core/helper/test_ssrf_proxy.py b/api/tests/unit_tests/core/helper/test_ssrf_proxy.py
index beae1d0358..7dafdce249 100644
--- a/api/tests/unit_tests/core/helper/test_ssrf_proxy.py
+++ b/api/tests/unit_tests/core/helper/test_ssrf_proxy.py
@@ -107,3 +107,24 @@ def test_host_header_preservation_with_user_header(mock_get_client):
     response = make_request("GET", "http://example.com", headers={"Host": custom_host})
 
     assert response.status_code == 200
+    # Verify build_request was called
+    mock_client.build_request.assert_called_once()
+    # Verify the Host header was set on the request object
+    assert mock_request.headers.get("Host") == custom_host
+    mock_client.send.assert_called_once_with(mock_request, follow_redirects=True)
+
+
+@patch("core.helper.ssrf_proxy._get_ssrf_client")
+@pytest.mark.parametrize("host_key", ["host", "HOST"])
+def test_host_header_preservation_case_insensitive(mock_get_client, host_key):
+    """Test that Host header is preserved regardless of case."""
+    mock_client = MagicMock()
+    mock_request = MagicMock()
+    mock_request.headers = {}
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_client.send.return_value = mock_response
+    mock_client.build_request.return_value = mock_request
+    mock_get_client.return_value = mock_client
+    response = make_request("GET", "http://example.com", headers={host_key: "api.example.com"})
+    assert mock_request.headers.get("Host") == "api.example.com"