refac

backend/open_webui/models/auth_sessions.py

@@ -89,7 +89,7 @@ class AuthSessionsTable:
             return [
                 {
                     **{key: getattr(row, key) for key in ('id', 'created_at', 'updated_at', 'expires_at')},
-                    'connection_id': row.data['connection_id'],
+                    'connection_id': row.data.get('connection_id'),
                     'name': row.data['name'],
                 }
                 for row in rows

backend/open_webui/routers/sessions.py

@@ -33,7 +33,7 @@ PROOF_KEYS = ('auth_time', 'auth_method', 'session_stamp', 'mfa_verified')
 
 class SessionForm(BaseModel):
     model_config = ConfigDict(extra='forbid', str_strip_whitespace=True)
-    connection_id: UUID
+    connection_id: UUID | None = None
     name: str = Field(min_length=1, max_length=100)
 
 
@@ -58,7 +58,6 @@ class AuthorizationRequest(SessionForm):
             or parsed.password
             or parsed.query
             or parsed.fragment
-            or not parsed.path.startswith('/callback/')
         ):
             raise ValueError('A loopback callback is required')
         return value
@@ -148,7 +147,7 @@ async def create_session(form: SessionForm, request: Request, user=Depends(get_h
     proof = await verify_proof(request.state.claims, request)
     row = await AuthSessions.create(
         user_id=user.id,
-        data={'connection_id': str(form.connection_id), 'name': form.name},
+        data=form.model_dump(mode='json', exclude_none=True),
         expires_at=await idle_expiry(),
         refresh_token_id=str(uuid4()),
     )
@@ -160,8 +159,7 @@ async def authorization_request(form: AuthorizationRequest, request: Request):
     await throttle(request)
     row = await AuthSessions.create(
         data={
-            'connection_id': str(form.connection_id),
-            'name': form.name,
+            **form.model_dump(mode='json', include={'connection_id', 'name'}, exclude_none=True),
             'authorization': {
                 'redirect_uri': form.redirect_uri,
                 'state': form.state,

backend/open_webui/routers/tether.py

@@ -31,13 +31,16 @@ async def connect(ws: WebSocket):
         row = await AuthSessions.get(claims['sid'])
         if not row or row.user_id != user.id or row.refresh_token_id is None:
             raise ValueError('Invalid session')
+        connection_id = row.data.get('connection_id')
+        if not connection_id:
+            raise ValueError('A Tether connection ID is required')
         await hub.start()
         presence = {
             'connection': str(uuid4()),
             'worker': hub.worker,
             'session': row.id,
             'user': user.id,
-            'connection_id': row.data['connection_id'],
+            'connection_id': connection_id,
         }
         connection = SimpleNamespace(ws=ws, presence=presence, routes={}, lock=asyncio.Lock())
 
@@ -48,10 +51,10 @@ async def connect(ws: WebSocket):
         connection.send = send
         # A reconnect replaces the previous connection, including on another worker.
         for old in list(hub.connections.values()):
-            if old.presence['user'] == user.id and old.presence['connection_id'] == row.data['connection_id']:
+            if old.presence['user'] == user.id and old.presence['connection_id'] == connection_id:
                 await old.ws.close(code=1000)
         hub.connections[presence['connection']] = connection
-        key = hub.key(user.id, row.data['connection_id'])
+        key = hub.key(user.id, connection_id)
         record = json.dumps(presence)
         if hub.redis:
             await hub.redis.set(key, record, ex=60)

backend/open_webui/utils/tether.py

@@ -219,7 +219,7 @@ class Hub:
             or row.user_id != user_id
             or row.refresh_token_id is None
             or row.expires_at <= time.time()
-            or row.data['connection_id'] != connection_id
+            or row.data.get('connection_id') != connection_id
         ):
             raise HTTPException(403, 'Tether session is no longer available.')
         if sum(s.presence['connection'] == presence['connection'] for s in self.streams.values()) >= 16:
@@ -272,7 +272,7 @@ async def get_terminal_connections(user):
     seen = set()
     for row in await AuthSessions.list(user_id):
         connection_id = row['connection_id']
-        if connection_id in seen:
+        if not connection_id or connection_id in seen:
             continue
         seen.add(connection_id)
         presence = await hub.presence(user_id, connection_id)