diff --git a/lang_agent/graphs/graph_states.py b/lang_agent/graphs/graph_states.py
index 4ff3963..47f92e2 100644
--- a/lang_agent/graphs/graph_states.py
+++ b/lang_agent/graphs/graph_states.py
@@ -14,3 +14,4 @@ class ChattyToolState(TypedDict):
                Dict[str, Dict[str, str|int]]]
     tool_messages: List[SystemMessage | HumanMessage | AIMessage]
     chatty_messages: List[SystemMessage | HumanMessage | AIMessage]
+    tool_done: bool  # Flag to signal when tool execution is complete