diff --git a/src/langbot/pkg/api/http/controller/groups/user.py b/src/langbot/pkg/api/http/controller/groups/user.py index 9dce0eca4..240d7a213 100644 --- a/src/langbot/pkg/api/http/controller/groups/user.py +++ b/src/langbot/pkg/api/http/controller/groups/user.py @@ -817,9 +817,7 @@ class UserRouterGroup(group.RouterGroup): rotate = bool(json_data.get('rotate', False)) try: - enrollment = await self.ap.totp_service.begin_enrollment( - account.uuid, account.user, rotate=rotate - ) + enrollment = await self.ap.totp_service.begin_enrollment(account.uuid, account.user, rotate=rotate) except totp_module.TotpError as e: return self.http_status(409, e.code, str(e)) @@ -985,9 +983,7 @@ class UserRouterGroup(group.RouterGroup): return self.http_status(404, 'account_not_found', 'Account not found') try: - enrollment = await self.ap.totp_service.begin_enrollment( - target_account_uuid, target.user, force=True - ) + enrollment = await self.ap.totp_service.begin_enrollment(target_account_uuid, target.user, force=True) except totp_module.TotpError as e: return self.http_status(409, e.code, str(e)) diff --git a/src/langbot/pkg/api/http/service/totp.py b/src/langbot/pkg/api/http/service/totp.py index 052f66835..5995e91ab 100644 --- a/src/langbot/pkg/api/http/service/totp.py +++ b/src/langbot/pkg/api/http/service/totp.py @@ -348,9 +348,7 @@ class TotpService: pending_secret = self._decrypt_secret(credential.secret_ciphertext, credential.key_version) return TotpEnrollment( uuid=credential.uuid, - qr_code_data_url=render_totp_qr_data_url( - build_totp_uri(pending_secret.decode('ascii'), user_email) - ), + qr_code_data_url=render_totp_qr_data_url(build_totp_uri(pending_secret.decode('ascii'), user_email)), algorithm=credential.algorithm, digits=credential.digits, period=credential.period, @@ -688,9 +686,13 @@ class TotpService: return bool(result.rowcount) async def count_unused_recovery_codes(self, account_uuid: str) -> int: - statement = sqlalchemy.select(sqlalchemy.func.count()).select_from(totp_entity.TotpRecoveryCode).where( - totp_entity.TotpRecoveryCode.account_uuid == account_uuid, - totp_entity.TotpRecoveryCode.used_at.is_(None), + statement = ( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(totp_entity.TotpRecoveryCode) + .where( + totp_entity.TotpRecoveryCode.account_uuid == account_uuid, + totp_entity.TotpRecoveryCode.used_at.is_(None), + ) ) async with self._session_factory()() as session: return int(await session.scalar(statement) or 0) diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 11aa3ee6c..5cd5766d0 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -630,9 +630,7 @@ class UserService: return False try: - verified = await totp_service.verify_code( - user_obj.uuid, code, allow_recovery=allow_recovery - ) + verified = await totp_service.verify_code(user_obj.uuid, code, allow_recovery=allow_recovery) except totp_service_module.TotpError: # Covers "not enrolled" and "invalid code" alike. Both are simply a # failed verification for this caller; neither should surface as a diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index e80624afb..c6a48df6b 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -64,6 +64,8 @@ _ALEMBIC_TENANT_TABLES = { 'model_providers', 'codex_credentials', 'passkey_credentials', + 'totp_credentials', + 'totp_recovery_codes', 'llm_models', 'embedding_models', 'rerank_models', diff --git a/tests/integration/persistence/test_rag_document_identity.py b/tests/integration/persistence/test_rag_document_identity.py index fe569c71d..1b939614c 100644 --- a/tests/integration/persistence/test_rag_document_identity.py +++ b/tests/integration/persistence/test_rag_document_identity.py @@ -11,11 +11,15 @@ import pytest_asyncio import sqlalchemy as sa from sqlalchemy.ext.asyncio import create_async_engine +from alembic.config import Config as AlembicConfig +from alembic.script import ScriptDirectory + from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.entity.persistence.base import Base from langbot.pkg.entity.persistence.rag import File, KnowledgeBase from langbot.pkg.entity.persistence.user import User from langbot.pkg.entity.persistence.workspace import Workspace +from langbot.pkg.persistence import alembic_runner from langbot.pkg.persistence.alembic_runner import ( get_alembic_current, run_alembic_downgrade, @@ -27,7 +31,21 @@ from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase from langbot.pkg.workspace.errors import WorkspaceNotFoundError OLD_HEAD = '0024_passkey_credentials' -NEW_HEAD = '0025_rag_document_identity' + + +def _current_script_head() -> str: + """Resolve the live Alembic head instead of pinning a revision number. + + Parallel migrations (the TOTP and RAG document identity branches) are joined + by a merge revision, so the head moves whenever either branch gains a new + migration. Resolving it here keeps this test from needing an edit each time. + """ + + cfg = AlembicConfig() + cfg.set_main_option('script_location', str(alembic_runner._ALEMBIC_DIR)) + return ScriptDirectory.from_config(cfg).get_current_head() + + CONTEXT = ExecutionContext(instance_uuid='instance-a', workspace_uuid='workspace-a', placement_generation=5) @@ -361,7 +379,7 @@ async def test_populated_legacy_migration_roundtrip(database): row = (await conn.execute(sa.text('SELECT * FROM knowledge_base_files'))).mappings().one() assert row['uuid'] == 'legacy' and row['status'] == 'completed' assert row['engine_document_id'] is None - assert await get_alembic_current(database) == NEW_HEAD + assert await get_alembic_current(database) == _current_script_head() await run_alembic_upgrade(database) await run_alembic_stamp(database, OLD_HEAD) await run_alembic_upgrade(database) @@ -379,7 +397,7 @@ async def test_fresh_metadata_then_migration_is_idempotent(database): await create_schema(database) await run_alembic_stamp(database, OLD_HEAD) await run_alembic_upgrade(database) - assert await get_alembic_current(database) == NEW_HEAD + assert await get_alembic_current(database) == _current_script_head() async with database.connect() as conn: assert 'engine_document_id' in await conn.run_sync( lambda sync: {col['name'] for col in sa.inspect(sync).get_columns('knowledge_base_files')} diff --git a/web/src/app/home/components/account-settings-dialog/AccountSettingsPanel.tsx b/web/src/app/home/components/account-settings-dialog/AccountSettingsPanel.tsx index b486326df..ea12d95f4 100644 --- a/web/src/app/home/components/account-settings-dialog/AccountSettingsPanel.tsx +++ b/web/src/app/home/components/account-settings-dialog/AccountSettingsPanel.tsx @@ -71,12 +71,12 @@ export default function AccountSettingsPanel({ const [registeringPasskey, setRegisteringPasskey] = useState(false); const [totpDialogOpen, setTotpDialogOpen] = useState(false); // Latched when the dialog opens so a status refresh cannot swap the flow. - const [totpDialogMode, setTotpDialogMode] = useState('enroll'); + const [totpDialogMode, setTotpDialogMode] = + useState('enroll'); // Owner/admin re-binding flow: the target Account is latched on open. const [adminResetOpen, setAdminResetOpen] = useState(false); - const [adminResetTarget, setAdminResetTarget] = useState( - null, - ); + const [adminResetTarget, setAdminResetTarget] = + useState(null); const [totpRows, setTotpRows] = useState([]); const [isManager, setIsManager] = useState(false); const [accountUuid, setAccountUuid] = useState(''); @@ -497,7 +497,9 @@ export default function AccountSettingsPanel({ size="sm" className="h-8 cursor-pointer" onClick={() => { - setTotpDialogMode(row.enabled ? 'manage' : 'enroll'); + setTotpDialogMode( + row.enabled ? 'manage' : 'enroll', + ); setTotpDialogOpen(true); }} disabled={!systemInfo.allow_modify_login_info} diff --git a/web/src/app/home/components/account-settings-dialog/TotpEnrollDialog.tsx b/web/src/app/home/components/account-settings-dialog/TotpEnrollDialog.tsx index d8f95c654..365ead5aa 100644 --- a/web/src/app/home/components/account-settings-dialog/TotpEnrollDialog.tsx +++ b/web/src/app/home/components/account-settings-dialog/TotpEnrollDialog.tsx @@ -196,7 +196,10 @@ export default function TotpEnrollDialog({ }; return ( - (next ? onOpenChange(true) : closeDialog())}> + (next ? onOpenChange(true) : closeDialog())} + > diff --git a/web/src/app/infra/http/BackendClient.ts b/web/src/app/infra/http/BackendClient.ts index 5ce96a23b..bf27187a0 100644 --- a/web/src/app/infra/http/BackendClient.ts +++ b/web/src/app/infra/http/BackendClient.ts @@ -1309,7 +1309,9 @@ export class BackendClient extends BaseHttpClient { ); } - public confirmTotpEnroll(code: string): Promise<{ recovery_codes: string[] }> { + public confirmTotpEnroll( + code: string, + ): Promise<{ recovery_codes: string[] }> { return this.post( '/api/v1/user/totp/enroll/confirm', { code }, @@ -1328,9 +1330,13 @@ export class BackendClient extends BaseHttpClient { } public disableTotp(code: string): Promise { - return this.post('/api/v1/user/totp/disable', { code }, { - skipWorkspace: true, - }); + return this.post( + '/api/v1/user/totp/disable', + { code }, + { + skipWorkspace: true, + }, + ); } // ============ TOTP oversight (Workspace owner/admin only) ============ diff --git a/web/src/app/login/page.tsx b/web/src/app/login/page.tsx index 2e8116fe5..d0134bb4c 100644 --- a/web/src/app/login/page.tsx +++ b/web/src/app/login/page.tsx @@ -234,24 +234,30 @@ export default function Login() { toast.success(t('common.loginSuccess')); } }) - .catch((error: { code?: string; msg?: string; data?: { challenge_token?: string } }) => { - // The backend answers `totp_required` when the password was correct but - // a second factor is still outstanding. It also hands back the - // challenge token that must accompany the code. - if (error?.code === 'totp_required') { - setPendingEmail(username); - setTotpChallengeToken(error?.data?.challenge_token || ''); - setTotpStep(true); - setUseRecoveryCode(false); - setTotpCode(''); - return; - } - if (error?.code === 'totp_invalid_code') { - toast.error(t('common.totpInvalidCode')); - return; - } - toast.error(t('common.loginFailed')); - }); + .catch( + (error: { + code?: string; + msg?: string; + data?: { challenge_token?: string }; + }) => { + // The backend answers `totp_required` when the password was correct but + // a second factor is still outstanding. It also hands back the + // challenge token that must accompany the code. + if (error?.code === 'totp_required') { + setPendingEmail(username); + setTotpChallengeToken(error?.data?.challenge_token || ''); + setTotpStep(true); + setUseRecoveryCode(false); + setTotpCode(''); + return; + } + if (error?.code === 'totp_invalid_code') { + toast.error(t('common.totpInvalidCode')); + return; + } + toast.error(t('common.loginFailed')); + }, + ); } async function handleTotpSubmit(event: React.FormEvent) { @@ -477,157 +483,157 @@ export default function Login() { ) : ( <> - {/* Space and password login are per-account capabilities. */} - {showSpaceLogin && ( -
- -
- )} + {/* Space and password login are per-account capabilities. */} + {showSpaceLogin && ( +
+ +
+ )} - {showPasskeyLogin && ( -
- -
- )} + {showPasskeyLogin && ( +
+ +
+ )} - {/* Divider - only show if both login methods are available */} - {(showSpaceLogin || showPasskeyLogin) && showLocalLogin && ( -
-
- -
-
- - {t('common.or')} - -
-
- )} + {/* Divider - only show if both login methods are available */} + {(showSpaceLogin || showPasskeyLogin) && showLocalLogin && ( +
+
+ +
+
+ + {t('common.or')} + +
+
+ )} - {/* Password login remains available to every account with a password. */} - {showLocalLogin && ( -
- - ( - - {t('common.email')} - -
- - -
-
- -
- )} - /> + {/* Password login remains available to every account with a password. */} + {showLocalLogin && ( + + + ( + + {t('common.email')} + +
+ + +
+
+ +
+ )} + /> - ( - -
- {t('common.password')} - - {t('common.forgotPassword')} - -
+ ( + +
+ {t('common.password')} + + {t('common.forgotPassword')} + +
- -
- - -
-
- -
- )} - /> + +
+ + +
+
+ +
+ )} + /> - + + + )} + +

+ {t('common.agreementNotice')}{' '} + - {t('common.loginWithPassword')} - - - - )} - -

- {t('common.agreementNotice')}{' '} - - {t('common.termsOfService')} - - {'、'} - - {t('common.privacyPolicy')} - {' '} - {t('common.and')}{' '} - - {t('common.dataCollectionPolicy')} - -

+ {t('common.termsOfService')} + + {'、'} + + {t('common.privacyPolicy')} + {' '} + {t('common.and')}{' '} + + {t('common.dataCollectionPolicy')} + +

)} diff --git a/web/src/app/reset-password/page.tsx b/web/src/app/reset-password/page.tsx index 7301a350a..2221c4e89 100644 --- a/web/src/app/reset-password/page.tsx +++ b/web/src/app/reset-password/page.tsx @@ -121,7 +121,11 @@ export default function ResetPassword() { }); } - const methodButton = (value: ResetMethod, label: string, Icon: typeof KeyRound) => ( + const methodButton = ( + value: ResetMethod, + label: string, + Icon: typeof KeyRound, + ) => (