From a140551506b89c6209ee576874ee1f65b0278d3e Mon Sep 17 00:00:00 2001
From: Rudy Celekli <rudy@gradiahq.com>
Date: Mon, 12 Oct 2026 21:09:30 -0400
Subject: [PATCH] fix(agent): filter nested function schemas by renderer

---
 apps/agent/src/visualization_router.py        | 14 ++++----
 .../test_router_tool_schemas_regression.py    | 35 +++++++++++++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)
 create mode 100644 apps/agent/tests/test_router_tool_schemas_regression.py

diff --git a/apps/agent/src/visualization_router.py b/apps/agent/src/visualization_router.py
index 928556b..fe5546e 100644
--- a/apps/agent/src/visualization_router.py
+++ b/apps/agent/src/visualization_router.py
@@ -188,12 +188,14 @@ def routed_request(request):
         "open_generative_ui": {"generateSandboxedUi"},
         "text": set(),
     }[renderer]
-    tools = [
-        t
-        for t in request.tools
-        if (t.get("name") if isinstance(t, dict) else t.name)
-        not in RENDER_TOOLS - allowed
-    ]
+    def tool_name(tool):
+        if not isinstance(tool, dict):
+            return tool.name
+        if tool.get("type") == "function" and isinstance(tool.get("function"), dict):
+            return tool["function"].get("name")
+        return tool.get("name")
+
+    tools = [t for t in request.tools if tool_name(t) not in RENDER_TOOLS - allowed]
     note = f"\n\nCRITICAL: Jev presentation decision for this user turn: renderer={renderer}; visualization={decision['visualization']}. Use this presentation when answering; do not substitute another rendering tool. If required data is unavailable, ask a concise clarification instead of fabricating it."
     if renderer == "a2ui":
         note += '\nUse the A2UI Table component from catalog copilotkit://open-generative-ui-tables. Its title, columns (string array), rows (array of string arrays, same width as columns), and source (honest data provenance) are required. When render_a2ui is available, call it with surfaceId (a unique string), catalogId="copilotkit://open-generative-ui-tables", and components as a flat array containing {"id":"root","component":"Table","title":...,"columns":[...],"rows":[[...]],"source":...}. The Table is the root for this custom catalog. Pass these structured arguments directly, not serialized operations or a2ui_json. If another A2UI tool is available instead, follow its advertised argument schema and use this same catalog and Table component. No HTML or scripts.'
diff --git a/apps/agent/tests/test_router_tool_schemas_regression.py b/apps/agent/tests/test_router_tool_schemas_regression.py
new file mode 100644
index 0000000..58fc8b5
--- /dev/null
+++ b/apps/agent/tests/test_router_tool_schemas_regression.py
@@ -0,0 +1,35 @@
+import pytest
+from langchain.agents.middleware import ModelRequest
+
+from src.visualization_router import routed_request
+
+
+@pytest.mark.parametrize(
+    ("renderer", "allowed"),
+    [("text", set()), ("a2ui", {"render_a2ui"}), ("open_generative_ui", {"generateSandboxedUi"})],
+)
+def test_openai_function_schemas_obey_the_selected_renderer(renderer, allowed):
+    names = ["render_a2ui", "generateSandboxedUi", "barChart", "query_data"]
+    tools = [
+        {"type": "function", "function": {"name": name, "description": name, "parameters": {"type": "object", "properties": {}}}}
+        for name in names
+    ]
+    request = ModelRequest(
+        model=None, messages=[], tools=tools,
+        state={"visualization_decision": {"renderer": renderer, "visualization": "text"}},
+    )
+
+    result = routed_request(request)
+
+    assert {tool["function"]["name"] for tool in result.tools} == allowed | {"query_data"}
+    assert all(tool in tools for tool in result.tools)
+
+
+def test_flat_and_nested_function_schemas_can_be_mixed():
+    data = {"name": "query_data"}
+    renderer = {"type": "function", "function": {"name": "barChart"}}
+    request = ModelRequest(
+        model=None, messages=[], tools=[data, renderer],
+        state={"visualization_decision": {"renderer": "text", "visualization": "text"}},
+    )
+    assert routed_request(request).tools == [data]