From 32abbb636f4455e965141d8d209b359dbfbb5aae Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:30:15 +0800 Subject: [PATCH] fix(oss): resolve workspace owner in scoped session --- src/langbot/pkg/api/http/service/user.py | 7 +++-- .../api/service/test_user_service.py | 29 ++++++++++--------- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 44bbc33f2..47bfe8c6c 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -199,7 +199,7 @@ class UserService: async def get_workspace_owner(self, workspace_uuid: str) -> user.User | None: """Resolve the active owner Account for a Workspace.""" - result = await self.ap.persistence_mgr.execute_async( + statement = ( sqlalchemy.select(user.User) .join(WorkspaceMembership, WorkspaceMembership.account_uuid == user.User.uuid) .where( @@ -209,7 +209,10 @@ class UserService: user.User.status == user.AccountStatus.ACTIVE.value, ) ) - return result.scalar_one_or_none() + current_session = self.ap.persistence_mgr.current_session() + if current_session is not None: + return await current_session.scalar(statement) + return await self._identity_scalar(statement, f'workspace-owner:{workspace_uuid}') def _session_factory(self) -> async_sessionmaker[AsyncSession]: return async_sessionmaker(self.ap.persistence_mgr.get_db_engine(), expire_on_commit=False) diff --git a/tests/unit_tests/api/service/test_user_service.py b/tests/unit_tests/api/service/test_user_service.py index 84f936a6b..8a00fad94 100644 --- a/tests/unit_tests/api/service/test_user_service.py +++ b/tests/unit_tests/api/service/test_user_service.py @@ -772,9 +772,7 @@ class TestUserServiceCreateOrUpdateSpaceUser: provider_service = SimpleNamespace(update_space_model_provider_api_keys=AsyncMock()) ap = SimpleNamespace( workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=False)), - workspace_collaboration_service=SimpleNamespace( - list_account_workspaces=AsyncMock(return_value=[access]) - ), + workspace_collaboration_service=SimpleNamespace(list_account_workspaces=AsyncMock(return_value=[access])), provider_service=provider_service, ) @@ -791,17 +789,13 @@ class TestUserServiceCreateOrUpdateSpaceUser: provider_service = SimpleNamespace(update_space_model_provider_api_keys=AsyncMock()) ap = SimpleNamespace( workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=False)), - workspace_collaboration_service=SimpleNamespace( - list_account_workspaces=AsyncMock(return_value=[access]) - ), + workspace_collaboration_service=SimpleNamespace(list_account_workspaces=AsyncMock(return_value=[access])), provider_service=provider_service, ) await UserService(ap)._update_space_provider_for_account(owner_account, 'owner-api-key') - provider_service.update_space_model_provider_api_keys.assert_awaited_once_with( - 'workspace-a', 'owner-api-key' - ) + provider_service.update_space_model_provider_api_keys.assert_awaited_once_with('workspace-a', 'owner-api-key') async def test_create_or_update_space_user_no_expiry(self): """Creates Space user without token expiry.""" @@ -849,12 +843,9 @@ class TestUserServiceCreateOrUpdateSpaceUser: assert result is not None assert result.space_account_uuid == 'noexpiry-uuid' - async def test_bind_space_account_rejects_different_email(self): service = UserService(SimpleNamespace()) - service.get_user_by_email = AsyncMock( - return_value=_create_mock_user(email='invited@example.com') - ) + service.get_user_by_email = AsyncMock(return_value=_create_mock_user(email='invited@example.com')) service.ap.space_service = SimpleNamespace( exchange_oauth_code=AsyncMock( return_value={'access_token': 'access', 'refresh_token': 'refresh', 'expires_in': 3600} @@ -874,6 +865,18 @@ class TestUserServiceCreateOrUpdateSpaceUser: service._identity_execute.assert_not_awaited() + @pytest.mark.asyncio + async def test_get_workspace_owner_returns_user_object_from_core_connection_result(self): + service = UserService(SimpleNamespace()) + owner = _create_mock_user('owner@example.com', password='pw') + service.ap.persistence_mgr = SimpleNamespace( + current_session=lambda: SimpleNamespace(scalar=AsyncMock(return_value=owner)), + ) + + resolved = await service.get_workspace_owner('workspace-1') + + assert resolved is owner + class TestUserServiceLoginCapabilities: async def test_capabilities_are_derived_from_all_accounts(self):