From 8ac6f1aa0643016dc9a5eb8263b5e59c80fc669e Mon Sep 17 00:00:00 2001 From: Abhishek Kumar Date: Thu, 9 Jul 2026 18:18:44 +0530 Subject: [PATCH] Fix auth and config validation regressions --- api/db/organization_configuration_client.py | 25 +++++ api/routes/user.py | 33 ++++++- .../configuration/ai_model_configuration.py | 47 +++++++-- api/tests/test_ai_model_configuration_v2.py | 30 ++++++ .../test_user_configuration_validation.py | 99 +++++++++++++++++++ .../components/layout/SidebarTeamSwitcher.tsx | 5 +- .../__tests__/SidebarTeamSwitcher.test.tsx | 32 ++++++ .../auth/providers/StackProviderWrapper.tsx | 43 ++------ .../__tests__/StackProviderWrapper.test.tsx | 81 +++++++++++++++ 9 files changed, 346 insertions(+), 49 deletions(-) create mode 100644 api/tests/test_user_configuration_validation.py create mode 100644 ui/src/lib/auth/providers/__tests__/StackProviderWrapper.test.tsx diff --git a/api/db/organization_configuration_client.py b/api/db/organization_configuration_client.py index 94cb7576..3be9d12e 100644 --- a/api/db/organization_configuration_client.py +++ b/api/db/organization_configuration_client.py @@ -1,3 +1,4 @@ +from datetime import UTC, datetime from typing import Any, Dict, List, Optional from sqlalchemy.future import select @@ -66,6 +67,30 @@ class OrganizationConfigurationClient(BaseDBClient): await session.refresh(config) return config + async def touch_configuration( + self, organization_id: int, key: str + ) -> Optional[OrganizationConfigurationModel]: + """Update the timestamp for an existing organization configuration.""" + async with self.async_session() as session: + result = await session.execute( + select(OrganizationConfigurationModel).where( + OrganizationConfigurationModel.organization_id == organization_id, + OrganizationConfigurationModel.key == key, + ) + ) + config = result.scalars().first() + if not config: + return None + + config.updated_at = datetime.now(UTC) + try: + await session.commit() + except Exception as e: + await session.rollback() + raise e + await session.refresh(config) + return config + async def delete_configuration(self, organization_id: int, key: str) -> bool: """Delete a configuration for an organization.""" async with self.async_session() as session: diff --git a/api/routes/user.py b/api/routes/user.py index 568d9e97..2622b167 100644 --- a/api/routes/user.py +++ b/api/routes/user.py @@ -18,6 +18,7 @@ from api.services.auth.depends import get_user from api.services.configuration.ai_model_configuration import ( convert_legacy_ai_model_configuration_to_v2, get_resolved_ai_model_configuration, + update_organization_ai_model_configuration_last_validated_at, upsert_organization_ai_model_configuration_v2, ) from api.services.configuration.check_validity import ( @@ -107,6 +108,24 @@ class UserConfigurationRequestResponseSchema(BaseModel): organization_pricing: dict[str, Union[float, str, bool]] | None = None +def _is_validation_cache_stale( + last_validated_at: datetime | None, + validity_ttl_seconds: int, +) -> bool: + if last_validated_at is None: + return True + + has_timezone = ( + last_validated_at.tzinfo is not None + and last_validated_at.utcoffset() is not None + ) + if has_timezone: + now = datetime.now(last_validated_at.tzinfo) + else: + now = datetime.now() + return last_validated_at < now - timedelta(seconds=validity_ttl_seconds) + + @router.get("/configurations/user") async def get_user_configurations( user: UserModel = Depends(get_user), @@ -255,10 +274,9 @@ async def validate_user_configurations( ) configurations = resolved_config.effective - if ( - not configurations.last_validated_at - or configurations.last_validated_at - < datetime.now() - timedelta(seconds=validity_ttl_seconds) + if _is_validation_cache_stale( + configurations.last_validated_at, + validity_ttl_seconds, ): validator = UserConfigurationValidator() try: @@ -267,6 +285,13 @@ async def validate_user_configurations( organization_id=user.selected_organization_id, created_by=user.provider_id, ) + if ( + resolved_config.source == "organization_v2" + and user.selected_organization_id is not None + ): + await update_organization_ai_model_configuration_last_validated_at( + user.selected_organization_id + ) return status except ValueError as e: raise HTTPException(status_code=422, detail=e.args[0]) diff --git a/api/services/configuration/ai_model_configuration.py b/api/services/configuration/ai_model_configuration.py index 92d512de..71a9e05e 100644 --- a/api/services/configuration/ai_model_configuration.py +++ b/api/services/configuration/ai_model_configuration.py @@ -11,7 +11,11 @@ from sqlalchemy.orm import selectinload from api.constants import MPS_API_URL from api.db import db_client -from api.db.models import WorkflowDefinitionModel, WorkflowModel +from api.db.models import ( + OrganizationConfigurationModel, + WorkflowDefinitionModel, + WorkflowModel, +) from api.enums import OrganizationConfigurationKey from api.schemas.ai_model_configuration import ( DOGRAH_DEFAULT_LANGUAGE, @@ -58,12 +62,19 @@ async def get_resolved_ai_model_configuration( organization_id: int | None, ) -> ResolvedAIModelConfiguration: """Resolve the effective model configuration for an organization.""" - organization_configuration = await get_organization_ai_model_configuration_v2( - organization_id + organization_configuration_row = ( + await _get_organization_ai_model_configuration_v2_row(organization_id) + ) + organization_configuration = _parse_organization_ai_model_configuration_v2( + organization_configuration_row, + organization_id, ) if organization_configuration is not None: + effective = compile_ai_model_configuration_v2(organization_configuration) + if organization_configuration_row is not None: + effective.last_validated_at = organization_configuration_row.updated_at return ResolvedAIModelConfiguration( - effective=compile_ai_model_configuration_v2(organization_configuration), + effective=effective, source="organization_v2", organization_configuration=organization_configuration, ) @@ -100,12 +111,34 @@ async def get_effective_ai_model_configuration_for_workflow( async def get_organization_ai_model_configuration_v2( organization_id: int | None, ) -> OrganizationAIModelConfigurationV2 | None: - if organization_id is None: - return None - row = await db_client.get_configuration( + row = await _get_organization_ai_model_configuration_v2_row(organization_id) + return _parse_organization_ai_model_configuration_v2(row, organization_id) + + +async def update_organization_ai_model_configuration_last_validated_at( + organization_id: int, +) -> None: + await db_client.touch_configuration( organization_id, OrganizationConfigurationKey.MODEL_CONFIGURATION_V2.value, ) + + +async def _get_organization_ai_model_configuration_v2_row( + organization_id: int | None, +) -> OrganizationConfigurationModel | None: + if organization_id is None: + return None + return await db_client.get_configuration( + organization_id, + OrganizationConfigurationKey.MODEL_CONFIGURATION_V2.value, + ) + + +def _parse_organization_ai_model_configuration_v2( + row: OrganizationConfigurationModel | None, + organization_id: int | None, +) -> OrganizationAIModelConfigurationV2 | None: if row is None or not row.value: return None try: diff --git a/api/tests/test_ai_model_configuration_v2.py b/api/tests/test_ai_model_configuration_v2.py index 151f2013..bdf3558c 100644 --- a/api/tests/test_ai_model_configuration_v2.py +++ b/api/tests/test_ai_model_configuration_v2.py @@ -15,6 +15,7 @@ from api.services.configuration.ai_model_configuration import ( WORKFLOW_MODEL_CONFIGURATION_V2_OVERRIDE_KEY, check_for_masked_keys_in_ai_model_configuration_v2, convert_legacy_ai_model_configuration_to_v2, + get_resolved_ai_model_configuration, mask_ai_model_configuration_v2, merge_ai_model_configuration_v2_secrets, migrate_workflow_configuration_model_override_to_v2, @@ -148,6 +149,35 @@ async def test_byok_realtime_validator_does_not_require_stt_or_tts(): } +@pytest.mark.asyncio +async def test_resolved_org_v2_uses_configuration_updated_at_as_validation_cache( + monkeypatch, +): + from datetime import UTC, datetime + + from api.services.configuration import ai_model_configuration + + last_validated_at = datetime.now(UTC) + config = OrganizationAIModelConfigurationV2( + mode="dograh", + dograh=DograhManagedAIModelConfiguration(api_key="mps-secret"), + ) + row = SimpleNamespace( + value=config.model_dump(mode="json", exclude_none=True), + updated_at=last_validated_at, + ) + monkeypatch.setattr( + ai_model_configuration.db_client, + "get_configuration", + AsyncMock(return_value=row), + ) + + resolved = await get_resolved_ai_model_configuration(organization_id=42) + + assert resolved.source == "organization_v2" + assert resolved.effective.last_validated_at == last_validated_at + + @pytest.mark.asyncio async def test_pipeline_validator_requires_stt_and_tts_when_not_realtime(): effective = EffectiveAIModelConfiguration( diff --git a/api/tests/test_user_configuration_validation.py b/api/tests/test_user_configuration_validation.py new file mode 100644 index 00000000..d4032d94 --- /dev/null +++ b/api/tests/test_user_configuration_validation.py @@ -0,0 +1,99 @@ +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from api.routes import user as user_routes +from api.schemas.ai_model_configuration import EffectiveAIModelConfiguration +from api.services.configuration.ai_model_configuration import ( + ResolvedAIModelConfiguration, +) + + +@pytest.mark.asyncio +async def test_validate_user_configurations_marks_stale_org_v2_config_validated( + monkeypatch, +): + stale_config = EffectiveAIModelConfiguration( + last_validated_at=datetime.now(UTC) - timedelta(seconds=120) + ) + resolved = ResolvedAIModelConfiguration( + effective=stale_config, + source="organization_v2", + ) + validate = AsyncMock(return_value={"status": [{"model": "all", "message": "ok"}]}) + touch_validation_cache = AsyncMock() + + class FakeValidator: + def __init__(self): + self.validate = validate + + monkeypatch.setattr( + user_routes, + "get_resolved_ai_model_configuration", + AsyncMock(return_value=resolved), + ) + monkeypatch.setattr(user_routes, "UserConfigurationValidator", FakeValidator) + monkeypatch.setattr( + user_routes, + "update_organization_ai_model_configuration_last_validated_at", + touch_validation_cache, + ) + + response = await user_routes.validate_user_configurations( + validity_ttl_seconds=60, + user=SimpleNamespace( + provider_id="provider-123", + selected_organization_id=42, + ), + ) + + assert response == {"status": [{"model": "all", "message": "ok"}]} + validate.assert_awaited_once_with( + stale_config, + organization_id=42, + created_by="provider-123", + ) + touch_validation_cache.assert_awaited_once_with(42) + + +@pytest.mark.asyncio +async def test_validate_user_configurations_uses_fresh_org_v2_validation_cache( + monkeypatch, +): + fresh_config = EffectiveAIModelConfiguration(last_validated_at=datetime.now(UTC)) + resolved = ResolvedAIModelConfiguration( + effective=fresh_config, + source="organization_v2", + ) + validate = AsyncMock() + touch_validation_cache = AsyncMock() + + class FakeValidator: + def __init__(self): + self.validate = validate + + monkeypatch.setattr( + user_routes, + "get_resolved_ai_model_configuration", + AsyncMock(return_value=resolved), + ) + monkeypatch.setattr(user_routes, "UserConfigurationValidator", FakeValidator) + monkeypatch.setattr( + user_routes, + "update_organization_ai_model_configuration_last_validated_at", + touch_validation_cache, + ) + + response = await user_routes.validate_user_configurations( + validity_ttl_seconds=60, + user=SimpleNamespace( + provider_id="provider-123", + selected_organization_id=42, + ), + ) + + assert response == {"status": []} + validate.assert_not_awaited() + touch_validation_cache.assert_not_awaited() diff --git a/ui/src/components/layout/SidebarTeamSwitcher.tsx b/ui/src/components/layout/SidebarTeamSwitcher.tsx index ed3372f7..33b6f957 100644 --- a/ui/src/components/layout/SidebarTeamSwitcher.tsx +++ b/ui/src/components/layout/SidebarTeamSwitcher.tsx @@ -2,10 +2,12 @@ import type { CurrentUser, Team } from "@stackframe/stack"; import React, { useState } from "react"; +import { toast } from "sonner"; import SpinLoader from "@/components/SpinLoader"; import { useAuth } from "@/lib/auth"; import { reloadApp } from "@/lib/browserReload"; +import logger from "@/lib/logger"; // Lazy load Stack's SelectedTeamSwitcher, but own the selected-team write here. // The stock component re-applies the selectedTeam prop asynchronously, which @@ -42,8 +44,9 @@ function SidebarTeamSwitcherContent({ user }: { user: CurrentUser }) { await user.setSelectedTeam(team); reloadApp(); } catch (error) { + logger.error("Failed to switch Stack team", error); + toast.error("Could not switch teams. Please try again."); setIsSwitching(false); - throw error; } }; diff --git a/ui/src/components/layout/__tests__/SidebarTeamSwitcher.test.tsx b/ui/src/components/layout/__tests__/SidebarTeamSwitcher.test.tsx index c2339493..d3075c04 100644 --- a/ui/src/components/layout/__tests__/SidebarTeamSwitcher.test.tsx +++ b/ui/src/components/layout/__tests__/SidebarTeamSwitcher.test.tsx @@ -33,11 +33,18 @@ import { TranslationContext } from "../../../../node_modules/@stackframe/stack/d import { SidebarTeamSwitcher } from "../SidebarTeamSwitcher"; const reloadAppMock = vi.hoisted(() => vi.fn()); +const toastErrorMock = vi.hoisted(() => vi.fn()); vi.mock("@/lib/browserReload", () => ({ reloadApp: reloadAppMock, })); +vi.mock("sonner", () => ({ + toast: { + error: toastErrorMock, + }, +})); + // Controls what the faked Stack useUser() returns; tests flip this to null to // simulate the session dying (token expiry / sign-out from another tab). const stackState: { user: ReturnType | null } = { @@ -157,6 +164,7 @@ describe("SidebarTeamSwitcher", () => { beforeEach(() => { stackState.user = makeStackUser(); reloadAppMock.mockClear(); + toastErrorMock.mockClear(); // React logs boundary-caught render errors; keep test output readable. vi.spyOn(console, "error").mockImplementation(() => {}); }); @@ -238,6 +246,30 @@ describe("SidebarTeamSwitcher", () => { }); }); + it("shows feedback and does not reload when Stack fails to persist the selected team", async () => { + stackState.user!.setSelectedTeam = vi.fn(async () => { + throw new Error("Stack write failed"); + }); + + render( + withProviders( + , + makeAuthValue(stackState.user as unknown as AuthUser) + ) + ); + + const switcher = await screen.findByRole("combobox"); + fireEvent.click(switcher); + + await waitFor(() => { + expect(toastErrorMock).toHaveBeenCalledWith( + "Could not switch teams. Please try again." + ); + }); + expect(screen.queryByRole("status")).toBeNull(); + expect(reloadAppMock).not.toHaveBeenCalled(); + }); + it("unmounts the switcher instead of crashing when the session dies (user flips to null)", async () => { const errors: Error[] = []; diff --git a/ui/src/lib/auth/providers/StackProviderWrapper.tsx b/ui/src/lib/auth/providers/StackProviderWrapper.tsx index e578b505..5398ff5e 100644 --- a/ui/src/lib/auth/providers/StackProviderWrapper.tsx +++ b/ui/src/lib/auth/providers/StackProviderWrapper.tsx @@ -1,6 +1,6 @@ 'use client'; -import { type CurrentUser, StackClientApp, StackProvider, StackTheme } from '@stackframe/stack'; +import { StackClientApp, StackProvider, StackTheme } from '@stackframe/stack'; import React, { useMemo, useRef } from 'react'; import { AuthContext } from './AuthProvider'; @@ -41,34 +41,7 @@ function StackAuthContextProvider({ children: React.ReactNode; app: StackClientApp; }) { - const [stackUser, setStackUser] = React.useState(null); - const [isLoading, setIsLoading] = React.useState(true); - - React.useEffect(() => { - let cancelled = false; - - setIsLoading(true); - app.getUser() - .then((user) => { - if (!cancelled) { - setStackUser(user); - } - }) - .catch(() => { - if (!cancelled) { - setStackUser(null); - } - }) - .finally(() => { - if (!cancelled) { - setIsLoading(false); - } - }); - - return () => { - cancelled = true; - }; - }, [app]); + const stackUser = app.useUser(); // Store user in ref for callbacks to access latest value without creating new callbacks const userRef = useRef(stackUser); @@ -123,21 +96,17 @@ function StackAuthContextProvider({ } }, []); - // Use primitive values in deps so user object reference churn does not - // invalidate the context value unless the effective auth state changes. - const userId = stackUser?.id; - const contextValue = useMemo(() => ({ - user: userRef.current, - isAuthenticated: !!userId, - loading: isLoading, + user: stackUser, + isAuthenticated: !!stackUser, + loading: false, getAccessToken, redirectToLogin, logout, provider: 'stack' as const, getSelectedTeam, listPermissions, - }), [userId, isLoading, getAccessToken, redirectToLogin, logout, getSelectedTeam, listPermissions]); + }), [stackUser, getAccessToken, redirectToLogin, logout, getSelectedTeam, listPermissions]); return ( diff --git a/ui/src/lib/auth/providers/__tests__/StackProviderWrapper.test.tsx b/ui/src/lib/auth/providers/__tests__/StackProviderWrapper.test.tsx new file mode 100644 index 00000000..46ac2af0 --- /dev/null +++ b/ui/src/lib/auth/providers/__tests__/StackProviderWrapper.test.tsx @@ -0,0 +1,81 @@ +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { useAuth } from "../AuthProvider"; +import { StackProviderWrapper } from "../StackProviderWrapper"; + +type MockStackUser = { + id: string; + selectedTeam: null; + getAuthJson: () => Promise<{ accessToken: string }>; +}; + +const stackState = vi.hoisted(() => ({ + user: null as MockStackUser | null, +})); + +vi.mock("@stackframe/stack", () => { + class MockStackClientApp { + useUser = vi.fn(() => stackState.user); + + constructor( + public options: { + projectId: string; + publishableClientKey: string; + } + ) {} + } + + return { + StackClientApp: MockStackClientApp, + StackProvider: ({ children }: { children: React.ReactNode }) => children, + StackTheme: ({ children }: { children: React.ReactNode }) => children, + }; +}); + +function AuthProbe() { + const auth = useAuth(); + return ( +
+ {auth.loading ? "loading" : auth.user?.id ?? "signed-out"} +
+ ); +} + +function makeUser(): MockStackUser { + return { + id: "stack-user-1", + selectedTeam: null, + getAuthJson: async () => ({ accessToken: "token" }), + }; +} + +describe("StackProviderWrapper", () => { + beforeEach(() => { + stackState.user = makeUser(); + }); + + it("updates AuthContext when Stack's reactive user becomes null", () => { + const props = { + projectId: "00000000-0000-4000-8000-000000000000", + publishableClientKey: "publishable-client-key", + }; + + const { rerender } = render( + + + + ); + expect(screen.getByTestId("auth-state").textContent).toBe("stack-user-1"); + + stackState.user = null; + rerender( + + + + ); + + expect(screen.getByTestId("auth-state").textContent).toBe("signed-out"); + }); +});