refac

backend/open_webui/env.py

@@ -806,6 +806,36 @@ else:
         AIOHTTP_CLIENT_TIMEOUT_TOOL_SERVER = AIOHTTP_CLIENT_TIMEOUT
 
 
+####################################
+# AIOHTTP Connection Pool
+####################################
+
+AIOHTTP_POOL_CONNECTIONS = os.environ.get('AIOHTTP_POOL_CONNECTIONS', '')
+if AIOHTTP_POOL_CONNECTIONS == '':
+    AIOHTTP_POOL_CONNECTIONS = None
+else:
+    try:
+        AIOHTTP_POOL_CONNECTIONS = int(AIOHTTP_POOL_CONNECTIONS)
+    except ValueError:
+        AIOHTTP_POOL_CONNECTIONS = None
+
+AIOHTTP_POOL_CONNECTIONS_PER_HOST = os.environ.get('AIOHTTP_POOL_CONNECTIONS_PER_HOST', '')
+if AIOHTTP_POOL_CONNECTIONS_PER_HOST == '':
+    AIOHTTP_POOL_CONNECTIONS_PER_HOST = None
+else:
+    try:
+        AIOHTTP_POOL_CONNECTIONS_PER_HOST = int(AIOHTTP_POOL_CONNECTIONS_PER_HOST)
+    except ValueError:
+        AIOHTTP_POOL_CONNECTIONS_PER_HOST = None
+
+AIOHTTP_POOL_DNS_TTL = os.environ.get('AIOHTTP_POOL_DNS_TTL', '300')
+try:
+    AIOHTTP_POOL_DNS_TTL = int(AIOHTTP_POOL_DNS_TTL)
+    if AIOHTTP_POOL_DNS_TTL < 0:
+        AIOHTTP_POOL_DNS_TTL = 300
+except ValueError:
+    AIOHTTP_POOL_DNS_TTL = 300
+
 RAG_EMBEDDING_TIMEOUT = os.environ.get('RAG_EMBEDDING_TIMEOUT', '')
 
 if RAG_EMBEDDING_TIMEOUT == '':

backend/open_webui/main.py

@@ -719,6 +719,10 @@ async def lifespan(app: FastAPI):
 
     yield
 
+    # Shutdown: clean up shared resources
+    from open_webui.utils.session_pool import close_session
+    await close_session()
+
     if hasattr(app.state, 'redis_task_command_listener'):
         app.state.redis_task_command_listener.cancel()
 

backend/open_webui/routers/ollama.py

@@ -50,7 +50,10 @@ from open_webui.models.groups import Groups
 from open_webui.utils.access_control import check_model_access
 from open_webui.utils.misc import (
     calculate_sha256,
+)
+from open_webui.utils.session_pool import (
     cleanup_response,
+    get_session,
     stream_wrapper,
 )
 from open_webui.utils.payload import (
@@ -122,10 +125,7 @@ async def send_request(
     r = None
     streaming = False
     try:
-        session = aiohttp.ClientSession(
+        session = await get_session()
-            trust_env=True,
-            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
-        )
 
         headers = {
             'Content-Type': 'application/json',
@@ -140,6 +140,7 @@ async def send_request(
         r = await session.request(
             method, url, data=payload, headers=headers,
             ssl=AIOHTTP_CLIENT_SESSION_SSL,
+            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
         )
 
         if not r.ok:
@@ -165,7 +166,7 @@ async def send_request(
 
             streaming = True
             return StreamingResponse(
-                stream_wrapper(r, session),
+                stream_wrapper(r),
                 status_code=r.status,
                 headers=response_headers,
             )
@@ -184,7 +185,7 @@ async def send_request(
         )
     finally:
         if not streaming:
-            await cleanup_response(r, session)
+            await cleanup_response(r)
 
 
 def get_api_key(idx, url, configs):

backend/open_webui/routers/openai.py

@@ -52,9 +52,12 @@ from open_webui.utils.payload import (
     apply_system_prompt_to_body,
 )
 from open_webui.utils.misc import (
-    cleanup_response,
     convert_logit_bias_input_to_json,
     stream_chunks_handler,
+)
+from open_webui.utils.session_pool import (
+    cleanup_response,
+    get_session,
     stream_wrapper,
 )
 
@@ -1174,12 +1177,11 @@ async def generate_chat_completion(
     payload = json.dumps(payload)
 
     r = None
-    session = None
     streaming = False
     response = None
 
     try:
-        session = aiohttp.ClientSession(trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT))
+        session = await get_session()
 
         r = await session.request(
             method='POST',
@@ -1188,13 +1190,14 @@ async def generate_chat_completion(
             headers=headers,
             cookies=cookies,
             ssl=AIOHTTP_CLIENT_SESSION_SSL,
+            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
         )
 
         # Check if response is SSE
         if 'text/event-stream' in r.headers.get('Content-Type', ''):
             streaming = True
             return StreamingResponse(
-                stream_wrapper(r, session, stream_chunks_handler),
+                stream_wrapper(r, content_handler=stream_chunks_handler),
                 status_code=r.status,
                 headers=dict(r.headers),
             )
@@ -1225,7 +1228,7 @@ async def generate_chat_completion(
         )
     finally:
         if not streaming:
-            await cleanup_response(r, session)
+            await cleanup_response(r)
 
 
 async def embeddings(request: Request, form_data: dict, user):
@@ -1261,27 +1264,24 @@ async def embeddings(request: Request, form_data: dict, user):
     )
 
     r = None
-    session = None
     streaming = False
 
     headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user)
     try:
-        session = aiohttp.ClientSession(
+        session = await get_session()
-            trust_env=True,
-            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
-        )
         r = await session.request(
             method='POST',
             url=f'{url}/embeddings',
             data=body,
             headers=headers,
             cookies=cookies,
+            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
         )
 
         if 'text/event-stream' in r.headers.get('Content-Type', ''):
             streaming = True
             return StreamingResponse(
-                stream_wrapper(r, session),
+                stream_wrapper(r),
                 status_code=r.status,
                 headers=dict(r.headers),
             )
@@ -1306,7 +1306,7 @@ async def embeddings(request: Request, form_data: dict, user):
         )
     finally:
         if not streaming:
-            await cleanup_response(r, session)
+            await cleanup_response(r)
 
 
 class ResponsesForm(BaseModel):
@@ -1365,7 +1365,6 @@ async def responses(
     )
 
     r = None
-    session = None
     streaming = False
 
     try:
@@ -1388,10 +1387,7 @@ async def responses(
         else:
             request_url = f'{url}/responses'
 
-        session = aiohttp.ClientSession(
+        session = await get_session()
-            trust_env=True,
-            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
-        )
         r = await session.request(
             method='POST',
             url=request_url,
@@ -1399,13 +1395,14 @@ async def responses(
             headers=headers,
             cookies=cookies,
             ssl=AIOHTTP_CLIENT_SESSION_SSL,
+            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
         )
 
         # Check if response is SSE
         if 'text/event-stream' in r.headers.get('Content-Type', ''):
             streaming = True
             return StreamingResponse(
-                stream_wrapper(r, session),
+                stream_wrapper(r),
                 status_code=r.status,
                 headers=dict(r.headers),
             )
@@ -1433,7 +1430,7 @@ async def responses(
         )
     finally:
         if not streaming:
-            await cleanup_response(r, session)
+            await cleanup_response(r)
 
 
 @router.api_route('/{path:path}', methods=['GET', 'POST', 'PUT', 'DELETE'])
@@ -1479,7 +1476,6 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
     )
 
     r = None
-    session = None
     streaming = False
 
     try:
@@ -1508,10 +1504,7 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
         else:
             request_url = f'{url}/{path}'
 
-        session = aiohttp.ClientSession(
+        session = await get_session()
-            trust_env=True,
-            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
-        )
         r = await session.request(
             method=request.method,
             url=request_url,
@@ -1519,13 +1512,14 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
             headers=headers,
             cookies=cookies,
             ssl=AIOHTTP_CLIENT_SESSION_SSL,
+            timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT),
         )
 
         # Check if response is SSE
         if 'text/event-stream' in r.headers.get('Content-Type', ''):
             streaming = True
             return StreamingResponse(
-                stream_wrapper(r, session),
+                stream_wrapper(r),
                 status_code=r.status,
                 headers=dict(r.headers),
             )
@@ -1553,4 +1547,4 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
         )
     finally:
         if not streaming:
-            await cleanup_response(r, session)
+            await cleanup_response(r)

backend/open_webui/utils/middleware.py

@@ -4060,11 +4060,12 @@ async def streaming_chat_response_handler(response, ctx):
                         if responses_api_tool_calls:
                             tool_calls.append(_split_tool_calls(responses_api_tool_calls))
 
+                try:
+                    await stream_body_handler(response, form_data)
+                finally:
                     if response.background:
                         await response.background()
 
-                await stream_body_handler(response, form_data)
-
                 tool_call_retries = 0
                 tool_call_sources = []  # Track citation sources from tool results
                 all_tool_call_sources = []  # Accumulated sources across all iterations

backend/open_webui/utils/misc.py

@@ -905,9 +905,11 @@ async def cleanup_response(
     session: Optional[aiohttp.ClientSession],
 ):
     if response:
-        response.close()
+        if not response.closed:
+            await response.close()
     if session:
-        await session.close()
+        if not session.closed:
+            await session.close()
 
 
 async def stream_wrapper(response, session, content_handler=None):

backend/open_webui/utils/session_pool.py

@@ -0,0 +1,114 @@
+"""Shared aiohttp ClientSession pool.
+
+Instead of creating a new ClientSession (and TCPConnector) per request,
+callers acquire a long-lived session from this module.  The pool manages
+a single TCPConnector with configurable limits, enabling TCP/SSL connection
+reuse, shared DNS cache, and bounded concurrency.
+
+All pool parameters are configurable via environment variables:
+    - AIOHTTP_POOL_CONNECTIONS (default 100) — max total connections
+    - AIOHTTP_POOL_CONNECTIONS_PER_HOST (default 30) — per-host limit
+    - AIOHTTP_POOL_DNS_TTL (default 300) — DNS cache TTL in seconds
+
+Usage:
+    from open_webui.utils.session_pool import get_session, cleanup_response
+
+    session = await get_session()
+    r = await session.request(...)
+    # When done with the *response* (not the session):
+    await cleanup_response(r)
+
+IMPORTANT: Callers must NOT close the shared session.  Only the response
+needs cleanup.  The session is closed once during application shutdown
+via ``close_session()``.
+"""
+
+import logging
+from typing import Optional
+
+import aiohttp
+
+from open_webui.env import (
+    AIOHTTP_CLIENT_TIMEOUT,
+    AIOHTTP_POOL_CONNECTIONS,
+    AIOHTTP_POOL_CONNECTIONS_PER_HOST,
+    AIOHTTP_POOL_DNS_TTL,
+)
+
+log = logging.getLogger(__name__)
+
+_session: Optional[aiohttp.ClientSession] = None
+
+
+async def get_session() -> aiohttp.ClientSession:
+    """Return the shared aiohttp ClientSession, creating it lazily."""
+    global _session
+    if _session is None or _session.closed:
+        connector_kwargs = {
+            'ttl_dns_cache': AIOHTTP_POOL_DNS_TTL,
+            'enable_cleanup_closed': True,
+        }
+        if AIOHTTP_POOL_CONNECTIONS is not None:
+            connector_kwargs['limit'] = AIOHTTP_POOL_CONNECTIONS
+        else:
+            connector_kwargs['limit'] = 0  # aiohttp: 0 = unlimited
+        if AIOHTTP_POOL_CONNECTIONS_PER_HOST is not None:
+            connector_kwargs['limit_per_host'] = AIOHTTP_POOL_CONNECTIONS_PER_HOST
+        else:
+            connector_kwargs['limit_per_host'] = 0  # aiohttp: 0 = unlimited
+        connector = aiohttp.TCPConnector(**connector_kwargs)
+        timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
+        _session = aiohttp.ClientSession(
+            connector=connector,
+            timeout=timeout,
+            trust_env=True,
+        )
+        log.info(
+            'Created shared aiohttp session pool '
+            '(limit=%s, per_host=%s, dns_ttl=%d)',
+            AIOHTTP_POOL_CONNECTIONS or 'unlimited',
+            AIOHTTP_POOL_CONNECTIONS_PER_HOST or 'unlimited',
+            AIOHTTP_POOL_DNS_TTL,
+        )
+    return _session
+
+
+async def close_session():
+    """Close the shared session.  Called during application shutdown."""
+    global _session
+    if _session and not _session.closed:
+        await _session.close()
+        log.info('Closed shared aiohttp session pool')
+        _session = None
+
+
+async def cleanup_response(
+    response: Optional[aiohttp.ClientResponse],
+    session: Optional[aiohttp.ClientSession] = None,
+):
+    """Release and close an aiohttp response, optionally closing the session.
+
+    When using the shared pool, ``session`` should be ``None`` (the pool
+    session is never closed per-request).  When a caller creates its own
+    one-off session, pass it here to close it after the response.
+    """
+    if response:
+        if not response.closed:
+            await response.close()
+    if session:
+        if not session.closed:
+            await session.close()
+
+
+async def stream_wrapper(response, session=None, content_handler=None):
+    """Wrap a stream to ensure cleanup happens even if streaming is interrupted.
+
+    This is more reliable than BackgroundTask which may not run if the client
+    disconnects.  When using the shared pool, ``session`` should be ``None``.
+    """
+    try:
+        stream = content_handler(response.content) if content_handler else response.content
+        async for chunk in stream:
+            yield chunk
+    finally:
+        await cleanup_response(response, session)