Displaying assembly and change its color

I am reading STEP assembly using STEPCAFControl_Reader and putting all the compounds and solids in a tree structure.

My tree node contains a name which I obtained from TDF_Label, and a shape which is AIS_ConnectedInteractive.

I display this tree by recursively going thru all the nodes and using m_Context->Display(shape, false), where m_Context is AIS_InteractiveContext.

Everything works but I cannot change the color of the assembly by using:

colorTool->SetColor(label, color, XCAFDoc_ColorSurf);

The color changes when it gets first loaded, but when I change it later, it stays the same color. I double checked the color by calling:

Quantity_Color col;
colorTool->GetColor(label, XCAFDoc_ColorSurf, col);

The col variable does update to the new color I just set, but the view somehow didn't update it.

In my render loop, I have the following:

m_Context->Update(shape, false);
m_View->Invalidate();
m_View->Redraw();
m_Context->UpdateCurrentViewer();

If I create a primitive shape such as a BRepPrimAPI_MakeBox, everything will work, displaying, color changing, etc. Just not assemblies.

So my question, do I need to use TPrsStd_AISViewer and TPrsStd_AISPresentation to achieve what I want? Or am I missing something obvious here?