fix(agent): reroute presentation after editing a user message

apps/agent/src/visualization_router.py

@@ -47,6 +47,7 @@ class ChoiceAnswer(BaseModel):
 
 class RoutingState(AgentState):
     visualization_turn: NotRequired[str]
+    visualization_request_hash: NotRequired[str]
     visualization_decision: NotRequired[dict]
 
 
@@ -93,7 +94,7 @@ def routing_input(state):
     human = next((m for m in reversed(messages) if m.type == "human"), None)
     if human is None:
         return None
-    # Use the user-message ID, not the content: repeating a prompt is a new turn.
+    # IDs distinguish repeated prompts; content detects edits that retain the ID.
     human_index = max(i for i, m in enumerate(messages) if m.type == "human")
     turn = (
         human.id
@@ -103,7 +104,12 @@ def routing_input(state):
             ).encode()
         ).hexdigest()
     )
-    if state.get("visualization_turn") == turn and state.get("visualization_decision"):
+    request_hash = hashlib.sha256(message_text(human).encode()).hexdigest()
+    if (
+        state.get("visualization_turn") == turn
+        and state.get("visualization_request_hash") == request_hash
+        and state.get("visualization_decision")
+    ):
         return None
     context = [
         {"role": m.type, "content": message_text(m)[:6000]}
@@ -219,15 +225,22 @@ def before_agent(self, state, runtime):
         if pending is None:
             return None
         turn, context = pending
-        return {"visualization_turn": turn, "visualization_decision": route(context)}
+        human = next(m for m in reversed(state["messages"]) if m.type == "human")
+        return {
+            "visualization_turn": turn,
+            "visualization_request_hash": hashlib.sha256(message_text(human).encode()).hexdigest(),
+            "visualization_decision": route(context),
+        }
 
     async def abefore_agent(self, state, runtime):
         pending = routing_input(state)
         if pending is None:
             return None
         turn, context = pending
+        human = next(m for m in reversed(state["messages"]) if m.type == "human")
         return {
             "visualization_turn": turn,
+            "visualization_request_hash": hashlib.sha256(message_text(human).encode()).hexdigest(),
             "visualization_decision": await aroute(context),
         }
 

apps/agent/tests/test_router_edited_messages_regression.py

@@ -0,0 +1,40 @@
+import asyncio
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+from langgraph.graph.message import add_messages
+
+import src.visualization_router as router
+
+
+@pytest.mark.parametrize("use_async", [False, True])
+def test_editing_a_human_message_reroutes_while_continuations_reuse_it(monkeypatch, use_async):
+    calls = []
+
+    def classify(context):
+        calls.append(next(item["content"] for item in reversed(context) if item["role"] == "human"))
+        return {"renderer": "text", "visualization": "text"}
+
+    async def aclassify(context):
+        return classify(context)
+
+    monkeypatch.setattr(router, "route", classify)
+    monkeypatch.setattr(router, "aroute", aclassify)
+    middleware = router.JevVisualizationMiddleware()
+    state = {"messages": [HumanMessage(content="Explain this", id="editable")]}
+
+    def run():
+        if use_async:
+            return asyncio.run(middleware.abefore_agent(state, None))
+        return middleware.before_agent(state, None)
+
+    state.update(run())
+    state["messages"] = add_messages(state["messages"], [AIMessage(content="Done", id="answer")])
+    assert run() is None
+
+    state["messages"] = add_messages(state["messages"], [HumanMessage(content="Put it in a table", id="editable")])
+    update = run()
+    assert update is not None
+    state.update(update)
+    assert run() is None
+    assert calls == ["Explain this", "Put it in a table"]

apps/agent/tests/test_visualization_router.py

@@ -1,3 +1,5 @@
+import hashlib
+
 import pytest
 from langchain_core.messages import HumanMessage, AIMessage
 from src.visualization_router import (
@@ -50,6 +52,7 @@ def test_continuation_reuses_decision_but_new_human_reroutes():
     state = {
         "messages": [HumanMessage(content="table", id="one")],
         "visualization_turn": "one",
+        "visualization_request_hash": hashlib.sha256(b"table").hexdigest(),
         "visualization_decision": decision_from_response(response("table")),
     }
     assert routing_input(state) is None
@@ -224,6 +227,7 @@ def test_idless_continuation_has_stable_turn_key():
     turn, _ = routing_input(state)
     state.update(
         visualization_turn=turn,
+        visualization_request_hash=hashlib.sha256(b"table").hexdigest(),
         visualization_decision=decision_from_response(response("table")),
     )
     state["messages"].append(AIMessage(content="done"))