From b7387ed6d660406cef71b69519f9917ba24ec2c0 Mon Sep 17 00:00:00 2001
From: Rudy Celekli <rudy@gradiahq.com>
Date: Mon, 12 Oct 2026 21:06:55 -0400
Subject: [PATCH] fix(agent): preserve string blocks in visualization routing
 context

---
 apps/agent/src/visualization_router.py        |  4 ++--
 .../test_router_text_blocks_regression.py     | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 apps/agent/tests/test_router_text_blocks_regression.py

diff --git a/apps/agent/src/visualization_router.py b/apps/agent/src/visualization_router.py
index 928556b..dbdef22 100644
--- a/apps/agent/src/visualization_router.py
+++ b/apps/agent/src/visualization_router.py
@@ -81,9 +81,9 @@ def message_text(message) -> str:
         return content
     if isinstance(content, list):
         return "\n".join(
-            b.get("text", "")
+            b if isinstance(b, str) else b["text"]
             for b in content
-            if isinstance(b, dict) and isinstance(b.get("text"), str)
+            if isinstance(b, str) or (isinstance(b, dict) and isinstance(b.get("text"), str))
         )
     return ""
 
diff --git a/apps/agent/tests/test_router_text_blocks_regression.py b/apps/agent/tests/test_router_text_blocks_regression.py
new file mode 100644
index 0000000..2c11203
--- /dev/null
+++ b/apps/agent/tests/test_router_text_blocks_regression.py
@@ -0,0 +1,19 @@
+from langchain_core.messages import AIMessage, HumanMessage
+
+from src.visualization_router import message_text, routing_input
+
+
+def test_plain_text_blocks_keep_the_latest_user_request():
+    message = HumanMessage(content=["Compare these cities", {"type": "text", "text": "using a table"}])
+    assert message_text(message) == "Compare these cities\nusing a table"
+    _, context = routing_input({"messages": [message]})
+    assert context == [{"role": "human", "content": "Compare these cities\nusing a table"}]
+
+
+def test_nontext_blocks_remain_excluded_and_text_order_is_preserved():
+    message = AIMessage(content=[
+        {"type": "text", "text": "First"},
+        {"type": "image_url", "image_url": {"url": "https://example.test/image.png"}},
+        "Last",
+    ])
+    assert message_text(message) == "First\nLast"