change logic of show AIMessages

app.py

@@ -99,21 +99,29 @@ async def chat_logic(
     # app_input["github_repo"] = github_repo
     # Stream the graph's response.
     # This revised logic handles intermediate messages and prevents duplication.
+    
+    final_reply = None          # buffer for the last AIMessage we see
+
     async for res in app.astream(app_input, config=config, stream_mode="values"):
-        if "messages" in res:
-            last_message = res["messages"][-1]
-            # We only stream content from AIMessages. Any intermediate AIMessages
-            # (e.g., "I will now use a tool") will be overwritten by subsequent
-            # AIMessages in the UI, so only the final answer is visible.
-            if isinstance(last_message, AIMessage):
-                yield last_message.content
-
-        elif "__interrupt__" in res:
-            # Handle interruptions where the agent needs human feedback.
-            interruption_message = res["__interrupt__"][0]
+        # ── 1) Handle human-interrupts immediately ─────────────────────────
+        if "__interrupt__" in res:
             session_state["is_message_command"] = True
-            yield interruption_message.value
-            return  # Stop the stream and wait for the user's next message.
+            # yield the interrupt text straight away
+            yield res["__interrupt__"][0].value
+            return                        # stop processing until user replies
+
+        # ── 2) Remember the latest AIMessage we’ve seen ────────────────────
+        if "messages" in res:
+            last = res["messages"][-1]
+            if isinstance(last, AIMessage):
+                final_reply = last.content
+
+    # ── 3) After the graph stops, emit the buffered final answer ───────────
+    if final_reply is not None:
+        yield final_reply                 # exactly one assistant chunk
+    else:
+        # fail-safe: graph produced no AIMessage
+        yield "✅ Done"
         
 
 def create_gradio_app():