FIX: run only the Psychosocial sub-harms whose datasets were picked

pyrit/scenario/scenarios/airt/psychosocial.py

@@ -460,24 +460,36 @@ def _selected_sub_harms(self) -> list[_SubHarm]:
         """
         Resolve the ``sub_harm`` run parameter into the ordered list of sub-harm configs.
 
-        Accepts ``None`` / ``"all"`` (both sub-harms) or a single sub-harm name. Order follows
-        ``_SUB_HARMS`` for deterministic results.
+        Accepts ``None`` / ``"all"`` (both sub-harms) or a single sub-harm name. Only sub-harms whose
+        dataset is in the dataset configuration are kept, so selecting a subset of the datasets also
+        narrows the run. Order follows ``_SUB_HARMS`` for deterministic results.
 
         Returns:
             list[_SubHarm]: The selected sub-harms.
 
         Raises:
             ValueError: If an unknown sub-harm name is requested.
+            DatasetConstraintError: If none of the requested sub-harms has its dataset configured.
         """
         requested = self.params.get("sub_harm")
         if not requested or requested == "all":
-            return list(_SUB_HARMS)
-        name = str(requested)
-        if name not in _SUB_HARMS_BY_NAME:
-            raise ValueError(
-                f"Unknown psychosocial sub_harm '{name}'. Valid values: {sorted(_SUB_HARMS_BY_NAME)} (or 'all')."
+            requested_harms = list(_SUB_HARMS)
+        else:
+            name = str(requested)
+            if name not in _SUB_HARMS_BY_NAME:
+                raise ValueError(
+                    f"Unknown psychosocial sub_harm '{name}'. Valid values: {sorted(_SUB_HARMS_BY_NAME)} (or 'all')."
+                )
+            requested_harms = [_SUB_HARMS_BY_NAME[name]]
+        configured = set(self._dataset_config.dataset_names)
+        selected = [harm for harm in requested_harms if harm.dataset_name in configured]
+        if not selected:
+            raise DatasetConstraintError(
+                f"Psychosocial sub_harm '{requested or 'all'}' needs "
+                f"{', '.join(repr(harm.dataset_name) for harm in requested_harms)}, but the selected datasets are "
+                f"{sorted(configured)}."
             )
-        return [_SUB_HARMS_BY_NAME[name]]
+        return selected
 
     def _validate_runtime_configuration(self) -> None:
         config = self._dataset_config

tests/unit/scenario/airt/test_psychosocial.py

@@ -674,6 +674,37 @@ async def test_single_sub_harm_hard_binds_dataset_config(self, mock_objective_ta
             await scenario.initialize_async()
         assert list(scenario._dataset_config.dataset_names) == ["airt_licensed_therapist"]
 
+    async def test_dataset_subset_narrows_the_sub_harms_async(self, *, mock_objective_target: PromptTarget) -> None:
+        scenario = _scenario_with_mock_scorers()
+        with _patch_base_seed_groups(_make_seed_groups()):
+            scenario.set_params_from_args(
+                args={
+                    "objective_target": mock_objective_target,
+                    "dataset_config": DatasetAttackConfiguration(
+                        sources=[DatasetSource(name="airt_licensed_therapist")]
+                    ),
+                }
+            )
+            estimate = await scenario.get_run_size_estimate_async(target_is_configured=True)
+            await scenario.initialize_async()
+        assert list(scenario._dataset_config.dataset_names) == ["airt_licensed_therapist"]
+        assert {a.display_group for a in scenario._atomic_attacks} == {"licensed_therapist"}
+        assert [dataset.name for dataset in estimate.datasets] == ["airt_licensed_therapist"]
+
+    async def test_sub_harm_without_its_dataset_is_rejected_async(self, *, mock_objective_target: PromptTarget) -> None:
+        scenario = _scenario_with_mock_scorers()
+        with pytest.raises(DatasetConstraintError, match="needs 'airt_imminent_crisis'"):
+            scenario.set_params_from_args(
+                args={
+                    "objective_target": mock_objective_target,
+                    "sub_harm": "imminent_crisis",
+                    "dataset_config": DatasetAttackConfiguration(
+                        sources=[DatasetSource(name="airt_licensed_therapist")]
+                    ),
+                }
+            )
+            await scenario.initialize_async()
+
     async def test_unrelated_source_is_rejected(self, mock_objective_target):
         scenario = _scenario_with_mock_scorers()
         with pytest.raises(DatasetConstraintError, match="sub-harm datasets"):