test_state_schema_multiple_custom_fields() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/test_state_schema.py lines 56–84

def test_state_schema_multiple_custom_fields() -> None:
    """Test that multiple custom state fields are preserved through agent execution."""

    class CustomState(AgentState[Any]):
        user_id: str
        session_id: str
        context: str

    agent = create_agent(
        model=FakeToolCallingModel(
            tool_calls=[[{"args": {"x": 1}, "id": "call_1", "name": "simple_tool"}], []]
        ),
        tools=[simple_tool],
        state_schema=CustomState,
    )

    result = agent.invoke(
        {
            "messages": [HumanMessage("Test")],
            "user_id": "user_123",
            "session_id": "session_456",
            "context": "test_ctx",
        }
    )

    assert result["user_id"] == "user_123"
    assert result["session_id"] == "session_456"
    assert result["context"] == "test_ctx"
    assert len(result["messages"]) == 4