Skip to content

fix(mcp): skip non-dict models/columns/rels in list and describe #2547

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Bartok9 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(mcp): skip non-dict models/columns/rels in list and describe #2547

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
14ad14f
fix(mcp): skip non-dict models/columns/rels in list and describe
Bartok9 Jul 20, 2026
fd107d3
style(mcp): ruff format non-dict model guards
Bartok9 Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 29 additions & 17 deletions core/wren/src/wren/mcp_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,14 +228,19 @@ def list_models() -> dict:
models = load_models(ctx.project)
result = []
for model in models:
if not isinstance(model, dict):
continue
description = model.get("description")
if description is None:
description = (model.get("properties") or {}).get("description")
columns = model.get("columns") or []
if not isinstance(columns, list):
columns = []
result.append(
{
"name": model.get("name"),
"description": description,
"column_count": len(model.get("columns", []) or []),
"column_count": len(columns),
}
)
return {"models": result}
Expand All @@ -248,25 +253,32 @@ def describe_model(name: str) -> dict:
from wren.context import load_models, load_relationships # noqa: PLC0415

models = load_models(ctx.project)
model = next((m for m in models if m.get("name") == name), None)
model = next(
(m for m in models if isinstance(m, dict) and m.get("name") == name),
None,
)
if model is None:
raise ValueError(f"Model '{name}' not found.")

columns = [
{
"name": col.get("name"),
"type": col.get("type"),
"description": col.get("description")
or (col.get("properties") or {}).get("description"),
}
for col in model.get("columns", []) or []
]

relationships = [
rel
for rel in load_relationships(ctx.project)
if name in (rel.get("models") or [])
]
columns = []
for col in model.get("columns") or []:
if not isinstance(col, dict):
continue
columns.append(
{
"name": col.get("name"),
"type": col.get("type"),
"description": col.get("description")
or (col.get("properties") or {}).get("description"),
}
)
Comment on lines +263 to +274

@coderabbitai coderabbitai Bot Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win

Incomplete guards against scalar YAML values.

Based on the PR objectives to prevent failures from malformed context output, the current type checks are insufficient because they do not protect against truthy scalar types. If malformed YAML values for columns, properties, or relationships parse as integers or booleans, or [] and or {} will fall back to those exact scalars, leading to TypeError during iteration/membership tests or AttributeError when chaining .get(). Furthermore, unvalidated strings can cause unintended substring matches during in checks.

  • core/wren/src/wren/mcp_server.py#L267-L278: Guard model.get("columns") as a list before iterating, and verify col.get("properties") is a dict before calling .get("description").
  • core/wren/src/wren/mcp_server.py#L233-L238: Verify model.get("properties") is a dict before calling .get("description").
  • core/wren/src/wren/mcp_server.py#L280-L285: Guard load_relationships(ctx.project) and rel.get("models") as lists before iterating or performing in checks.
  • core/wren/tests/unit/test_mcp_model_guards.py#L9-L14: Mirror the properties dict check in the test helper.
  • core/wren/tests/unit/test_mcp_model_guards.py#L26-L30: Mirror the columns list check in the test helper.
🛡️ Proposed fixes

core/wren/src/wren/mcp_server.py

# Lines 233-238 (list_models fix)
-            description = model.get("description")
-            if description is None:
-                description = (model.get("properties") or {}).get("description")
-            columns = model.get("columns") or []
-            if not isinstance(columns, list):
-                columns = []
+            description = model.get("description")
+            if description is None:
+                props = model.get("properties")
+                if isinstance(props, dict):
+                    description = props.get("description")
+            columns = model.get("columns")
+            if not isinstance(columns, list):
+                columns = []

# Lines 267-278 (describe_model columns fix)
-        columns = []
-        for col in model.get("columns") or []:
-            if not isinstance(col, dict):
-                continue
-            columns.append(
-                {
-                    "name": col.get("name"),
-                    "type": col.get("type"),
-                    "description": col.get("description")
-                    or (col.get("properties") or {}).get("description"),
-                }
-            )
+        columns = []
+        raw_columns = model.get("columns")
+        if isinstance(raw_columns, list):
+            for col in raw_columns:
+                if not isinstance(col, dict):
+                    continue
+                desc = col.get("description")
+                if not desc:
+                    props = col.get("properties")
+                    if isinstance(props, dict):
+                        desc = props.get("description")
+                columns.append(
+                    {
+                        "name": col.get("name"),
+                        "type": col.get("type"),
+                        "description": desc,
+                    }
+                )

# Lines 280-285 (describe_model relationships fix)
-        relationships = []
-        for rel in load_relationships(ctx.project):
-            if not isinstance(rel, dict):
-                continue
-            if name in (rel.get("models") or []):
-                relationships.append(rel)
+        relationships = []
+        raw_rels = load_relationships(ctx.project)
+        if isinstance(raw_rels, list):
+            for rel in raw_rels:
+                if not isinstance(rel, dict):
+                    continue
+                models_list = rel.get("models")
+                if isinstance(models_list, list) and name in models_list:
+                    relationships.append(rel)

core/wren/tests/unit/test_mcp_model_guards.py

# Lines 9-14 (test helper fix)
-        description = model.get("description")
-        if description is None:
-            description = (model.get("properties") or {}).get("description")
-        columns = model.get("columns") or []
-        if not isinstance(columns, list):
-            columns = []
+        description = model.get("description")
+        if description is None:
+            props = model.get("properties")
+            if isinstance(props, dict):
+                description = props.get("description")
+        columns = model.get("columns")
+        if not isinstance(columns, list):
+            columns = []

# Lines 26-30 (test helper fix)
-    columns = []
-    for col in model.get("columns") or []:
-        if not isinstance(col, dict):
-            continue
-        columns.append({"name": col.get("name")})
+    columns = []
+    raw_columns = model.get("columns")
+    if isinstance(raw_columns, list):
+        for col in raw_columns:
+            if not isinstance(col, dict):
+                continue
+            columns.append({"name": col.get("name")})
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
columns = []
for col in model.get("columns") or []:
if not isinstance(col, dict):
continue
columns.append(
{
"name": col.get("name"),
"type": col.get("type"),
"description": col.get("description")
or (col.get("properties") or {}).get("description"),
}
)
columns = []
raw_columns = model.get("columns")
if isinstance(raw_columns, list):
for col in raw_columns:
if not isinstance(col, dict):
continue
desc = col.get("description")
if not desc:
props = col.get("properties")
if isinstance(props, dict):
desc = props.get("description")
columns.append(
{
"name": col.get("name"),
"type": col.get("type"),
"description": desc,
}
)
📍 Affects 2 files
  • core/wren/src/wren/mcp_server.py#L267-L278 (this comment)
  • core/wren/src/wren/mcp_server.py#L233-L238
  • core/wren/src/wren/mcp_server.py#L280-L285
  • core/wren/tests/unit/test_mcp_model_guards.py#L9-L14
  • core/wren/tests/unit/test_mcp_model_guards.py#L26-L30
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@core/wren/src/wren/mcp_server.py` around lines 267 - 278, Prevent malformed
scalar YAML values from reaching collection operations: in
core/wren/src/wren/mcp_server.py lines 267-278, require model.get("columns") to
be a list and require col.get("properties") to be a dict before reading
description; in lines 233-238, require model.get("properties") to be a dict
before accessing description; in lines 280-285, require
load_relationships(ctx.project) and rel.get("models") to be lists before
iteration or membership checks. Mirror the properties guard in
core/wren/tests/unit/test_mcp_model_guards.py lines 9-14 and the columns guard
in lines 26-30.


relationships = []
for rel in load_relationships(ctx.project):
if not isinstance(rel, dict):
continue
if name in (rel.get("models") or []):
relationships.append(rel)

return {
"name": model.get("name"),
Expand Down
50 changes: 50 additions & 0 deletions core/wren/tests/unit/test_mcp_model_guards.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
"""Guards for non-dict MDL rows in MCP list/describe helpers (mirrors mcp_server)."""


def list_models_payload(models):
result = []
for model in models:
if not isinstance(model, dict):
continue
description = model.get("description")
if description is None:
description = (model.get("properties") or {}).get("description")
columns = model.get("columns") or []
if not isinstance(columns, list):
columns = []
result.append(
{
"name": model.get("name"),
"description": description,
"column_count": len(columns),
}
)
return {"models": result}


def describe_columns(model):
columns = []
for col in model.get("columns") or []:
if not isinstance(col, dict):
continue
columns.append({"name": col.get("name")})
return columns


def test_list_skips_junk():
out = list_models_payload(
[
{"name": "a", "columns": [{"name": "id"}]},
None,
"x",
{"name": "b", "columns": "bad"},
]
)
assert out["models"][0]["column_count"] == 1
assert out["models"][1]["column_count"] == 0
assert len(out["models"]) == 2


def test_describe_skips_non_dict_cols():
cols = describe_columns({"columns": [None, {"name": "id"}, "x"]})
assert cols == [{"name": "id"}]
Loading