diff --git a/backend/open_webui/models/auth_sessions.py b/backend/open_webui/models/auth_sessions.py index 968afb484b..d2e6eec7db 100644 --- a/backend/open_webui/models/auth_sessions.py +++ b/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 diff --git a/backend/open_webui/routers/sessions.py b/backend/open_webui/routers/sessions.py index 9996891f6a..1de8997c41 100644 --- a/backend/open_webui/routers/sessions.py +++ b/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, diff --git a/backend/open_webui/routers/tether.py b/backend/open_webui/routers/tether.py index f82bdb46e3..562fcbd100 100644 --- a/backend/open_webui/routers/tether.py +++ b/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) diff --git a/backend/open_webui/utils/tether.py b/backend/open_webui/utils/tether.py index 99728f9b86..6592601c7f 100644 --- a/backend/open_webui/utils/tether.py +++ b/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)