Merge 14ab1d9100 into 41336b09e2

litellm/proxy/proxy_server.py

@@ -18576,6 +18576,30 @@ async def login(request: Request):
     return redirect_response
 
 
+async def _resume_return_to_url(request: Request, jwt_token: str) -> str | None:
+    """Where a JSON sign-in should send the browser for the one-shot ``litellm_cp_return_to`` cookie,
+    or None to land on the dashboard. Goes through the same resumer as /login and the SSO callback,
+    so the two shapes it accepts (same-origin path, control-plane URL) stay in sync across every
+    sign-in route. Never raises: a stale cookie must not fail a completed sign-in."""
+    from litellm.proxy.management_endpoints.ui_sso import read_return_to_cookie, sso_return_to_redirect
+
+    return_to: Final = read_return_to_cookie(request)
+    if not return_to:
+        return None
+    try:
+        resumed: Final = await sso_return_to_redirect(
+            return_to=return_to,
+            jwt_token=jwt_token,
+            redis_usage_cache=redis_usage_cache,
+            user_api_key_cache=user_api_key_cache,
+            request=request,
+        )
+    except Exception:  # noqa: BLE001  # resuming must NEVER block a completed sign-in
+        verbose_proxy_logger.info("Ignoring stale litellm_cp_return_to cookie; landing on dashboard")
+        return None
+    return None if resumed is None else resumed.headers.get("location")
+
+
 @router.post("/v2/login", include_in_schema=False)  # hidden helper for UI logins via API
 async def login_v2(request: Request):
     global premium_user, general_settings, master_key
@@ -18609,14 +18633,21 @@ async def login_v2(request: Request):
         litellm_dashboard_ui = litellm_dashboard_ui.rstrip("/")
         litellm_dashboard_ui += "/ui?login=success"
 
+        # The deprecated /sso/key/generate form points users at this login page (/ui/login/), so a
+        # return_to it preserved (e.g. the MCP gateway's /authorize round-trip) has to be honored here
+        # too, or the native client's loopback callback never gets its code.
+        redirect_url: Final = await _resume_return_to_url(request=request, jwt_token=jwt_token) or litellm_dashboard_ui
+
         # Token is included in the response body so the UI can set a JS-accessible
         # cookie even when a reverse proxy (e.g. nginx-ingress) adds HttpOnly to the
         # server-set cookie, which would otherwise cause an infinite login redirect.
         json_response: Final = JSONResponse(
-            content={"redirect_url": litellm_dashboard_ui, "token": jwt_token},
+            content={"redirect_url": redirect_url, "token": jwt_token},
             status_code=status.HTTP_200_OK,
         )
         set_session_token_cookie(json_response, request, jwt_token)
+        if "litellm_cp_return_to" in request.cookies:
+            json_response.delete_cookie(key="litellm_cp_return_to")
         return json_response
     except Exception as e:
         verbose_proxy_logger.exception("litellm.proxy.proxy_server.login_v2(): Exception occurred - %s", e)

tests/unit/proxy/proxy_server/test_routes_login_sso.py

@@ -499,6 +499,69 @@ def test_login_form_ignores_open_redirect_return_to(client, monkeypatch):
     assert "/ui" in location  # dashboard fallback
 
 
+def test_login_v2_honors_same_origin_return_to_cookie(client, monkeypatch):
+    """/sso/key/generate's deprecation banner sends password users to /ui/login/, which signs in via
+    /v2/login. A return_to the banner page preserved (the MCP gateway /authorize round-trip) must come
+    back as redirect_url, or the native client's loopback callback never gets its code (#48090)."""
+    _install_login_mocks(monkeypatch)
+    return_to = "/authorize?response_type=code&client_id=llm_dcrc_abc&redirect_uri=http%3A%2F%2F127.0.0.1%3A5555%2F"
+    response = client.post(
+        "/v2/login",
+        json={"username": "admin", "password": "password"},
+        cookies={"litellm_cp_return_to": _sealed_return_to(return_to)},
+    )
+    set_cookie = response.headers.get("set-cookie", "")
+    assert response.status_code == 200
+    assert response.json()["redirect_url"] == return_to
+    assert "token=" in set_cookie
+    assert 'litellm_cp_return_to=""' in set_cookie  # one-shot: cleared after use
+
+
+def test_login_v2_honors_control_plane_return_to_cookie(client, monkeypatch):
+    """/v2/login resumes through the same resumer as /login, so a control-plane return_to gets the
+    one-time login code instead of the dashboard."""
+    import litellm.proxy.proxy_server as ps
+
+    _install_login_mocks(monkeypatch)
+    monkeypatch.setitem(ps.general_settings, "control_plane_url", "https://cp.example.com")
+    response = client.post(
+        "/v2/login",
+        json={"username": "admin", "password": "password"},
+        cookies={"litellm_cp_return_to": _sealed_return_to("https://cp.example.com/console")},
+    )
+    redirect_url = response.json()["redirect_url"]
+    assert response.status_code == 200
+    assert redirect_url.startswith("https://cp.example.com/console")
+    assert "code=" in redirect_url and "login=success" in redirect_url
+
+
+@pytest.mark.parametrize(
+    "cookie",
+    [
+        pytest.param("/authorize?client_id=llm_dcrc_abc", id="plaintext-not-sealed-here"),
+        pytest.param("sealed:https://evil.example.com/steal", id="open-redirect"),
+        pytest.param("sealed:https://old-cp.example.com/console", id="stale-control-plane"),
+    ],
+)
+def test_login_v2_falls_back_to_dashboard_for_unusable_return_to(client, monkeypatch, cookie):
+    """A return_to the resumer won't honor never fails the sign-in: /v2/login lands on the dashboard
+    and clears the cookie."""
+    import litellm.proxy.proxy_server as ps
+
+    _install_login_mocks(monkeypatch)
+    monkeypatch.setitem(ps.general_settings, "control_plane_url", "https://new-cp.example.com")
+    if cookie.startswith("sealed:"):
+        cookie = _sealed_return_to(cookie.removeprefix("sealed:"))
+    response = client.post(
+        "/v2/login",
+        json={"username": "admin", "password": "password"},
+        cookies={"litellm_cp_return_to": cookie},
+    )
+    assert response.status_code == 200
+    assert response.json()["redirect_url"].endswith("/ui?login=success")
+    assert 'litellm_cp_return_to=""' in response.headers.get("set-cookie", "")
+
+
 # ---------------------------------------------------------------------------
 # Failed-login accounting across the login routes (LIT-5285)
 # ---------------------------------------------------------------------------