handle redirect cookie association

Author: mayank-at-sauce

py/selenium/webdriver/common/api_request_context.py

@@ -231,6 +231,22 @@ def _get_set_cookie_headers(resp: urllib3.BaseHTTPResponse) -> list[str]:
     return [sc] if sc else []
 
 
+def _resolve_redirect_url(resp: urllib3.BaseHTTPResponse, original_url: str) -> str:
+    """Return the final URL after any redirects.
+
+    urllib3's retry history records each hop.  When redirects occurred,
+    the last entry's redirect_location resolved against its URL gives
+    the final destination.  When no redirects occurred, the original
+    request URL is returned unchanged.
+    """
+    history = resp.retries.history if resp.retries else ()
+    if history:
+        last = history[-1]
+        if last.url and last.redirect_location:
+            return urllib.parse.urljoin(last.url, last.redirect_location)
+    return original_url
+
+
 class _BaseRequestContext:
     """Base class with shared HTTP request logic for API request contexts."""
 
@@ -473,12 +489,16 @@ def _fetch(self, url: str, method: str, **kwargs: Any) -> APIResponse:
         url = self._append_params(url, kwargs)
         resp = self._execute_request(method, url, headers, body, kwargs)
 
+        # After redirects, associate cookies with the final destination's
+        # origin, not the initial request URL.
+        final_url = _resolve_redirect_url(resp, url)
+
         # Process response cookies
         set_cookie_headers = _get_set_cookie_headers(resp)
         if set_cookie_headers:
-            self._handle_response_cookies(set_cookie_headers, url)
+            self._handle_response_cookies(set_cookie_headers, final_url)
 
-        response = self._build_response(resp, url)
+        response = self._build_response(resp, final_url)
 
         fail = kwargs.get("fail_on_status_code", self._fail_on_status_code)
         if fail and not response.ok:

py/test/unit/selenium/webdriver/common/api_request_context_tests.py

@@ -98,6 +98,10 @@ def do_GET(self):
             self.send_response(302)
             self.send_header("Location", "/ok")
             self.end_headers()
+        elif path == "/redirect_with_cookies":
+            self.send_response(302)
+            self.send_header("Location", "/set_cookie?name=redirected&value=yes")
+            self.end_headers()
         elif path == "/redirect_chain":
             import urllib.parse
 
@@ -1474,6 +1478,22 @@ def test_e2e_redirect_chain_exceeds_limit(base_url):
     ctx.dispose()
 
 
+def test_e2e_redirect_cookies_associated_with_final_url(base_url):
+    """Cookies from a redirected response must be associated with the final URL's origin."""
+    ctx = _IsolatedAPIRequestContext(base_url=base_url)
+    r = ctx.get("/redirect_with_cookies")
+    assert r.status == 200
+    assert r.text() == "cookie set"
+    # The response URL should be the final destination, not the redirect source
+    assert "/set_cookie" in r.url
+    assert "/redirect_with_cookies" not in r.url
+    # The cookie should be stored with the correct domain (from the final URL)
+    assert len(ctx._cookies) == 1
+    assert ctx._cookies[0]["name"] == "redirected"
+    assert ctx._cookies[0]["value"] == "yes"
+    ctx.dispose()
+
+
 def test_e2e_cookie_set_and_sent(base_url):
     ctx = _IsolatedAPIRequestContext(base_url=base_url)
     # Server sets a cookie