diff --git a/src/agentlab/ui_assistant.py b/src/agentlab/ui_assistant.py
index 57916543..36769bd9 100644
--- a/src/agentlab/ui_assistant.py
+++ b/src/agentlab/ui_assistant.py
@@ -7,9 +7,9 @@
 from agentlab.experiments.loop import EnvArgs, ExpArgs
 
 
-def make_exp_args(agent_args: AgentArgs, start_url="https://www.google.com"):
+def make_exp_args(agent_args: GenericAgentArgs, start_url="https://www.google.com"):
     try:
-        agent_args.flags.action.demo_mode = "default"
+        agent_args.flags.action.action_set.demo_mode = "default"
     except AttributeError:
         pass
 
diff --git a/tests/test_ui_assistant.py b/tests/test_ui_assistant.py
index 8b3a6511..85d0fcd8 100644
--- a/tests/test_ui_assistant.py
+++ b/tests/test_ui_assistant.py
@@ -6,4 +6,4 @@ def test_make_exp_args():
     """Basic unit test to detect refactoring errors."""
     exp_args = make_exp_args(AGENT_4o)
 
-    assert exp_args.agent_args.flags.action.demo_mode == "default"
+    assert exp_args.agent_args.flags.action.action_set.demo_mode == "default"