diff --git a/sdk/wren-pydantic/src/wren_pydantic/_tools.py b/sdk/wren-pydantic/src/wren_pydantic/_tools.py index dd49861bda..2599814a8d 100644 --- a/sdk/wren-pydantic/src/wren_pydantic/_tools.py +++ b/sdk/wren-pydantic/src/wren_pydantic/_tools.py @@ -169,11 +169,26 @@ def _run_list_models(toolkit: WrenToolkit) -> list[ModelSummary]: raise to_model_retry(exc) from exc models = manifest.get("models") or [] - return [ - ModelSummary( - name=m["name"], - column_count=len(m.get("columns") or []), - description=(m.get("properties") or {}).get("description"), + # Skip non-dict rows and tolerate missing name/columns/properties so a + # single corrupted model entry cannot KeyError/TypeError the tool. + summaries: list[ModelSummary] = [] + for m in models: + if not isinstance(m, dict): + continue + name = m.get("name") + if not name: + continue + columns = m.get("columns") or [] + if not isinstance(columns, list): + columns = [] + props = m.get("properties") or {} + if not isinstance(props, dict): + props = {} + summaries.append( + ModelSummary( + name=name, + column_count=len(columns), + description=props.get("description"), + ) ) - for m in models - ] + return summaries diff --git a/sdk/wren-pydantic/tests/unit/test_run_list_models_guard.py b/sdk/wren-pydantic/tests/unit/test_run_list_models_guard.py new file mode 100644 index 0000000000..526714a7a7 --- /dev/null +++ b/sdk/wren-pydantic/tests/unit/test_run_list_models_guard.py @@ -0,0 +1,25 @@ +"""_run_list_models skips non-dict / incomplete model rows.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from wren_pydantic._tools import _run_list_models + + +def test_skips_bad_models(): + toolkit = MagicMock() + toolkit._mdl_source.load_manifest.return_value = { + "models": [ + None, + "x", + {"columns": []}, # no name + {"name": "ok", "columns": None, "properties": "nope"}, + {"name": "t", "columns": [{"n": 1}], "properties": {"description": "d"}}, + ] + } + out = _run_list_models(toolkit) + assert [m.name for m in out] == ["ok", "t"] + assert out[0].column_count == 0 + assert out[1].column_count == 1 + assert out[1].description == "d"