fix cloud monitoring and invitation sign-in (#2381)

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>

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,

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('')),
         ),

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()) {
+        const invitationToken = getPendingInvitationToken();
-          navigate('/invitations/accept', { replace: true });
+        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({

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,
+  });
+});