diff --git a/src/langbot/pkg/api/http/context.py b/src/langbot/pkg/api/http/context.py index eb5c8ac3b..f9c7b4258 100644 --- a/src/langbot/pkg/api/http/context.py +++ b/src/langbot/pkg/api/http/context.py @@ -22,6 +22,7 @@ class PrincipalContext: account_uuid: str | None = None actor_account_uuid: str | None = None api_key_uuid: str | None = None + support_session_id: str | None = None @dataclasses.dataclass(frozen=True, slots=True) diff --git a/src/langbot/pkg/api/http/controller/group.py b/src/langbot/pkg/api/http/controller/group.py index 946eec78c..0917c6c6d 100644 --- a/src/langbot/pkg/api/http/controller/group.py +++ b/src/langbot/pkg/api/http/controller/group.py @@ -348,6 +348,7 @@ class RouterGroup(abc.ABC): principal=PrincipalContext( principal_type=PrincipalType.SUPPORT_ADMIN, actor_account_uuid=identity.actor_account_uuid, + support_session_id=identity.grant_jti_hash, ), workspace=WorkspaceContext( workspace_uuid=identity.workspace_uuid, diff --git a/src/langbot/pkg/api/http/controller/groups/platform/adapters.py b/src/langbot/pkg/api/http/controller/groups/platform/adapters.py index edfbf9615..75a046e8c 100644 --- a/src/langbot/pkg/api/http/controller/groups/platform/adapters.py +++ b/src/langbot/pkg/api/http/controller/groups/platform/adapters.py @@ -22,6 +22,7 @@ class _AdapterSessionScope: principal_type: str account_uuid: str | None api_key_uuid: str | None + support_session_id: str | None @classmethod def from_request_context(cls, request_context: RequestContext) -> '_AdapterSessionScope': @@ -33,6 +34,7 @@ class _AdapterSessionScope: principal_type=principal.principal_type.value, account_uuid=principal.account_uuid, api_key_uuid=principal.api_key_uuid, + support_session_id=principal.support_session_id, ) def matches(self, request_context: RequestContext) -> bool: diff --git a/tests/unit_tests/api/test_adapter_session_scoping.py b/tests/unit_tests/api/test_adapter_session_scoping.py index 3e8642aad..4eb28d1b4 100644 --- a/tests/unit_tests/api/test_adapter_session_scoping.py +++ b/tests/unit_tests/api/test_adapter_session_scoping.py @@ -149,6 +149,36 @@ async def test_session_scope_matches_exact_tenant_placement_and_principal(): assert sessions == {} +async def test_support_admin_sessions_are_scoped_to_the_persisted_grant(): + def support_context(grant_jti_hash: str) -> RequestContext: + return RequestContext( + instance_uuid='instance-test', + placement_generation=1, + request_id='request-test', + auth_type='support-admin', + principal=PrincipalContext( + principal_type=PrincipalType.SUPPORT_ADMIN, + actor_account_uuid='support-actor', + support_session_id=grant_jti_hash, + ), + workspace=WorkspaceContext( + workspace_uuid='workspace-a', + membership_uuid=None, + role='owner', + permissions=frozenset({'resource.manage'}), + ), + ) + + first_context = support_context('a' * 64) + second_context = support_context('b' * 64) + sessions: dict[str, dict] = {'session-test': {'status': 'waiting'}} + _bind_session_scope(sessions['session-test'], first_context) + + assert _get_owned_session(sessions, 'session-test', second_context) is None + assert _pop_owned_session(sessions, 'session-test', second_context) is None + assert _get_owned_session(sessions, 'session-test', first_context) is sessions['session-test'] + + async def test_session_capacity_evicts_oldest_session_in_same_workspace(): owner_context = _request_context() sessions: dict[str, dict] = {}