diff --git a/backend/agent/coding_agent.py b/backend/agent/coding_agent.py
index 65307ad..f2d8898 100644
--- a/backend/agent/coding_agent.py
+++ b/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"{question}\n"
+ "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."
+ ),
+ })
+
def reject(self, feedback: str) -> None:
"""Reject the pending action and inject user feedback.
diff --git a/frontend/chat.py b/frontend/chat.py
index c39f4f5..7716a04 100644
--- a/frontend/chat.py
+++ b/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 ───────────────────────────────────────────────────────────────