diff --git a/headroom/providers/anthropic.py b/headroom/providers/anthropic.py index 8d4b8313e..fb7278ab4 100644 --- a/headroom/providers/anthropic.py +++ b/headroom/providers/anthropic.py @@ -276,6 +276,11 @@ def _load_custom_model_config() -> dict[str, Any]: # Try to parse as JSON string loaded = json.loads(env_config) + if not isinstance(loaded, dict): + raise ValueError( + f"HEADROOM_MODEL_LIMITS must be a JSON object, got {type(loaded).__name__}" + ) + # Check for anthropic-specific config, fall back to root level anthropic_config = loaded.get("anthropic", loaded) if "context_limits" in anthropic_config: @@ -284,7 +289,10 @@ def _load_custom_model_config() -> dict[str, Any]: config["pricing"].update(anthropic_config["pricing"]) logger.debug(f"Loaded custom model config from HEADROOM_MODEL_LIMITS: {loaded}") - except (json.JSONDecodeError, OSError) as e: + except (ValueError, OSError) as e: + # ValueError covers json.JSONDecodeError (a subclass) and the + # non-object guard above, so a malformed value warns and falls back + # to defaults instead of crashing provider init. logger.warning(f"Failed to load HEADROOM_MODEL_LIMITS: {e}") # Check config file. Prefer the canonical config-dir location, then fall @@ -299,6 +307,9 @@ def _load_custom_model_config() -> dict[str, Any]: with open(config_file, encoding="utf-8") as f: loaded = json.load(f) + if not isinstance(loaded, dict): + raise ValueError(f"{config_file} must contain a JSON object") + # Only load anthropic-specific config anthropic_config = loaded.get("anthropic", loaded) if "context_limits" in anthropic_config: @@ -312,7 +323,7 @@ def _load_custom_model_config() -> dict[str, Any]: config["pricing"][model] = pricing logger.debug(f"Loaded custom model config from {config_file}") - except (json.JSONDecodeError, OSError) as e: + except (ValueError, OSError) as e: logger.warning(f"Failed to load {config_file}: {e}") return config diff --git a/headroom/providers/openai.py b/headroom/providers/openai.py index bb61e9300..b0e2b9afc 100644 --- a/headroom/providers/openai.py +++ b/headroom/providers/openai.py @@ -200,6 +200,11 @@ def _load_custom_model_config() -> dict[str, Any]: # Try to parse as JSON string loaded = json.loads(env_config) + if not isinstance(loaded, dict): + raise ValueError( + f"HEADROOM_MODEL_LIMITS must be a JSON object, got {type(loaded).__name__}" + ) + openai_config = loaded.get("openai", loaded) if "context_limits" in openai_config: config["context_limits"].update(openai_config["context_limits"]) @@ -209,7 +214,10 @@ def _load_custom_model_config() -> dict[str, Any]: config["encodings"].update(openai_config["encodings"]) logger.debug("Loaded custom OpenAI model config from HEADROOM_MODEL_LIMITS") - except (json.JSONDecodeError, OSError) as e: + except (ValueError, OSError) as e: + # ValueError covers json.JSONDecodeError (a subclass) and the + # non-object guard above, so a malformed value warns and falls back + # to defaults instead of crashing provider init. logger.warning(f"Failed to load HEADROOM_MODEL_LIMITS: {e}") # Check config file. Prefer the canonical config-dir location, then fall @@ -224,6 +232,9 @@ def _load_custom_model_config() -> dict[str, Any]: with open(config_file, encoding="utf-8") as f: loaded = json.load(f) + if not isinstance(loaded, dict): + raise ValueError(f"{config_file} must contain a JSON object") + openai_config = loaded.get("openai", {}) if "context_limits" in openai_config: for model, limit in openai_config["context_limits"].items(): @@ -239,7 +250,7 @@ def _load_custom_model_config() -> dict[str, Any]: config["encodings"][model] = encoding logger.debug(f"Loaded custom OpenAI model config from {config_file}") - except (json.JSONDecodeError, OSError) as e: + except (ValueError, OSError) as e: logger.warning(f"Failed to load {config_file}: {e}") return config diff --git a/tests/test_provider_model_fallback.py b/tests/test_provider_model_fallback.py index b19b4e49b..9a83ec0ac 100644 --- a/tests/test_provider_model_fallback.py +++ b/tests/test_provider_model_fallback.py @@ -237,6 +237,25 @@ class TestAnthropicConfigLoading: # Env var should win assert loaded["context_limits"]["test-model"] == 100000 + @pytest.mark.parametrize("raw", ["[1, 2, 3]", '"gpt-4"', "42", "true", "null"]) + def test_non_object_env_var_falls_back_to_defaults(self, raw): + """A valid-JSON-but-not-an-object env var must warn and use defaults, + not crash provider init with AttributeError on ``loaded.get``.""" + with patch.dict(os.environ, {"HEADROOM_MODEL_LIMITS": raw}): + loaded = anthropic_load_config() + assert loaded == {"context_limits": {}, "pricing": {}} + + def test_non_object_config_file_falls_back_to_defaults(self): + """A models.json whose top level is not an object must not crash.""" + with tempfile.TemporaryDirectory() as tmpdir: + config_dir = Path(tmpdir) / ".headroom" + config_dir.mkdir() + (config_dir / "models.json").write_text("[1, 2, 3]") + + with patch.object(Path, "home", return_value=Path(tmpdir)): + loaded = anthropic_load_config() + assert loaded == {"context_limits": {}, "pricing": {}} + class TestOpenAIModelFallback: """Tests for OpenAI provider model fallback.""" @@ -353,6 +372,14 @@ class TestOpenAIConfigLoading: loaded = openai_load_config() assert loaded["pricing"]["test-model"] == [5.0, 15.0] + @pytest.mark.parametrize("raw", ["[1, 2, 3]", '"gpt-4"', "42", "true", "null"]) + def test_non_object_env_var_falls_back_to_defaults(self, raw): + """A valid-JSON-but-not-an-object env var must warn and use defaults, + not crash provider init with AttributeError on ``loaded.get``.""" + with patch.dict(os.environ, {"HEADROOM_MODEL_LIMITS": raw}): + loaded = openai_load_config() + assert loaded == {"context_limits": {}, "pricing": {}, "encodings": {}} + class TestCrossProviderConsistency: """Tests for consistency across providers."""