diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 53be402fa..0e460f335 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -209,7 +209,7 @@ _ALLOWED_SCOPED_BUILTIN_FUNCTION_TYPES = { 'now': sqlalchemy.sql.functions.now, 'sum': sqlalchemy.sql.functions.sum, } -_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'length', 'nullif'}) +_ALLOWED_SCOPED_GENERIC_FUNCTIONS = frozenset({'date_trunc', 'length', 'nullif'}) _ALLOWED_SCOPED_CUSTOM_OPERATORS = frozenset({'<=>'}) _ALLOWED_SCOPED_STATEMENT_TYPES = ( sqlalchemy.sql.dml.UpdateBase, diff --git a/tests/unit_tests/persistence/test_tenant_uow.py b/tests/unit_tests/persistence/test_tenant_uow.py index 7be62bf2e..bc8369f08 100644 --- a/tests/unit_tests/persistence/test_tenant_uow.py +++ b/tests/unit_tests/persistence/test_tenant_uow.py @@ -961,6 +961,7 @@ async def test_scoped_session_rejects_raw_or_unapproved_sql( sa.select(sa.func.coalesce(sa.func.sum(sa.literal(1)), sa.literal(0))), sa.select( sa.func.now(), + sa.func.date_trunc('hour', sa.column('timestamp')), sa.func.length(sa.literal('value')), sa.func.nullif(sa.literal('value'), sa.literal('')), ), diff --git a/web/src/app/auth/space/callback/page.tsx b/web/src/app/auth/space/callback/page.tsx index e73e2dea4..c2c142589 100644 --- a/web/src/app/auth/space/callback/page.tsx +++ b/web/src/app/auth/space/callback/page.tsx @@ -5,6 +5,7 @@ import { beginAuthenticatedSession, beginSupportAdminSession, bootstrapWorkspaceSession, + clearPendingInvitationToken, getPendingInvitationToken, } from '@/app/infra/http'; import { toast } from 'sonner'; @@ -112,8 +113,31 @@ function SpaceOAuthCallbackContent() { } beginAuthenticatedSession(response.token, response.user); - if (getPendingInvitationToken()) { - navigate('/invitations/accept', { replace: true }); + const invitationToken = getPendingInvitationToken(); + if (invitationToken) { + let invitation; + try { + invitation = + await httpClient.acceptWorkspaceInvitation(invitationToken); + } catch (error) { + const code = (error as { code?: string }).code; + const path = code + ? `/invitations/accept?error=${encodeURIComponent(code)}` + : '/invitations/accept'; + navigate(path, { replace: true }); + return; + } + + beginAuthenticatedSession(invitation.token, response.user); + clearPendingInvitationToken(); + const workspaceResult = await bootstrapWorkspaceSession({ + preferredWorkspaceUuid: invitation.workspace_uuid, + }); + if (workspaceResult.status === 'unavailable') { + navigate('/workspace-unavailable', { replace: true }); + return; + } + navigate('/home', { replace: true }); return; } const workspaceResult = await bootstrapWorkspaceSession({ diff --git a/web/tests/e2e/invitations.spec.ts b/web/tests/e2e/invitations.spec.ts index 5e8b217ec..7cd1cc1cc 100644 --- a/web/tests/e2e/invitations.spec.ts +++ b/web/tests/e2e/invitations.spec.ts @@ -165,3 +165,87 @@ test('an authenticated OSS invitation requires logout before registration', asyn invitation: 'logout-invitation', }); }); + +test('Space OAuth accepts a pending invitation with the freshly authenticated account', async ({ + page, +}) => { + await installLangBotApiMocks(page, { + authenticated: false, + storage: { + token: 'stale-other-account-token', + userEmail: 'other@example.com', + }, + }); + await page.addInitScript(() => { + sessionStorage.setItem( + 'langbot_pending_invitation_token', + 'matching-invitation', + ); + }); + + await page.route('**/api/v1/user/space/callback', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + token: 'fresh-invited-account-token', + user: 'invited@example.com', + }, + msg: 'ok', + }), + }); + }); + await page.route('**/api/v1/user/info', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + account_uuid: 'invited-account', + user: 'invited@example.com', + account_type: 'space', + has_password: false, + }, + msg: 'ok', + }), + }); + }); + + let acceptanceAuthorization = ''; + await page.route('**/api/v1/invitations/accept', async (route) => { + acceptanceAuthorization = route.request().headers().authorization ?? ''; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + code: 0, + data: { + token: 'accepted-invited-account-token', + workspace_uuid: 'workspace-playwright', + }, + msg: 'ok', + }), + }); + }); + + await page.goto('/auth/space/callback?code=oauth-code&state=oauth-state'); + + await expect(page).toHaveURL(/\/home(?:\/monitoring)?$/, { + timeout: 5_000, + }); + expect(acceptanceAuthorization).toBe('Bearer fresh-invited-account-token'); + expect( + await page.evaluate(() => ({ + token: localStorage.getItem('token'), + userEmail: localStorage.getItem('userEmail'), + invitation: sessionStorage.getItem('langbot_pending_invitation_token'), + })), + ).toEqual({ + token: 'accepted-invited-account-token', + userEmail: 'invited@example.com', + invitation: null, + }); +});