diff --git a/apps/agent/src/visualization_router.py b/apps/agent/src/visualization_router.py index 928556b..5793d43 100644 --- a/apps/agent/src/visualization_router.py +++ b/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), } diff --git a/apps/agent/tests/test_router_edited_messages_regression.py b/apps/agent/tests/test_router_edited_messages_regression.py new file mode 100644 index 0000000..c29375d --- /dev/null +++ b/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"] diff --git a/apps/agent/tests/test_visualization_router.py b/apps/agent/tests/test_visualization_router.py index af1a5d0..beb1ccc 100644 --- a/apps/agent/tests/test_visualization_router.py +++ b/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"))