refac

backend/open_webui/main.py

@@ -482,9 +482,11 @@ from open_webui.utils.auth import (
 )
 from open_webui.utils.plugin import install_tool_and_function_dependencies
 from open_webui.utils.oauth import (
+    get_oauth_client_info_with_dynamic_client_registration,
+    encrypt_data,
+    decrypt_data,
     OAuthManager,
     OAuthClientManager,
-    decrypt_data,
     OAuthClientInformationFull,
 )
 from open_webui.utils.security_headers import SecurityHeadersMiddleware
@@ -1987,6 +1989,64 @@ except Exception as e:
     )
 
 
+async def register_client(self, request, client_id: str) -> bool:
+    server_type, server_id = client_id.split(":", 1)
+
+    connection = None
+    connection_idx = None
+
+    for idx, conn in enumerate(request.app.state.config.TOOL_SERVER_CONNECTIONS or []):
+        if conn.get("type", "openapi") == server_type:
+            info = conn.get("info", {})
+            if info.get("id") == server_id:
+                connection = conn
+                connection_idx = idx
+                break
+
+    if connection is None or connection_idx is None:
+        log.warning(
+            f"Unable to locate MCP tool server configuration for client {client_id} during re-registration"
+        )
+        return False
+
+    server_url = connection.get("url")
+    oauth_server_key = (connection.get("config") or {}).get("oauth_server_key")
+
+    try:
+        oauth_client_info = (
+            await get_oauth_client_info_with_dynamic_client_registration(
+                request,
+                client_id,
+                server_url,
+                oauth_server_key,
+            )
+        )
+    except Exception as e:
+        log.error(f"Dynamic client re-registration failed for {client_id}: {e}")
+        return False
+
+    try:
+        request.app.state.config.TOOL_SERVER_CONNECTIONS[connection_idx] = {
+            **connection,
+            "info": {
+                **connection.get("info", {}),
+                "oauth_client_info": encrypt_data(
+                    oauth_client_info.model_dump(mode="json")
+                ),
+            },
+        }
+    except Exception as e:
+        log.error(
+            f"Failed to persist updated OAuth client info for tool server {client_id}: {e}"
+        )
+        return False
+
+    oauth_client_manager.remove_client(client_id)
+    oauth_client_manager.add_client(client_id, oauth_client_info)
+    log.info(f"Re-registered OAuth client {client_id} for tool server")
+    return True
+
+
 @app.get("/oauth/clients/{client_id}/authorize")
 async def oauth_client_authorize(
     client_id: str,
@@ -1994,6 +2054,41 @@ async def oauth_client_authorize(
     response: Response,
     user=Depends(get_verified_user),
 ):
+    # ensure_valid_client_registration
+    client = oauth_client_manager.get_client(client_id)
+    client_info = oauth_client_manager.get_client_info(client_id)
+    if client is None or client_info is None:
+        raise HTTPException(status.HTTP_404_NOT_FOUND)
+
+    if not await oauth_client_manager._preflight_authorization_url(client, client_info):
+        log.info(
+            "Detected invalid OAuth client %s; attempting re-registration",
+            client_id,
+        )
+
+        re_registered = await register_client(request, client_id)
+        if not re_registered:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Failed to re-register OAuth client",
+            )
+
+        client = oauth_client_manager.get_client(client_id)
+        client_info = oauth_client_manager.get_client_info(client_id)
+        if client is None or client_info is None:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="OAuth client unavailable after re-registration",
+            )
+
+        if not await oauth_client_manager._preflight_authorization_url(
+            client, client_info
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="OAuth client registration is still invalid after re-registration",
+            )
+
     return await oauth_client_manager.handle_authorize(request, client_id=client_id)