From b5bb6ae177dcdc4e8274d7e5ffa50bc8099fd466 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Mon, 8 Sep 2025 18:50:23 +0400
Subject: [PATCH] refac

---
 backend/open_webui/functions.py        | 10 ++++++++++
 backend/open_webui/main.py             |  8 ++++++++
 backend/open_webui/routers/users.py    | 14 ++++++++++++++
 backend/open_webui/utils/middleware.py | 13 ++++++++++++-
 backend/open_webui/utils/tools.py      | 12 +-----------
 5 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/backend/open_webui/functions.py b/backend/open_webui/functions.py
index db367ccbd0..4122cbbe0d 100644
--- a/backend/open_webui/functions.py
+++ b/backend/open_webui/functions.py
@@ -219,6 +219,15 @@ async def generate_function_chat_completion(
         __task__ = metadata.get("task", None)
         __task_body__ = metadata.get("task_body", None)
 
+    oauth_token = None
+    try:
+        oauth_token = request.app.state.oauth_manager.get_oauth_token(
+            user.id,
+            request.cookies.get("oauth_session_id", None),
+        )
+    except Exception as e:
+        log.error(f"Error getting OAuth token: {e}")
+
     extra_params = {
         "__event_emitter__": __event_emitter__,
         "__event_call__": __event_call__,
@@ -230,6 +239,7 @@ async def generate_function_chat_completion(
         "__files__": files,
         "__user__": user.model_dump() if isinstance(user, UserModel) else {},
         "__metadata__": metadata,
+        "__oauth_token__": oauth_token,
         "__request__": request,
     }
     extra_params["__tools__"] = await get_tools(
diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index de7dcae086..ea60900c9c 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -1408,6 +1408,14 @@ async def chat_completion(
     model_item = form_data.pop("model_item", {})
     tasks = form_data.pop("background_tasks", None)
 
+    oauth_token = None
+    try:
+        oauth_token = request.app.state.oauth_manager.get_oauth_token(
+            user.id, request.cookies.get("oauth_session_id", None)
+        )
+    except Exception as e:
+        log.error(f"Error getting OAuth token: {e}")
+
     metadata = {}
     try:
         if not model_item.get("direct", False):
diff --git a/backend/open_webui/routers/users.py b/backend/open_webui/routers/users.py
index 4d2539a18e..5b331dce73 100644
--- a/backend/open_webui/routers/users.py
+++ b/backend/open_webui/routers/users.py
@@ -10,6 +10,8 @@ from pydantic import BaseModel
 
 
 from open_webui.models.auths import Auths
+from open_webui.models.oauth_sessions import OAuthSessions
+
 from open_webui.models.groups import Groups
 from open_webui.models.chats import Chats
 from open_webui.models.users import (
@@ -340,6 +342,18 @@ async def get_user_by_id(user_id: str, user=Depends(get_verified_user)):
         )
 
 
+@router.get("/{user_id}/oauth/sessions", response_model=Optional[dict])
+async def get_user_oauth_sessions_by_id(user_id: str, user=Depends(get_admin_user)):
+    sessions = OAuthSessions.get_sessions_by_user_id(user_id)
+    if sessions and len(sessions) > 0:
+        return sessions
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=ERROR_MESSAGES.USER_NOT_FOUND,
+        )
+
+
 ############################
 # GetUserProfileImageById
 ############################
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 406ee1b16f..ae2c96c6da 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -818,7 +818,8 @@ async def process_chat_payload(request, form_data, user, metadata, model):
     oauth_token = None
     try:
         oauth_token = request.app.state.oauth_manager.get_oauth_token(
-            user.id, request.cookies.get("oauth_session_id", None)
+            user.id,
+            request.cookies.get("oauth_session_id", None),
         )
     except Exception as e:
         log.error(f"Error getting OAuth token: {e}")
@@ -1493,11 +1494,21 @@ async def process_chat_response(
     ):
         return response
 
+    oauth_token = None
+    try:
+        oauth_token = request.app.state.oauth_manager.get_oauth_token(
+            user.id,
+            request.cookies.get("oauth_session_id", None),
+        )
+    except Exception as e:
+        log.error(f"Error getting OAuth token: {e}")
+
     extra_params = {
         "__event_emitter__": event_emitter,
         "__event_call__": event_caller,
         "__user__": user.model_dump() if isinstance(user, UserModel) else {},
         "__metadata__": metadata,
+        "__oauth_token__": oauth_token,
         "__request__": request,
         "__model__": model,
     }
diff --git a/backend/open_webui/utils/tools.py b/backend/open_webui/utils/tools.py
index efdb7c85b1..f7e7f7acef 100644
--- a/backend/open_webui/utils/tools.py
+++ b/backend/open_webui/utils/tools.py
@@ -134,17 +134,7 @@ async def get_tools(
                         )
                     elif auth_type == "oauth":
                         cookies = request.cookies
-                        oauth_token = None
-                        try:
-                            oauth_token = (
-                                request.app.state.oauth_manager.get_oauth_token(
-                                    user.id,
-                                    request.cookies.get("oauth_session_id", None),
-                                )
-                            )
-                        except Exception as e:
-                            log.error(f"Error getting OAuth token: {e}")
-
+                        oauth_token = extra_params.get("__oauth_token__", None)
                         headers["Authorization"] = (
                             f"Bearer {oauth_token.get('access_token', '')}"
                         )