diff --git a/ai-hub/app/core/services/preference.py b/ai-hub/app/core/services/preference.py index 387663f..4ccdffb 100644 --- a/ai-hub/app/core/services/preference.py +++ b/ai-hub/app/core/services/preference.py @@ -57,8 +57,17 @@ merged_statuses.update(user_statuses) def is_provider_healthy(section: str, provider_id: str, p_data: dict = None) -> bool: + # M6: Check health using either the full ID or the base ID (in case of personal suffix) status_key = f"{section}_{provider_id}" - is_success = user_statuses.get(status_key) == "success" or system_statuses.get(status_key) == "success" + base_id = provider_id.replace("_personal", "") + base_status_key = f"{section}_{base_id}" + + is_success = ( + user_statuses.get(status_key) == "success" or + user_statuses.get(base_status_key) == "success" or + system_statuses.get(status_key) == "success" or + system_statuses.get(base_status_key) == "success" + ) has_key = p_data and p_data.get("api_key") and p_data.get("api_key") not in ("None", "none", "") return is_success or bool(has_key) @@ -161,10 +170,17 @@ return providers = effective[section_key]["providers"] - filtered_eff = {k: v for k, v in providers.items() if k in allowed} + # M6: Allow explicitly whitelisted IDs OR their personal suffixed versions + def is_allowed(pid): + if pid in allowed: return True + if pid.endswith("_personal") and pid.replace("_personal", "") in allowed: return True + return False + + filtered_eff = {k: v for k, v in providers.items() if is_allowed(k)} effective[section_key]["providers"] = filtered_eff - if effective[section_key].get("active_provider") not in allowed: + curr_active = effective[section_key].get("active_provider") + if curr_active and not is_allowed(curr_active): effective[section_key]["active_provider"] = next(iter(filtered_eff), None) or "" apply_policy("llm", "llm")