Merge pull request #2681 from langbot-app/test/quality-auth-boundary-oracles-20261017

test(auth): verify revocation and authorization ordering

tests/integration/api/test_pipelines.py

@@ -462,7 +462,31 @@ class TestPipelineExtensionsEndpoint:
 @pytest.mark.usefixtures('mock_circular_import_chain')
 class TestPipelineDashboardWebSocket:
     @pytest.mark.asyncio
-    async def test_websocket_authenticates_before_registering(self, quart_test_client, fake_pipeline_app):
+    async def test_websocket_authenticates_before_registering(self, quart_test_client, fake_pipeline_app, monkeypatch):
+        import langbot.pkg.api.http.controller.groups.pipelines.websocket_chat as chat
+
+        authenticated = False
+        original_authenticate = chat.WebSocketChatRouterGroup._authenticate_websocket
+        original_add = chat.ws_connection_manager.add_connection
+        proxy = fake_pipeline_app.platform_mgr.get_websocket_proxy_bot
+
+        async def authenticate(router):
+            nonlocal authenticated
+            result = await original_authenticate(router)
+            authenticated = True
+            return result
+
+        async def get_proxy(*_args, **_kwargs):
+            assert authenticated, 'Runtime lookup must follow authentication'
+            return proxy.return_value
+
+        async def add_connection(*args, **kwargs):
+            assert authenticated, 'Connection registration must follow authentication'
+            return await original_add(*args, **kwargs)
+
+        monkeypatch.setattr(chat.WebSocketChatRouterGroup, '_authenticate_websocket', authenticate)
+        monkeypatch.setattr(proxy, 'side_effect', get_proxy)
+        monkeypatch.setattr(chat.ws_connection_manager, 'add_connection', add_connection)
         async with quart_test_client.websocket(
             '/api/v1/pipelines/test-pipeline-uuid/ws/connect?session_type=person',
             headers={'Origin': 'http://localhost'},
@@ -488,7 +512,13 @@ class TestPipelineDashboardWebSocket:
         fake_pipeline_app.platform_mgr.get_websocket_proxy_bot.assert_awaited()
 
     @pytest.mark.asyncio
-    async def test_websocket_rejects_non_auth_first_frame(self, quart_test_client):
+    async def test_websocket_rejects_non_auth_first_frame(self, quart_test_client, fake_pipeline_app, monkeypatch):
+        import langbot.pkg.api.http.controller.groups.pipelines.websocket_chat as chat
+
+        proxy = fake_pipeline_app.platform_mgr.get_websocket_proxy_bot
+        proxy.reset_mock()
+        add_connection = AsyncMock(wraps=chat.ws_connection_manager.add_connection)
+        monkeypatch.setattr(chat.ws_connection_manager, 'add_connection', add_connection)
         async with quart_test_client.websocket(
             '/api/v1/pipelines/test-pipeline-uuid/ws/connect?session_type=person',
             headers={'Origin': 'http://localhost'},
@@ -497,6 +527,9 @@ class TestPipelineDashboardWebSocket:
             response = json.loads(await websocket.receive())
             assert response == {'type': 'error', 'message': 'Unauthorized'}
 
+        proxy.assert_not_awaited()
+        add_connection.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_websocket_rechecks_revocable_membership_before_each_message(
         self,

tests/integration/api/test_support_admin_launch.py

@@ -405,7 +405,7 @@ async def test_support_admin_persisted_expiry_and_revocation_are_enforced(suppor
 
 
 async def test_support_admin_websocket_preserves_actor_and_revalidates(support_admin_api):
-    app, _client, _engine, private_key = support_admin_api
+    app, _client, engine, private_key = support_admin_api
     launch = await _issue_support_token(app, private_key)
     captured_contexts = []
 
@@ -414,7 +414,6 @@ async def test_support_admin_websocket_preserves_actor_and_revalidates(support_a
             del data
             captured_contexts.append(connection.execution_context)
             await connection.send_queue.put({'type': 'handled'})
-            connection.is_active = False
 
     app.pipeline_service = SimpleNamespace(get_pipeline=AsyncMock(return_value=SimpleNamespace(uuid='pipeline-1')))
     app.platform_mgr = SimpleNamespace(
@@ -439,8 +438,17 @@ async def test_support_admin_websocket_preserves_actor_and_revalidates(support_a
         await websocket.send(json.dumps({'type': 'message', 'message': [{'type': 'text', 'text': 'hi'}]}))
         handled = json.loads(await websocket.receive())
         assert handled['type'] == 'handled'
+        async with engine.begin() as connection:
+            await connection.execute(
+                sqlalchemy.update(SupportAdminTemporarySession)
+                .where(SupportAdminTemporarySession.grant_jti_hash == launch['grant_jti_hash'])
+                .values(revoked_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None))
+            )
+        await websocket.send(json.dumps({'type': 'message', 'message': [{'type': 'text', 'text': 'after revoke'}]}))
+        rejected = json.loads(await websocket.receive())
+        assert rejected == {'type': 'error', 'message': 'Unauthorized'}
 
-    assert captured_contexts
+    assert len(captured_contexts) == 1
     principal = captured_contexts[0].trigger_principal
     assert principal is not None
     assert principal.principal_type == PrincipalType.SUPPORT_ADMIN

tests/integration/api/test_user_space_oauth.py

@@ -390,6 +390,9 @@ async def test_oss_login_callback_does_not_request_cloud_reconciliation(space_oa
     )
 
     assert response.status_code == 200
+    payload = await response.get_json()
+    assert payload['code'] == 0
+    assert payload['data']['token'] == 'space-login-token'
     application.directory_projection_service.reconcile_workspaces.assert_not_awaited()
 
 
@@ -524,6 +527,7 @@ async def test_bind_callback_uses_opaque_state_and_never_treats_it_as_jwt(space_
 async def test_direct_launch_assertion_does_not_consume_normal_oauth_state(space_oauth_api):
     application, client = space_oauth_api
     application.user_service.consume_space_oauth_state.reset_mock()
+    application.user_service.consume_space_oauth_state_details.reset_mock()
     application.space_service.exchange_oauth_code.reset_mock()
     application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock())
 
@@ -545,6 +549,7 @@ async def test_direct_launch_assertion_does_not_consume_normal_oauth_state(space
         expected_workspace_uuid=WORKSPACE_UUID,
     )
     application.user_service.consume_space_oauth_state.assert_not_awaited()
+    application.user_service.consume_space_oauth_state_details.assert_not_awaited()
     application.space_service.exchange_oauth_code.assert_not_awaited()
 
 
@@ -558,7 +563,24 @@ async def test_direct_launch_reconciles_exact_workspace_before_resolving_access(
         status='active',
     )
     application.user_service.get_user_by_uuid = AsyncMock(return_value=projected_account)
-    application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock())
+    events = []
+
+    async def reconcile(_workspaces):
+        events.append('reconcile')
+
+    async def get_account(_account_uuid):
+        events.append('account')
+        return projected_account
+
+    access = application.workspace_collaboration_service.resolve_account_workspace.return_value
+
+    async def resolve_access(*_args, **_kwargs):
+        events.append('access')
+        return access
+
+    application.directory_projection_service = SimpleNamespace(reconcile_workspaces=AsyncMock(side_effect=reconcile))
+    application.user_service.get_user_by_uuid.side_effect = get_account
+    application.workspace_collaboration_service.resolve_account_workspace.side_effect = resolve_access
 
     response = await client.post(
         '/api/v1/user/space/callback',
@@ -572,3 +594,5 @@ async def test_direct_launch_reconciles_exact_workspace_before_resolving_access(
     assert (await response.get_json())['data']['workspace_uuid'] == WORKSPACE_UUID
     application.directory_projection_service.reconcile_workspaces.assert_awaited_once_with((WORKSPACE_UUID,))
     application.user_service.get_user_by_uuid.assert_awaited_once_with('account-a')
+
+    assert events == ['reconcile', 'account', 'access']

tests/integration/api/test_workspaces.py

@@ -623,6 +623,12 @@ async def test_api_key_context_returns_bound_identity_without_workspace_permissi
     )
     assert expired_capabilities.status_code == 401
 
+    await application.persistence_mgr.execute_async(
+        sqlalchemy.update(apikey.ApiKey).where(apikey.ApiKey.uuid == created['uuid']).values(expires_at=None)
+    )
+    valid_again = await client.get('/api/v1/system/context', headers={'X-API-Key': created['key']})
+    assert valid_again.status_code == 200
+
     revoke_response = await client.delete(
         f'/api/v1/apikeys/{created["id"]}',
         headers=_auth(owner_token, workspace_uuid),