Unable to Map message Field to Custom AgentState in LangGraph Studio Chat UI

[bilalcodehub]

Hi all,

I’ve been troubleshooting a LangChain agentic workflow using LangGraph Studio with the new Chat interface (see screenshot below).

My custom AgentState is defined as shown below. I want the question field in the state to automatically receive input from the Chat UI’s message box (the single input field shown at the bottom of the interface), so that when a user types a message, it seamlessly populates the question field.

from typing import Optional, List
from langchain_core.documents import Document
from langgraph.graph import MessagesState

class AgentState(MessagesState):
    """Carries all data passed from node to node in the graph.

    Inherits the `messages` attribute (with automatic merging via
    `add_messages`) from `MessagesState`, and adds agent–specific fields.
    """

    destination_country: Optional[str] = None
    question: Optional[str] = None
    documents: Optional[List[Document]] = None
    answer: Optional[str] = None

To make the graph compatible with the Chat UI — which only passes a top-level message field — I’ve added a prepare_state node at the very start of the graph workflow.add_edge(START, "prepare_state") to convert message to question, like so:

def prepare_state(state: AgentState) -> AgentState:
    # Prefer 'question' if it's set; fallback to 'message'
    question = state.get("question") or state.get("message", "").strip()
    return {"question": question, "answer": None, "documents": []}

This works fine in the Graph UI in the studio where I can directly set the question field, but when I input the message in Message box in the studio's Chat UI, the question field isn’t getting populated — it’s like message is not being passed into the state.

Can someone guide how to map the message input from the Chat UI into the agent's state question field? Do I need to do anything special to get the message field to appear in the initial state passed into the graph?

I will appreciate any guidance.

Many thanks and
Kind regards,
Bilal