backend/agent/coding_agent.py

@@ -463,6 +463,24 @@ class CodingAgent:
             "is_done":   False,
         }
 
+    def follow_up(self, question: str) -> None:
+        """Continue a completed task with a follow-up question.
+
+        Resets the done-flag and appends the user's question to the existing
+        conversation, so the agent retains full context of what was already done.
+        Call propose_next_action() afterwards to continue the loop.
+        """
+        self.is_done = False
+        self.messages.append({
+            "role": "user",
+            "content": (
+                f"<human_message>{question}</human_message>\n"
+                "<replan>The user has a follow-up question or correction regarding "
+                "the task you just completed. Review what you already did and "
+                "address their question accordingly.</replan>"
+            ),
+        })
+
     def reject(self, feedback: str) -> None:
         """Reject the pending action and inject user feedback.
 

frontend/chat.py

@@ -51,6 +51,15 @@ def _reject_action(feedback: str):
     st.session_state.agent_status = "waiting_approval"
 
 
+def _followup_agent(question: str):
+    """Inject a follow-up question into the finished agent and resume the loop."""
+    agent = st.session_state.coding_agent
+    agent.follow_up(question)
+    action = agent.propose_next_action()
+    st.session_state.agent_pending_action = action
+    st.session_state.agent_status = "waiting_approval"
+
+
 def _reset_agent():
     """Reset all agent state back to idle."""
     st.session_state.coding_agent = None
@@ -72,19 +81,18 @@ def render_agent_mode():
     if agent_log:
         with st.expander(f"Agent Log — {len(agent_log)} step(s) completed", expanded=False):
             for i, step in enumerate(agent_log):
-                st.markdown(f"**Step {i + 1} — `{step['tool']}`**")
-                st.caption(f"Thought: {step['thought']}")
-                if step.get("arguments"):
-                    with st.container():
+                with st.chat_message("assistant"):
+                    st.markdown(f"**Step {i + 1} — `{step['tool']}`**")
+                    st.caption(f"Thought: {step['thought']}")
+                    if step.get("arguments"):
                         st.json(step["arguments"])
-                result_text = step.get("result", "")
-                if result_text.startswith("ERROR") or result_text.startswith("SYNTAX ERROR"):
-                    st.error(result_text)
-                elif result_text.startswith("OK") or result_text.startswith("DONE"):
-                    st.success(result_text)
-                else:
-                    st.code(result_text, language=None)
-                st.divider()
+                    result_text = step.get("result", "")
+                    if result_text.startswith("ERROR") or result_text.startswith("SYNTAX ERROR"):
+                        st.error(result_text)
+                    elif result_text.startswith("OK") or result_text.startswith("DONE"):
+                        st.success(result_text)
+                    else:
+                        st.code(result_text, language=None)
 
     # ── Idle: task input ──────────────────────────────────────────────────────
     if agent_status == "idle":
@@ -106,17 +114,12 @@ def render_agent_mode():
     elif agent_status == "waiting_approval":
         pending = st.session_state.get("agent_pending_action", {})
 
-        st.markdown("**The agent wants to execute the following step:**")
-
-        with st.container(border=True):
-            st.markdown("**Thought**")
-            st.markdown(pending.get("thought", "—"))
-
+        with st.status("Agent proposes the following step:", expanded=True):
+            st.markdown(f"**Thought:** {pending.get('thought', '—')}")
             st.markdown(f"**Tool:** `{pending.get('tool', '—')}`")
 
             args = pending.get("arguments", {})
             if args:
-                st.markdown("**Arguments:**")
                 if "content" in args:
                     display_args = {k: v for k, v in args.items() if k != "content"}
                     if display_args:
@@ -131,7 +134,7 @@ def render_agent_mode():
             placeholder="e.g. Use a different approach...",
         )
 
-        col1, col2, col3 = st.columns([2, 2, 3])
+        col1, col2, col3 = st.columns([3, 2, 2])
         with col1:
             if st.button("Approve", type="primary", use_container_width=True):
                 with st.spinner("Executing and planning next step..."):
@@ -151,49 +154,74 @@ def render_agent_mode():
     elif agent_status == "done":
         last_result = agent_log[-1]["result"] if agent_log else ""
         st.success(f"Task completed! {last_result}")
-        if st.button("New Task", use_container_width=True):
-            _reset_agent()
-            st.rerun()
+
+        st.divider()
+
+        followup = st.text_area(
+            "Follow-up question or correction:",
+            key="agent_followup_input",
+            height=80,
+            placeholder="e.g. The output is wrong — it should sort descending. Can you fix that?",
+        )
+
+        col1, col2 = st.columns(2)
+        with col1:
+            if st.button("Ask Follow-up", type="primary", use_container_width=True):
+                if followup.strip():
+                    with st.spinner("Agent is thinking..."):
+                        _followup_agent(followup.strip())
+                    st.rerun()
+                else:
+                    st.warning("Please enter a follow-up question first.")
+        with col2:
+            if st.button("New Task", use_container_width=True):
+                _reset_agent()
+                st.rerun()
 
 
 # ── Normal Chat ───────────────────────────────────────────────────────────────
 
 def render_normal_chat():
-    chat_section  = st.container()
+    # Chat history as bubbles
+    for message in st.session_state.chat_history:
+        role = message["role"]
+        if role == "system":
+            continue
+        with st.chat_message(role):
+            st.markdown(message["content"])
 
-    with chat_section:
-        if st.session_state.chat_history:
-            for message in st.session_state.chat_history:
-                st.markdown(f"**{message['role'].capitalize()}:** {message['content']}")
+    # Chat input — Enter to send, no extra button needed
+    user_input = st.chat_input("Type your message here...")
+    if user_input:
+        chat_manager = st.session_state.chat_manager
 
-        st.toggle("Agent Mode", key="agent_mode")
+        # Inject system prompt on the first message
+        if not chat_manager.get_history():
+            system_prompt = SystemPrompter.generate_prompt()
+            chat_manager.add_message("system", system_prompt)
 
-        # Clear the input field before the widget is rendered (Streamlit requirement)
-        if st.session_state.get("_clear_chat_input"):
-            st.session_state.chat_input = ""
-            st.session_state._clear_chat_input = False
+        # Show user message immediately without waiting for response
+        with st.chat_message("user"):
+            st.markdown(user_input)
 
-        user_input = st.text_input("Type your message here:", key="chat_input", placeholder="Ask me anything or give me a task to do!", use_container_width=True)
+        # Show response with spinner while API is called
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                try:
+                    ai_response = chat_manager.send_message(user_input)
+                except Exception as e:
+                    ai_response = f"Error: {e}"
+            st.markdown(ai_response)
 
-        if st.button("Send", key="send_button", use_container_width=True) and user_input:
-            chat_manager = st.session_state.chat_manager
-
-            # Inject system prompt on the first message
-            if not chat_manager.get_history():
-                system_prompt = SystemPrompter.generate_prompt()
-                chat_manager.add_message("system", system_prompt)
-
-            st.session_state.chat_history.append({"role": "user", "content": user_input})
-
-            try:
-                ai_response = chat_manager.send_message(user_input)
-            except Exception as e:
-                ai_response = f"Error: {e}"
-
-            st.session_state.chat_history.append({"role": "assistant", "content": ai_response})
-            st.session_state._clear_chat_input = True
-            st.rerun()
+        st.session_state.chat_history.append({"role": "user", "content": user_input})
+        st.session_state.chat_history.append({"role": "assistant", "content": ai_response})
+        st.rerun()
 
+    # Rendered in the normal flow; JS above clones them to fixed positions
+    # and hides these originals.
+    st.toggle("Agent Mode", key="agent_mode")
+    with st.expander("Settings", expanded=False):
+        st.toggle("Use debug system prompt", key="use_system_prompt", value=True)
 
 
 # ── Entry point ───────────────────────────────────────────────────────────────