diff --git a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
index 6fdbd922..5465ccfd 100644
--- a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
+++ b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
@@ -83,6 +83,7 @@ public static void chat(
                     new ArrayList<>((List<ChatMessage>) toolCallContext.get(initialRequestId));
 
             messageContext.add(response);
+            toolCallContext.put(initialRequestId, messageContext);
             stm.set(TOOL_CALL_CONTEXT, toolCallContext);
 
             ToolRequestEvent toolRequestEvent =
@@ -182,6 +183,7 @@ public static void processChatRequestOrToolResponse(Event event, RunnerContext c
                                     extraArgs));
                 }
             }
+            toolCallContext.put(initialRequestId, messages);
             // overwrite tool call context
             stm.set(TOOL_CALL_CONTEXT, toolCallContext);