Adding headers to HttpRequestResponse when making http_request

Author: sunishsheth2009

databricks/sdk/mixins/open_ai_client.py

@@ -4,8 +4,8 @@
 from requests import Response
 
 from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
-                                            ServingEndpointsAPI)
-
+                                            ServingEndpointsAPI,
+                                            HttpRequestResponse)
 
 class ServingEndpointsExt(ServingEndpointsAPI):
 
@@ -88,15 +88,30 @@ def http_request(
         """
         response = Response()
         response.status_code = 200
-        server_response = super().http_request(
-            connection_name=conn,
-            method=method,
-            path=path,
-            headers=js.dumps(headers) if headers is not None else None,
-            json=js.dumps(json) if json is not None else None,
-            params=js.dumps(params) if params is not None else None,
+
+        # We currently don't call super.http_request because we need to pass in response_headers
+        # This is a temporary fix to get the headers we need for the MCP session id
+        # TODO: Remove this once we have a better way to get back the response headers
+        headers_to_capture = ["mcp-session-id"]
+        res = self._api.do(
+            "POST",
+            "/api/2.0/external-function",
+            body={
+                "connection_name": conn,
+                "method": method.value,
+                "path": path,
+                "headers": js.dumps(headers) if headers is not None else None,
+                "json": js.dumps(json) if json is not None else None,
+                "params": js.dumps(params) if params is not None else None,
+            },
+            headers={"Accept": "text/plain", "Content-Type": "application/json"},
+            raw=True,
+            response_headers=headers_to_capture
         )
 
+        # Create HttpRequestResponse from the raw response
+        server_response = HttpRequestResponse.from_dict(res)
+
         # Read the content from the HttpRequestResponse object
         if hasattr(server_response, "contents") and hasattr(server_response.contents, "read"):
             raw_content = server_response.contents.read()  # Read the bytes
@@ -109,4 +124,9 @@ def http_request(
         else:
             raise ValueError("Contents must be bytes.")
 
+        # Copy headers from raw response to Response
+        for header_name in headers_to_capture:
+            if header_name in res:
+                response.headers[header_name] = res[header_name]
+
         return response

tests/test_open_ai_mixin.py

@@ -37,6 +37,7 @@ def test_http_request(w, requests_mock):
         "Accept": "text/plain",
         "Content-Type": "application/json",
     }
+    headers_to_capture = ["mcp-session-id"]
     mocked_url = "http://localhost/api/2.0/external-function"
     blob_response = BytesIO(b"The request was successful")
 
@@ -45,6 +46,7 @@ def test_http_request(w, requests_mock):
         request_headers=headers,
         content=blob_response.getvalue(),
         status_code=200,
+        response_headers=headers_to_capture,
     )
     response = w.serving_endpoints.http_request(
         conn="test_conn",
@@ -55,3 +57,4 @@ def test_http_request(w, requests_mock):
     assert requests_mock.called
     assert response.status_code == 200  # Verify the response status
     assert response.text == "The request was successful"  # Ensure the response body matches the mocked data
+    assert response.headers["mcp-session-id"] is not None