xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 12:40:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): harden shared cloud runtime boundaries

Browse Source
This commit is contained in:
Junyan Qin
2026-07-20 01:47:42 +08:00
parent 3f01ffe63b
commit 627eb6b8ef
121 changed files with 18142 additions and 5615 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/langbot/pkg/api/http/service/apikey.py
+71 -19
View File
@@ -170,37 +170,89 @@ class ApiKeyService:
if not secret.startswith('lbk_'):
return None
secret_hash = self._hash_secret(secret)
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(apikey.ApiKey).where(apikey.ApiKey.key_hash == secret_hash)
)
key = result.first()
if key is None or key.status != apikey.ApiKeyStatus.ACTIVE.value:
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
discovery_uow = getattr(self.ap.persistence_mgr, 'api_key_discovery_uow', None)
if current_session() is None and callable(discovery_uow):
async with discovery_uow(secret_hash) as discovery:
key = await discovery.session.scalar(
sqlalchemy.select(apikey.ApiKey).where(apikey.ApiKey.key_hash == secret_hash)
)
else:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(apikey.ApiKey).where(apikey.ApiKey.key_hash == secret_hash)
)
key = result.first()
if key is None:
return None
discovered_workspace_uuid = key.workspace_uuid
discovered_key_id = key.id
now = self._utcnow()
if key.expires_at is not None and key.expires_at <= now:
return None
raw_scopes = list(key.scopes or [])
async def bind_and_record_use() -> tuple[typing.Any, typing.Any] | None:
# Re-read inside the tenant transaction. A revoke/expiry racing
# discovery must not result in an authenticated identity.
active_session = current_session()
if active_session is not None:
scoped_key = await active_session.scalar(
sqlalchemy.select(apikey.ApiKey).where(
apikey.ApiKey.id == discovered_key_id,
apikey.ApiKey.workspace_uuid == discovered_workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
)
)
else: # compatibility for isolated service tests
scoped_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(apikey.ApiKey).where(
apikey.ApiKey.id == discovered_key_id,
apikey.ApiKey.workspace_uuid == discovered_workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
)
)
scoped_key = scoped_result.first()
if scoped_key is None or scoped_key.status != apikey.ApiKeyStatus.ACTIVE.value:
return None
if scoped_key.expires_at is not None and scoped_key.expires_at <= now:
return None
binding = await self.ap.workspace_service.get_execution_binding(discovered_workspace_uuid)
updated = await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(apikey.ApiKey)
.where(
apikey.ApiKey.id == scoped_key.id,
apikey.ApiKey.workspace_uuid == discovered_workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
apikey.ApiKey.status == apikey.ApiKeyStatus.ACTIVE.value,
)
.values(last_used_at=now)
.returning(apikey.ApiKey.id)
)
# Authentication and revocation race on this atomic predicate. If
# revoke won, no active row is returned and the stale object read
# above must never become an authenticated identity.
if updated.scalar_one_or_none() is None:
return None
return binding, scoped_key
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if current_session() is None and callable(tenant_uow):
async with tenant_uow(discovered_workspace_uuid):
bound = await bind_and_record_use()
else:
bound = await bind_and_record_use()
if bound is None:
return None
binding, scoped_key = bound
raw_scopes = list(scoped_key.scopes or [])
permissions = (
frozenset(permission.value for permission in Permission)
if '*' in raw_scopes
else frozenset(self._normalize_scopes(raw_scopes))
)
binding = await self.ap.workspace_service.get_execution_binding(key.workspace_uuid)
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(apikey.ApiKey)
.where(
apikey.ApiKey.id == key.id,
apikey.ApiKey.workspace_uuid == key.workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
)
.values(last_used_at=now)
)
return ApiKeyIdentity(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
api_key_uuid=key.uuid,
api_key_uuid=scoped_key.uuid,
permissions=permissions,
)
src/langbot/pkg/api/http/service/maintenance.py
+21
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import datetime
import functools
import os
import re
from pathlib import Path
@@ -22,6 +23,25 @@ DEFAULT_LOG_RETENTION_DAYS = 3
UPLOAD_OWNER_TYPES = ('upload_image', 'upload_document', 'upload')
def _workspace_scope(method):
"""Bind maintenance work to a Workspace without spanning external I/O."""
@functools.wraps(method)
async def wrapped(self, context, *args, **kwargs):
workspace_uuid = require_workspace_uuid(context)
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(tenant_scope):
raise RuntimeError('Cloud maintenance requires an explicit tenant scope')
async with tenant_scope(workspace_uuid):
return await method(self, context, *args, **kwargs)
return await method(self, context, *args, **kwargs)
return wrapped
class MaintenanceService:
"""Storage maintenance and diagnostics."""
@@ -30,6 +50,7 @@ class MaintenanceService:
def __init__(self, ap: app.Application) -> None:
self.ap = ap
@_workspace_scope
async def cleanup_expired_files(self, context: ExecutionContext) -> dict[str, int]:
if not isinstance(context, ExecutionContext):
raise WorkspaceRequiredError('Storage cleanup requires an ExecutionContext')
src/langbot/pkg/api/http/service/mcp.py
+9 -3
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import asyncio
import copy
import re
import uuid
@@ -8,6 +7,7 @@ import uuid
import sqlalchemy
from ....core import app, taskmgr
from ....core.task_boundary import create_detached_task
from ....entity.persistence import mcp as persistence_mcp
from ....entity.persistence import plugin as persistence_plugin
from ....provider.tools.loaders.mcp import MCPSessionStatus, RuntimeMCPSession
@@ -249,7 +249,10 @@ class MCPService:
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_mcp.MCPServer).values(payload))
created = await self._get_mcp_server_by_uuid_raw(execution_context, payload['uuid'])
if created and self.ap.tool_mgr.mcp_tool_loader:
task = asyncio.create_task(self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created))
task = create_detached_task(
self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created),
after_commit_manager=self.ap.persistence_mgr,
)
self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
return payload['uuid']
@@ -351,7 +354,10 @@ class MCPService:
if old_enable and loader.has_session(execution_context, old_name):
await loader.remove_mcp_server(execution_context, old_name)
if new_enable:
task = asyncio.create_task(loader.host_mcp_server(execution_context, updated))
task = create_detached_task(
loader.host_mcp_server(execution_context, updated),
after_commit_manager=self.ap.persistence_mgr,
)
loader._hosted_mcp_tasks.append(task)
async def delete_mcp_server(self, context: TenantContext, server_uuid: str) -> None:
src/langbot/pkg/api/http/service/monitoring.py
+138 -107
View File
@@ -2,8 +2,11 @@ from __future__ import annotations
import uuid
import datetime
import functools
import json
import sqlalchemy
from sqlalchemy.dialects import postgresql as postgresql_dialect
from sqlalchemy.dialects import sqlite as sqlite_dialect
from ....core import app
from ....entity.persistence import monitoring as persistence_monitoring
@@ -12,6 +15,21 @@ from ..context import ExecutionContext
from .tenant import TenantContext, require_workspace_uuid
def _workspace_transaction(method):
"""Run an explicit service entrypoint in one Workspace transaction."""
@functools.wraps(method)
async def wrapped(self, context, *args, **kwargs):
workspace_uuid = require_workspace_uuid(context)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if callable(tenant_uow):
async with tenant_uow(workspace_uuid):
return await method(self, context, *args, **kwargs)
return await method(self, context, *args, **kwargs)
return wrapped
class MonitoringService:
"""Monitoring service"""
@@ -49,7 +67,7 @@ class MonitoringService:
Returns:
A dict mapping table name to the number of deleted rows.
"""
self._require_write_context(context)
workspace_uuid = self._require_write_context(context)
if retention_days < 1:
raise ValueError('retention_days must be >= 1')
if batch_size < 1:
@@ -104,17 +122,27 @@ class MonitoringService:
),
]
deleted_counts: dict[str, int] = {}
async def delete_records() -> dict[str, int]:
deleted_counts: dict[str, int] = {}
for table_name, model_cls, ts_column, pk_column in tables_and_columns:
deleted_counts[table_name] = await self._delete_expired_in_batches(
context=context,
model_cls=model_cls,
ts_column=ts_column,
pk_column=pk_column,
cutoff=cutoff,
batch_size=batch_size,
)
return deleted_counts
for table_name, model_cls, ts_column, pk_column in tables_and_columns:
deleted_counts[table_name] = await self._delete_expired_in_batches(
context=context,
model_cls=model_cls,
ts_column=ts_column,
pk_column=pk_column,
cutoff=cutoff,
batch_size=batch_size,
)
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
if callable(tenant_scope):
# Carry the Workspace across the complete cleanup without holding a
# connection. Each select+delete batch opens and commits its own UoW.
async with tenant_scope(workspace_uuid):
deleted_counts = await delete_records()
else:
deleted_counts = await delete_records()
if sum(deleted_counts.values()) > 0:
await self._release_sqlite_space()
@@ -134,25 +162,36 @@ class MonitoringService:
deleted_total = 0
while True:
select_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(pk_column)
.where(model_cls.workspace_uuid == workspace_uuid, ts_column < cutoff)
.limit(batch_size)
)
pk_values = list(select_result.scalars().all())
if not pk_values:
break
delete_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(model_cls).where(
model_cls.workspace_uuid == workspace_uuid,
pk_column.in_(pk_values),
async def delete_batch() -> tuple[int, int]:
select_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(pk_column)
.where(model_cls.workspace_uuid == workspace_uuid, ts_column < cutoff)
.limit(batch_size)
)
)
deleted = delete_result.rowcount or 0
deleted_total += deleted
pk_values = list(select_result.scalars().all())
if not pk_values:
return 0, 0
if len(pk_values) < batch_size:
delete_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(model_cls).where(
model_cls.workspace_uuid == workspace_uuid,
pk_column.in_(pk_values),
)
)
return len(pk_values), int(delete_result.rowcount or 0)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if callable(tenant_uow):
async with tenant_uow(workspace_uuid):
selected, deleted = await delete_batch()
else:
selected, deleted = await delete_batch()
deleted_total += deleted
if selected == 0:
break
if selected < batch_size:
break
return deleted_total
@@ -192,22 +231,30 @@ class MonitoringService:
session_id: str | None = None,
):
workspace_uuid = self._require_write_context(context)
context_columns = (
persistence_monitoring.MonitoringMessage.id,
persistence_monitoring.MonitoringMessage.bot_id,
persistence_monitoring.MonitoringMessage.bot_name,
persistence_monitoring.MonitoringMessage.pipeline_id,
persistence_monitoring.MonitoringMessage.pipeline_name,
persistence_monitoring.MonitoringMessage.session_id,
)
if message_id:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_monitoring.MonitoringMessage).where(
sqlalchemy.select(*context_columns).where(
persistence_monitoring.MonitoringMessage.workspace_uuid == workspace_uuid,
persistence_monitoring.MonitoringMessage.id == message_id,
)
)
row = result.first()
if row:
return row[0]
return row
if not session_id:
return None
user_query = (
sqlalchemy.select(persistence_monitoring.MonitoringMessage)
sqlalchemy.select(*context_columns)
.where(
sqlalchemy.and_(
persistence_monitoring.MonitoringMessage.session_id == session_id,
@@ -221,10 +268,10 @@ class MonitoringService:
result = await self.ap.persistence_mgr.execute_async(user_query)
row = result.first()
if row:
return row[0]
return row
any_query = (
sqlalchemy.select(persistence_monitoring.MonitoringMessage)
sqlalchemy.select(*context_columns)
.where(
persistence_monitoring.MonitoringMessage.workspace_uuid == workspace_uuid,
persistence_monitoring.MonitoringMessage.session_id == session_id,
@@ -234,10 +281,11 @@ class MonitoringService:
)
result = await self.ap.persistence_mgr.execute_async(any_query)
row = result.first()
return row[0] if row else None
return row
# ========== Recording Methods ==========
@_workspace_transaction
async def record_message(
self,
context: ExecutionContext,
@@ -285,6 +333,7 @@ class MonitoringService:
return message_id
@_workspace_transaction
async def record_llm_call(
self,
context: ExecutionContext,
@@ -331,6 +380,7 @@ class MonitoringService:
return call_id
@_workspace_transaction
async def record_tool_call(
self,
context: ExecutionContext,
@@ -389,6 +439,7 @@ class MonitoringService:
return call_id
@_workspace_transaction
async def record_embedding_call(
self,
context: ExecutionContext,
@@ -432,6 +483,7 @@ class MonitoringService:
return call_id
@_workspace_transaction
async def record_session_start(
self,
context: ExecutionContext,
@@ -466,6 +518,7 @@ class MonitoringService:
sqlalchemy.insert(persistence_monitoring.MonitoringSession).values(session_data)
)
@_workspace_transaction
async def update_session_activity(
self,
context: ExecutionContext,
@@ -503,6 +556,7 @@ class MonitoringService:
# Check if any rows were updated
return result.rowcount > 0
@_workspace_transaction
async def record_error(
self,
context: ExecutionContext,
@@ -540,6 +594,7 @@ class MonitoringService:
return error_id
@_workspace_transaction
async def update_message_status(
self,
context: ExecutionContext,
@@ -1802,81 +1857,57 @@ class MonitoringService:
)
return None
# Check if record with this feedback_id already exists
existing_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(MonitoringFeedback).where(
MonitoringFeedback.workspace_uuid == workspace_uuid,
MonitoringFeedback.feedback_id == feedback_id,
)
)
existing_row = existing_result.first()
if existing_row:
# UPDATE existing record
existing = existing_row[0] if isinstance(existing_row, tuple) else existing_row
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(MonitoringFeedback)
.where(
MonitoringFeedback.workspace_uuid == workspace_uuid,
MonitoringFeedback.feedback_id == feedback_id,
)
.values(
timestamp=now,
feedback_type=feedback_type,
feedback_content=feedback_content,
inaccurate_reasons=reasons_json,
bot_id=bot_id or existing.bot_id,
bot_name=bot_name or existing.bot_name,
pipeline_id=pipeline_id or existing.pipeline_id,
pipeline_name=pipeline_name or existing.pipeline_name,
session_id=session_id or existing.session_id,
message_id=message_id or existing.message_id,
stream_id=stream_id or existing.stream_id,
user_id=user_id or existing.user_id,
platform=platform or existing.platform,
)
)
return existing.id
record_data = {
'id': str(uuid.uuid4()),
'workspace_uuid': workspace_uuid,
'timestamp': now,
'feedback_id': feedback_id,
'feedback_type': feedback_type,
'feedback_content': feedback_content,
'inaccurate_reasons': reasons_json,
'bot_id': bot_id,
'bot_name': bot_name,
'pipeline_id': pipeline_id,
'pipeline_name': pipeline_name,
'session_id': session_id,
'message_id': message_id,
'stream_id': stream_id,
'user_id': user_id,
'platform': platform,
}
dialect_name = self.ap.persistence_mgr.get_db_engine().dialect.name
if dialect_name == 'postgresql':
statement = postgresql_dialect.insert(MonitoringFeedback).values(record_data)
elif dialect_name == 'sqlite':
statement = sqlite_dialect.insert(MonitoringFeedback).values(record_data)
else:
# INSERT new record with IntegrityError defense
record_id = str(uuid.uuid4())
record_data = {
'id': record_id,
'workspace_uuid': workspace_uuid,
'timestamp': now,
'feedback_id': feedback_id,
'feedback_type': feedback_type,
'feedback_content': feedback_content,
'inaccurate_reasons': reasons_json,
'bot_id': bot_id,
'bot_name': bot_name,
'pipeline_id': pipeline_id,
'pipeline_name': pipeline_name,
'session_id': session_id,
'message_id': message_id,
'stream_id': stream_id,
'user_id': user_id,
'platform': platform,
}
try:
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(MonitoringFeedback).values(record_data))
return record_id
except Exception:
# UNIQUE constraint conflict (concurrent feedback for same feedback_id)
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(MonitoringFeedback)
.where(
MonitoringFeedback.workspace_uuid == workspace_uuid,
MonitoringFeedback.feedback_id == feedback_id,
)
.values(
timestamp=now,
feedback_type=feedback_type,
feedback_content=feedback_content,
inaccurate_reasons=reasons_json,
)
)
return feedback_id
raise RuntimeError(f'Monitoring feedback upsert does not support {dialect_name!r}')
excluded = statement.excluded
def preserve_existing(column):
return sqlalchemy.func.coalesce(sqlalchemy.func.nullif(getattr(excluded, column.key), ''), column)
statement = statement.on_conflict_do_update(
index_elements=[MonitoringFeedback.workspace_uuid, MonitoringFeedback.feedback_id],
set_={
'timestamp': excluded.timestamp,
'feedback_type': excluded.feedback_type,
'feedback_content': excluded.feedback_content,
'inaccurate_reasons': excluded.inaccurate_reasons,
'bot_id': preserve_existing(MonitoringFeedback.bot_id),
'bot_name': preserve_existing(MonitoringFeedback.bot_name),
'pipeline_id': preserve_existing(MonitoringFeedback.pipeline_id),
'pipeline_name': preserve_existing(MonitoringFeedback.pipeline_name),
'session_id': preserve_existing(MonitoringFeedback.session_id),
'message_id': preserve_existing(MonitoringFeedback.message_id),
'stream_id': preserve_existing(MonitoringFeedback.stream_id),
'user_id': preserve_existing(MonitoringFeedback.user_id),
'platform': preserve_existing(MonitoringFeedback.platform),
},
).returning(MonitoringFeedback.id)
result = await self.ap.persistence_mgr.execute_async(statement)
return str(result.scalar_one())
async def get_feedback_stats(
self,
src/langbot/pkg/api/http/service/skill.py
+53 -5
View File
@@ -4,6 +4,7 @@ import io
import inspect
import os
import posixpath
import stat
import zipfile
from typing import Optional
from urllib.parse import quote, unquote, urlparse
@@ -34,6 +35,12 @@ _GITHUB_ASSET_HOSTS = {
'raw.githubusercontent.com',
'codeload.github.com',
}
_MAX_GITHUB_ARCHIVE_BYTES = 10 * 1024 * 1024
_MAX_GITHUB_ARCHIVE_ENTRIES = 4096
_MAX_SKILL_ARCHIVE_FILES = 1024
_MAX_SKILL_FILE_BYTES = 10 * 1024 * 1024
_MAX_SKILL_UNCOMPRESSED_BYTES = 50 * 1024 * 1024
_MAX_SKILL_COMPRESSION_RATIO = 200
class SkillService:
@@ -322,9 +329,22 @@ class SkillService:
async def _download_github_asset(self, asset_url: str) -> bytes:
async with httpx.AsyncClient(follow_redirects=True, timeout=120) as client:
resp = await client.get(asset_url)
resp.raise_for_status()
return resp.content
async with client.stream('GET', asset_url) as resp:
resp.raise_for_status()
content_length = resp.headers.get('content-length')
if content_length is not None:
try:
if int(content_length) > _MAX_GITHUB_ARCHIVE_BYTES:
raise ValueError('GitHub skill archive exceeds the compressed size limit')
except ValueError as exc:
if 'exceeds' in str(exc):
raise
content = bytearray()
async for chunk in resp.aiter_bytes():
content.extend(chunk)
if len(content) > _MAX_GITHUB_ARCHIVE_BYTES:
raise ValueError('GitHub skill archive exceeds the compressed size limit')
return bytes(content)
async def _download_github_skill_directory_as_zip(
self, asset_url: str, *, owner: str, repo: str
@@ -339,7 +359,11 @@ class SkillService:
raise ValueError('GitHub repository archive must be a valid .zip archive') from exc
with source_archive as source_zip:
if len(source_zip.infolist()) > _MAX_GITHUB_ARCHIVE_ENTRIES:
raise ValueError('GitHub repository archive contains too many entries')
skill_entry = self._find_github_skill_archive_entry(source_zip, info['file_path'])
if skill_entry.file_size > _MAX_SKILL_FILE_BYTES:
raise ValueError('GitHub SKILL.md exceeds the file size limit')
try:
skill_md_content = source_zip.read(skill_entry).decode('utf-8')
except UnicodeDecodeError as exc:
@@ -377,6 +401,7 @@ class SkillService:
normalized_source_dir = posixpath.normpath(source_skill_dir)
source_prefix = f'{normalized_source_dir}/'
copied_files = 0
copied_bytes = 0
for member in source_zip.infolist():
normalized_member = posixpath.normpath(member.filename)
@@ -399,10 +424,33 @@ class SkillService:
if member.is_dir():
target_zip.writestr(target_info, b'')
continue
target_zip.writestr(target_info, source_zip.read(member))
if member.flag_bits & 0x1:
raise ValueError('Encrypted GitHub skill archive entries are not supported')
unix_mode = member.external_attr >> 16
if stat.S_IFMT(unix_mode) == stat.S_IFLNK:
raise ValueError(f'GitHub archive contains a symbolic link: {member.filename}')
if member.file_size > _MAX_SKILL_FILE_BYTES:
raise ValueError(f'GitHub skill file exceeds the size limit: {member.filename}')
if member.file_size and member.file_size > max(member.compress_size, 1) * _MAX_SKILL_COMPRESSION_RATIO:
raise ValueError(f'GitHub skill file exceeds the compression-ratio limit: {member.filename}')
copied_files += 1
copied_bytes += member.file_size
if copied_files > _MAX_SKILL_ARCHIVE_FILES:
raise ValueError('GitHub skill directory contains too many files')
if copied_bytes > _MAX_SKILL_UNCOMPRESSED_BYTES:
raise ValueError('GitHub skill directory exceeds the uncompressed size limit')
# Copy in bounded chunks instead of materialising a potentially
# large member in Core memory. The Box Runtime independently
# revalidates the resulting archive before installation.
with source_zip.open(member, 'r') as source_file, target_zip.open(target_info, 'w') as target_file:
remaining = member.file_size
while remaining:
chunk = source_file.read(min(64 * 1024, remaining))
if not chunk:
raise ValueError(f'GitHub skill file is truncated: {member.filename}')
target_file.write(chunk)
remaining -= len(chunk)
if copied_files == 0:
raise ValueError('GitHub skill directory is empty')
src/langbot/pkg/api/http/service/user.py
+69 -34
View File
@@ -151,10 +151,11 @@ class UserService:
)
async def is_initialized(self) -> bool:
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(user.User).limit(1))
result_list = result.all()
return result_list is not None and len(result_list) > 0
account = await self._identity_scalar(
sqlalchemy.select(user.User).limit(1),
f'instance:{self._jwt_identity()[1]}',
)
return account is not None
def _session_factory(self) -> async_sessionmaker[AsyncSession]:
return async_sessionmaker(self.ap.persistence_mgr.get_db_engine(), expire_on_commit=False)
@@ -247,28 +248,24 @@ class UserService:
async def get_user_by_email(self, user_email: str) -> user.User | None:
normalized_email = user_email.strip().casefold()
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(user.User).where(user.User.normalized_email == normalized_email)
return await self._identity_scalar(
sqlalchemy.select(user.User).where(user.User.normalized_email == normalized_email),
f'email:{normalized_email}',
)
result_list = result.all()
return result_list[0] if result_list is not None and len(result_list) > 0 else None
async def get_user_by_uuid(self, account_uuid: str) -> user.User | None:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(user.User).where(user.User.uuid == account_uuid)
return await self._identity_scalar(
sqlalchemy.select(user.User).where(user.User.uuid == account_uuid),
f'uuid:{account_uuid}',
)
return result.first()
async def get_user_by_space_account_uuid(self, space_account_uuid: str) -> user.User | None:
"""Get user by Space account UUID"""
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(user.User).where(user.User.space_account_uuid == space_account_uuid)
return await self._identity_scalar(
sqlalchemy.select(user.User).where(user.User.space_account_uuid == space_account_uuid),
f'space:{space_account_uuid}',
)
result_list = result.all()
return result_list[0] if result_list is not None and len(result_list) > 0 else None
async def authenticate(self, user_email: str, password: str) -> str | None:
user_obj = await self.get_user_by_email(user_email)
if user_obj is None:
@@ -389,11 +386,13 @@ class UserService:
async def reset_password(self, user_email: str, new_password: str) -> None:
hashed_password = await self._hash_password(new_password)
normalized_email = normalize_email(user_email)
await self.ap.persistence_mgr.execute_async(
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.values(password=hashed_password)
.where(user.User.normalized_email == normalized_email)
.values(password=hashed_password),
f'email:{normalized_email}',
)
async def change_password(self, user_email: str, current_password: str, new_password: str) -> None:
@@ -407,11 +406,13 @@ class UserService:
await self._verify_password(user_obj.password, current_password)
hashed_password = await self._hash_password(new_password)
normalized_email = normalize_email(user_email)
await self.ap.persistence_mgr.execute_async(
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.values(password=hashed_password)
.where(user.User.normalized_email == normalized_email)
.values(password=hashed_password),
f'email:{normalized_email}',
)
# Space user management
@@ -435,7 +436,7 @@ class UserService:
if existing_user:
# Update existing user's tokens
await self.ap.persistence_mgr.execute_async(
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.space_account_uuid == space_account_uuid)
.values(
@@ -443,7 +444,8 @@ class UserService:
space_refresh_token=refresh_token,
space_api_key=api_key,
space_access_token_expires_at=expires_at,
)
),
f'space:{space_account_uuid}',
)
await self._update_space_provider_for_account(existing_user, api_key)
return await self.get_user_by_space_account_uuid(space_account_uuid)
@@ -538,9 +540,38 @@ class UserService:
async def get_first_user(self) -> user.User | None:
"""Get the first user (for single-user mode)"""
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(user.User).limit(1))
result_list = result.all()
return result_list[0] if result_list else None
return await self._identity_scalar(
sqlalchemy.select(user.User).limit(1),
f'instance:{self._jwt_identity()[1]}',
)
async def _identity_scalar(
self,
statement: typing.Any,
identity: str,
) -> user.User | None:
"""Execute one exact Account lookup in an explicit discovery transaction."""
digest = hashlib.sha256(identity.encode('utf-8')).hexdigest()
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
identity_uow = getattr(self.ap.persistence_mgr, 'identity_discovery_uow', None)
if current_session() is None and callable(identity_uow):
async with identity_uow(digest) as discovery:
return await discovery.session.scalar(statement)
result = await self.ap.persistence_mgr.execute_async(statement)
rows = result.all()
return rows[0] if rows else None
async def _identity_execute(self, statement: typing.Any, identity: str) -> typing.Any:
"""Execute one exact Account mutation in an explicit transaction."""
digest = hashlib.sha256(identity.encode('utf-8')).hexdigest()
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
identity_uow = getattr(self.ap.persistence_mgr, 'identity_discovery_uow', None)
if current_session() is None and callable(identity_uow):
async with identity_uow(digest) as discovery:
return await discovery.session.execute(statement)
return await self.ap.persistence_mgr.execute_async(statement)
async def set_password(self, user_email: str, new_password: str, current_password: str | None = None) -> None:
"""Set or change password for a user"""
@@ -557,10 +588,12 @@ class UserService:
await self._verify_password(user_obj.password, current_password)
hashed_password = await self._hash_password(new_password)
await self.ap.persistence_mgr.execute_async(
normalized_email = normalize_email(user_email)
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.values(password=hashed_password)
.where(user.User.normalized_email == normalized_email)
.values(password=hashed_password),
f'email:{normalized_email}',
)
async def bind_space_account(self, user_email: str, code: str) -> user.User:
@@ -596,9 +629,10 @@ class UserService:
raise ValueError('This Space account is already bound to another user')
# Update local account to Space account
await self.ap.persistence_mgr.execute_async(
normalized_email = normalize_email(user_email)
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.where(user.User.normalized_email == normalized_email)
.values(
user=normalize_email(space_email), # Update email to Space email
normalized_email=normalize_email(space_email),
@@ -608,7 +642,8 @@ class UserService:
space_refresh_token=refresh_token,
space_api_key=api_key,
space_access_token_expires_at=expires_at,
)
),
f'email:{normalized_email}',
)
# Update Space model provider API keys