diff --git a/api/main.py b/api/main.py
index 1760458..a22405f 100644
--- a/api/main.py
+++ b/api/main.py
@@ -9,7 +9,7 @@ import asyncio
 
 from workflows.agent_goal_workflow import AgentGoalWorkflow
 from models.data_types import CombinedInput, AgentGoalWorkflowParams
-from tools.goal_registry import goal_match_train_invoice, goal_event_flight_invoice, goal_choose_agent_type
+from tools.goal_registry import goal_list
 from fastapi.middleware.cors import CORSMiddleware
 from shared.config import get_temporal_client, TEMPORAL_TASK_QUEUE, AGENT_GOAL
 
@@ -22,17 +22,13 @@ load_dotenv()
 
 def get_initial_agent_goal():
     """Get the agent goal from environment variables."""
-    goals = {
-        "goal_match_train_invoice": goal_match_train_invoice,
-        "goal_event_flight_invoice": goal_event_flight_invoice,
-        "goal_choose_agent_type": goal_choose_agent_type,
-    }
-
     if AGENT_GOAL is not None:
-        return goals.get(AGENT_GOAL)
+        for listed_goal in goal_list:
+            if listed_goal.id == AGENT_GOAL:
+                return listed_goal
     else:
-        #if no goal is set in the env file, default to event/flight use case
-        return goals.get("goal_event_flight_invoice", goal_event_flight_invoice)
+        #if no goal is set in the config file, default to choosing an agent
+        return goal_list.get("goal_choose_agent_type")
 
 
 @app.on_event("startup")
diff --git a/tools/goal_registry.py b/tools/goal_registry.py
index 00be732..39e2187 100644
--- a/tools/goal_registry.py
+++ b/tools/goal_registry.py
@@ -50,6 +50,7 @@ goal_match_train_invoice = AgentGoal(
         search_trains_tool,
         book_trains_tool,
         create_invoice_tool,
+        list_agents_tool, #last tool must be list_agents to fasciliate changing back to picking an agent again at the end
     ],
     description="The user wants to book a trip to a city in the UK around the dates of a premier league match. "
     "Help the user find a premier league match to attend, search and book trains for that match and offers to invoice them for the cost of train tickets. "
@@ -95,7 +96,7 @@ goal_event_flight_invoice = AgentGoal(
         find_events_tool,
         search_flights_tool,
         create_invoice_tool,
-        list_agents_tool,
+        list_agents_tool, #last tool must be list_agents to fasciliate changing back to picking an agent again at the end
     ],
     description="Help the user gather args for these tools in order: "
     "1. FindEvents: Find an event to travel to "
diff --git a/workflows/agent_goal_workflow.py b/workflows/agent_goal_workflow.py
index 56f34a0..c433463 100644
--- a/workflows/agent_goal_workflow.py
+++ b/workflows/agent_goal_workflow.py
@@ -180,7 +180,6 @@ class AgentGoalWorkflow:
                     if await helpers.handle_missing_args(current_tool, args, tool_data, self.prompt_queue):
                         continue
 
-# Would swapping these two get rid of the confirm button step?
                     waiting_for_confirm = True
                     self.confirm = False
                     workflow.logger.info("Waiting for user confirm signal...")