sessions history, chat/unified/stream

api_routes_v2.py

@@ -505,6 +505,26 @@ def _normalize_history_for_api(chat_id: str) -> List[Message]:
     return history
 
 
+def _add_and_mirror_message(chat_id: str, role: str, content: str):
+    """
+    Add a message via session_manager AND mirror it into the session doc
+    so history is always available in responses.
+    """
+    session_manager.add_message(chat_id, role, content)
+    try:
+        s = session_manager.get_session(chat_id) or {}
+        msgs = list(s.get("messages", []))
+        msgs.append({
+            "role": role,
+            "content": content if isinstance(content, str) else json.dumps(content, ensure_ascii=False),
+            "timestamp": datetime.utcnow().isoformat() + "Z",
+        })
+        session_manager.update_session(chat_id, {"messages": msgs})
+    except Exception:
+        # Don't block the flow on mirror issues
+        pass
+
+
 def _assistant_response_payload(
     chat_id: str,
     friendly_response: str,
@@ -513,7 +533,7 @@ def _assistant_response_payload(
     state: str
 ) -> ChatResponse:
     # Persist assistant message
-    session_manager.add_message(chat_id, "assistant", friendly_response)
+    _add_and_mirror_message(chat_id, "assistant", friendly_response)
     # Return full history in role/content shape
     history = _normalize_history_for_api(chat_id)
     return ChatResponse(
@@ -610,7 +630,7 @@ def upload_stream_to_s3(chat_id: str, file: UploadFile) -> str:
 
     s3_uri = f"s3://{S3_BUCKET}/{key}"
     session_manager.update_session(chat_id, {"current_file": s3_uri, "state": "initial"})
-    session_manager.add_message(chat_id, "system", f"File uploaded to S3: {s3_uri}")
+    _add_and_mirror_message(chat_id, "system", f"File uploaded to S3: {s3_uri}")
     return s3_uri
 
 
@@ -651,7 +671,7 @@ async def chat_unified(
     # If JSON included a file_path (e.g., s3://...), attach it
     if file_path_from_json:
         session_manager.update_session(chat_id, {"current_file": file_path_from_json})
-        session_manager.add_message(chat_id, "system", f"File attached: {file_path_from_json}")
+        _add_and_mirror_message(chat_id, "system", f"File attached: {file_path_from_json}")
         session = _get_session_or_init(chat_id)
 
     # If a file is included in the form, upload to S3 and attach it
@@ -674,7 +694,7 @@ async def chat_unified(
         return _assistant_response_payload(chat_id, friendly, {"intent": "missing_message"}, api_data, session.get("state", "initial"))
 
     # Add user message
-    session_manager.add_message(chat_id, "user", message)
+    _add_and_mirror_message(chat_id, "user", message)
 
     # Classify intent
     intent_data = intent_classifier.classify_intent(message)
@@ -907,7 +927,7 @@ async def chat_unified(
 # ========================
 
 @router.post("/chat/unified/stream")
-def chat_unified_stream(
+async def chat_unified_stream(
     request: Request,
     chat_id: Optional[str] = Form(None),
     message: Optional[str] = Form(None),
@@ -921,277 +941,258 @@ def chat_unified_stream(
     - On approval, streams execution progress and final result.
     """
 
-    async def prepare():
-        # Parse JSON if needed
-        content_type = (request.headers.get("content-type") or "").lower()
-        file_path_from_json = None
-        _chat_id, _message, _prefer_bedrock, _file = chat_id, message, prefer_bedrock, file
-        if "application/json" in content_type:
-            body = await request.json()
-            _chat_id = body.get("chat_id") or _chat_id
-            _message = body.get("message") if "message" in body else _message
-            _prefer_bedrock = body.get("prefer_bedrock", True) if "prefer_bedrock" in body else _prefer_bedrock
-            file_path_from_json = body.get("file_path")
-
-        _chat_id = _ensure_chat(_chat_id)
-        _session = _get_session_or_init(_chat_id)
-
-        # Attach JSON file path if provided
-        if file_path_from_json:
-            session_manager.update_session(_chat_id, {"current_file": file_path_from_json})
-            session_manager.add_message(_chat_id, "system", f"File attached: {file_path_from_json}")
-            _session = _get_session_or_init(_chat_id)
-
-        # Upload file if provided (form)
-        uploaded_file_info = None
-        if _file is not None:
-            s3_uri = upload_stream_to_s3(_chat_id, _file)
-            uploaded_file_info = {"bucket": S3_BUCKET, "key": s3_uri.split(f"s3://{S3_BUCKET}/", 1)[1], "s3_uri": s3_uri}
-            _session = _get_session_or_init(_chat_id)
-
-        return _chat_id, _message, _prefer_bedrock, uploaded_file_info
-
-    def make_stream(chat_id_local: str, msg: Optional[str], prefer_bedrock_local: bool, uploaded_file_info: Optional[Dict[str, Any]]):
-        def emit(obj: Dict[str, Any]):
-            obj.setdefault("chat_id", chat_id_local)
-            obj.setdefault("state", session_manager.get_session(chat_id_local).get("state", "initial"))
-            line = json.dumps(obj, ensure_ascii=False).encode("utf-8") + b"\n"
-            return line
+    # Parse JSON if needed
+    content_type = (request.headers.get("content-type") or "").lower()
+    file_path_from_json = None
+    if "application/json" in content_type:
+        body = await request.json()
+        chat_id = body.get("chat_id") or chat_id
+        message = body.get("message") if "message" in body else message
+        prefer_bedrock = body.get("prefer_bedrock", True) if "prefer_bedrock" in body else prefer_bedrock
+        file_path_from_json = body.get("file_path")
+
+    chat_id = _ensure_chat(chat_id)
+    session = _get_session_or_init(chat_id)
+
+    # Attach JSON file path if provided
+    if file_path_from_json:
+        session_manager.update_session(chat_id, {"current_file": file_path_from_json})
+        _add_and_mirror_message(chat_id, "system", f"File attached: {file_path_from_json}")
+        session = _get_session_or_init(chat_id)
 
-        def gen() -> Generator[bytes, None, None]:
-            session = _get_session_or_init(chat_id_local)
+    # Upload file if provided
+    uploaded_file_info = None
+    if file is not None:
+        s3_uri = upload_stream_to_s3(chat_id, file)
+        uploaded_file_info = {"bucket": S3_BUCKET, "key": s3_uri.split(f"s3://{S3_BUCKET}/", 1)[1], "s3_uri": s3_uri}
+        session = _get_session_or_init(chat_id)
+
+    def emit(obj: Dict[str, Any]) -> bytes:
+        obj.setdefault("chat_id", chat_id)
+        obj.setdefault("state", session_manager.get_session(chat_id).get("state", "initial"))
+        return (json.dumps(obj, ensure_ascii=False) + "\n").encode("utf-8")
+
+    def stream_gen() -> Generator[bytes, None, None]:
+        session_local = _get_session_or_init(chat_id)
+
+        # Only-file case
+        if (message is None or str(message).strip() == "") and uploaded_file_info:
+            friendly = "📁 File uploaded successfully. Tell me what you'd like to do with it (e.g., extract text, get tables, summarize)."
+            _add_and_mirror_message(chat_id, "assistant", friendly)
+            yield emit({"type": "assistant_final", "content": friendly, "file": uploaded_file_info, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+            return
+
+        # No message
+        if message is None or str(message).strip() == "":
+            friendly = "Please provide a message (e.g., 'extract text', 'get tables', 'summarize')."
+            _add_and_mirror_message(chat_id, "assistant", friendly)
+            yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+            return
+
+        # Add user message
+        _add_and_mirror_message(chat_id, "user", message)
+
+        # Classify
+        intent_data = intent_classifier.classify_intent(message)
+        current_state = session_local.get("state", "initial")
+
+        # Casual / question / unclear
+        if intent_data["intent"] in {"casual_chat", "question", "unclear"} and current_state == "initial":
+            friendly = intent_classifier.get_friendly_response(intent_data["intent"], message)
+            _add_and_mirror_message(chat_id, "assistant", friendly)
+            yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+            return
 
-            # If only a file was uploaded and no message, acknowledge
-            if (msg is None or str(msg).strip() == "") and uploaded_file_info:
-                friendly = "📁 File uploaded successfully. Tell me what you'd like to do with it (e.g., extract text, get tables, summarize)."
-                session_manager.add_message(chat_id_local, "assistant", friendly)
-                yield emit({"type": "assistant_final", "content": friendly, "file": uploaded_file_info, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+        # Initial: nudge or generate plan
+        if current_state == "initial":
+            if not intent_data.get("requires_pipeline", False):
+                friendly = (
+                    "I'm here to help process documents! Please tell me what you'd like to do with your document.\n\n"
+                    "For example:\n- 'extract text and summarize'\n- 'get tables from pages 2-5'\n- 'translate to Spanish'\n\n"
+                    "Type 'help' to see all capabilities!"
+                )
+                _add_and_mirror_message(chat_id, "assistant", friendly)
+                yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
-            # If no message at all, nudge
-            if msg is None or str(msg).strip() == "":
-                friendly = "Please provide a message (e.g., 'extract text', 'get tables', 'summarize')."
-                session_manager.add_message(chat_id_local, "assistant", friendly)
-                yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+            if not session_local.get("current_file"):
+                friendly = "📁 Please upload a document first before I can process it!"
+                _add_and_mirror_message(chat_id, "assistant", friendly)
+                yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
-            # Add user message
-            session_manager.add_message(chat_id_local, "user", msg)
+            yield emit({"type": "status", "message": "Analyzing request and creating a pipeline..."})
+            try:
+                pipeline = generate_pipeline(
+                    user_input=message,
+                    file_path=session_local.get("current_file"),
+                    prefer_bedrock=bool(prefer_bedrock),
+                )
+                session_manager.update_session(chat_id, {"proposed_pipeline": pipeline, "state": "pipeline_proposed"})
 
-            # Classify
-            intent_data = intent_classifier.classify_intent(msg)
-            current_state = session.get("state", "initial")
+                pipeline_name = pipeline.get("pipeline_name", "Document Processing")
+                steps_list = pipeline.get("pipeline_steps", [])
+                steps_summary = "\n".join([f"  {i+1}. {step.get('tool', 'Unknown')}" for i, step in enumerate(steps_list)])
 
-            # Casual / question / unclear at initial
-            if intent_data["intent"] in {"casual_chat", "question", "unclear"} and current_state == "initial":
-                friendly = intent_classifier.get_friendly_response(intent_data["intent"], msg)
-                session_manager.add_message(chat_id_local, "assistant", friendly)
-                yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+                friendly = (
+                    f"🎯 **Pipeline Created: {pipeline_name}**\n"
+                    f"Here's what I'll do:\n{steps_summary}\n"
+                    f"**Ready to proceed?**\n"
+                    f"- Type 'approve' or 'yes' to execute\n"
+                    f"- Type 'reject' or 'no' to cancel\n"
+                    f"- Describe changes to modify the plan"
+                )
+                _add_and_mirror_message(chat_id, "assistant", friendly)
+                yield emit({"type": "assistant_final", "content": friendly, "pipeline": pipeline, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+                return
+            except Exception as e:
+                friendly = f"❌ Error generating pipeline: {str(e)}"
+                _add_and_mirror_message(chat_id, "assistant", friendly)
+                yield emit({"type": "assistant_final", "content": friendly, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
-            # Initial: pipeline request or nudge
-            if current_state == "initial":
-                if not intent_data.get("requires_pipeline", False):
-                    friendly = (
-                        "I'm here to help process documents! Please tell me what you'd like to do with your document.\n\n"
-                        "For example:\n- 'extract text and summarize'\n- 'get tables from pages 2-5'\n- 'translate to Spanish'\n\n"
-                        "Type 'help' to see all capabilities!"
-                    )
-                    session_manager.add_message(chat_id_local, "assistant", friendly)
-                    yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                    return
+        # Pipeline proposed
+        if current_state == "pipeline_proposed":
+            if intent_data["intent"] == "approval":
+                session_manager.update_session(chat_id, {"state": "executing"})
+                plan = session_local.get("proposed_pipeline", {})
+                initial = (
+                    f"✅ Approved! Starting execution of: **{plan.get('pipeline_name', 'pipeline')}**\n\n"
+                    f"🚀 Processing, please wait...\n_(Using {plan.get('_generator', 'AI')} - {plan.get('_model', 'model')})_"
+                )
+                yield emit({"type": "assistant_delta", "content": initial})
 
-                if not session.get("current_file"):
-                    friendly = "📁 Please upload a document first before I can process it!\n\nClick 'Upload Document' to get started."
-                    session_manager.add_message(chat_id_local, "assistant", friendly)
-                    yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                    return
+                steps_completed, final_payload, executor_used = [], None, "unknown"
+                accumulated = initial
+                file_ref = session_local.get("current_file")
+                local_path, cleanup = download_to_temp_file(file_ref)
 
-                # Generate pipeline (no need to download file)
-                yield emit({"type": "status", "message": "Analyzing request and creating a pipeline..."})
                 try:
-                    pipeline = generate_pipeline(
-                        user_input=msg,
-                        file_path=session.get("current_file"),
-                        prefer_bedrock=bool(prefer_bedrock_local)
-                    )
-                    session_manager.update_session(chat_id_local, {"proposed_pipeline": pipeline, "state": "pipeline_proposed"})
-
-                    pipeline_name = pipeline.get("pipeline_name", "Document Processing")
-                    steps_list = pipeline.get("pipeline_steps", [])
-                    steps_summary = "\n".join([f"  {i+1}. {step.get('tool', 'Unknown')}" for i, step in enumerate(steps_list)])
-
-                    friendly = (
-                        f"🎯 **Pipeline Created: {pipeline_name}**\n"
-                        f"Here's what I'll do:\n{steps_summary}\n"
-                        f"**Ready to proceed?**\n"
-                        f"- Type 'approve' or 'yes' to execute\n"
-                        f"- Type 'reject' or 'no' to cancel\n"
-                        f"- Describe changes to modify the plan"
-                    )
+                    for event in execute_pipeline_streaming(
+                        pipeline=plan,
+                        file_path=local_path,
+                        session_id=chat_id,
+                        prefer_bedrock=bool(prefer_bedrock)
+                    ):
+                        etype = event.get("type")
 
-                    session_manager.add_message(chat_id_local, "assistant", friendly)
-                    yield emit({"type": "assistant_final", "content": friendly, "pipeline": pipeline, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                    return
-                except Exception as e:
-                    friendly = f"❌ Error generating pipeline: {str(e)}"
-                    session_manager.add_message(chat_id_local, "assistant", friendly)
-                    yield emit({"type": "assistant_final", "content": friendly, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                    return
+                        if etype == "info":
+                            msg2 = f"ℹ️ {event.get('message')} _(Executor: {event.get('executor', 'unknown')})_"
+                            accumulated += "\n\n" + msg2
+                            yield emit({"type": "assistant_delta", "content": accumulated})
 
-            # Pipeline proposed: handle approval, rejection, or edit
-            if current_state == "pipeline_proposed":
-                if intent_data["intent"] == "approval":
-                    session_manager.update_session(chat_id_local, {"state": "executing"})
-                    plan = session.get("proposed_pipeline", {})
-                    initial = (
-                        f"✅ Approved! Starting execution of: **{plan.get('pipeline_name', 'pipeline')}**\n\n"
-                        f"🚀 Processing, please wait...\n_(Using {plan.get('_generator', 'AI')} - {plan.get('_model', 'model')})_"
-                    )
-                    yield emit({"type": "assistant_delta", "content": initial})
+                        elif etype == "step":
+                            step_num = event.get("step", 0)
+                            tool_name = event.get("tool", "processing")
+                            status = event.get("status", "running")
+                            if status == "completed" and "observation" in event:
+                                obs_preview = str(event.get("observation"))[:80]
+                                step_msg = f"✅ Step {step_num}: {tool_name} - Completed!\n   Preview: {obs_preview}..."
+                            elif status == "executing":
+                                step_msg = f"⏳ Step {step_num}: {tool_name} - Processing..."
+                            else:
+                                step_msg = f"📍 Step {step_num}: {tool_name}"
+
+                            steps_completed.append({
+                                "step": step_num,
+                                "tool": tool_name,
+                                "status": status,
+                                "executor": event.get("executor", "unknown"),
+                                "observation": event.get("observation"),
+                                "input": event.get("input"),
+                            })
+                            executor_used = event.get("executor", executor_used)
 
-                    steps_completed = []
-                    final_payload = None
-                    executor_used = "unknown"
-                    accumulated = initial
+                            accumulated += "\n\n" + step_msg
+                            yield emit({"type": "assistant_delta", "content": accumulated})
 
-                    file_ref = session.get("current_file")
-                    local_path, cleanup = download_to_temp_file(file_ref)
+                        elif etype == "final":
+                            final_payload = event.get("data")
+                            executor_used = event.get("executor", executor_used)
 
-                    try:
-                        for event in execute_pipeline_streaming(
-                            pipeline=plan,
-                            file_path=local_path,
-                            session_id=chat_id_local,
-                            prefer_bedrock=bool(prefer_bedrock_local)
-                        ):
-                            etype = event.get("type")
-
-                            if etype == "info":
-                                msg2 = f"ℹ️ {event.get('message')} _(Executor: {event.get('executor', 'unknown')})_"
-                                accumulated += "\n\n" + msg2
-                                yield emit({"type": "assistant_delta", "content": accumulated})
-
-                            elif etype == "step":
-                                step_num = event.get("step", 0)
-                                tool_name = event.get("tool", "processing")
-                                status = event.get("status", "running")
-                                if status == "completed" and "observation" in event:
-                                    obs_preview = str(event.get("observation"))[:80]
-                                    step_msg = f"✅ Step {step_num}: {tool_name} - Completed!\n   Preview: {obs_preview}..."
-                                elif status == "executing":
-                                    step_msg = f"⏳ Step {step_num}: {tool_name} - Processing..."
-                                else:
-                                    step_msg = f"📍 Step {step_num}: {tool_name}"
-
-                                steps_completed.append({
-                                    "step": step_num,
-                                    "tool": tool_name,
-                                    "status": status,
-                                    "executor": event.get("executor", "unknown"),
-                                    "observation": event.get("observation"),
-                                    "input": event.get("input"),
-                                })
-                                executor_used = event.get("executor", executor_used)
-
-                                accumulated += "\n\n" + step_msg
-                                yield emit({"type": "assistant_delta", "content": accumulated})
-
-                            elif etype == "final":
-                                final_payload = event.get("data")
-                                executor_used = event.get("executor", executor_used)
-
-                            elif etype == "error":
-                                err = event.get("error", "Unknown error")
-                                friendly_err = f"❌ Pipeline Failed\n\nError: {err}\n\nCompleted {len(steps_completed)} step(s) before failure."
-                                session_manager.update_session(chat_id_local, {"state": "initial"})
-                                session_manager.add_message(chat_id_local, "assistant", friendly_err)
-                                yield emit({"type": "assistant_final", "content": friendly_err, "error": err, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                                return
-
-                        # Finalize
-                        if final_payload:
-                            session_manager.update_session(chat_id_local, {"pipeline_result": final_payload, "state": "initial"})
-                            session_manager.save_pipeline_execution(
-                                session_id=chat_id_local,
-                                pipeline=plan,
-                                result=final_payload,
-                                file_path=file_ref,
-                                executor=executor_used
-                            )
-                            success_count = len([s for s in steps_completed if s.get("status") == "completed"])
-                            friendly_final = (
-                                f"🎉 Pipeline Completed Successfully!\n"
-                                f"- Pipeline: {plan.get('pipeline_name', 'Document Processing')}\n"
-                                f"- Total Steps: {len(steps_completed)}\n"
-                                f"- Successful: {success_count}\n"
-                                f"- Executor: {executor_used}\n"
-                                f"✅ All done! What else would you like me to help you with?"
-                            )
-                            session_manager.add_message(chat_id_local, "assistant", friendly_final)
-                            yield emit({"type": "assistant_final", "content": friendly_final, "result": final_payload, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                            return
-                        else:
-                            done = f"✅ Pipeline Completed! Executed {len(steps_completed)} steps using {executor_used}."
-                            session_manager.update_session(chat_id_local, {"state": "initial"})
-                            session_manager.add_message(chat_id_local, "assistant", done)
-                            yield emit({"type": "assistant_final", "content": done, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+                        elif etype == "error":
+                            err = event.get("error", "Unknown error")
+                            friendly_err = f"❌ Pipeline Failed\n\nError: {err}\n\nCompleted {len(steps_completed)} step(s) before failure."
+                            session_manager.update_session(chat_id, {"state": "initial"})
+                            _add_and_mirror_message(chat_id, "assistant", friendly_err)
+                            yield emit({"type": "assistant_final", "content": friendly_err, "error": err, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                             return
 
-                    except Exception as e:
-                        friendly_err = f"❌ Pipeline Execution Failed\n\nError: {str(e)}\n\nCompleted {len(steps_completed)} step(s) before failure."
-                        session_manager.update_session(chat_id_local, {"state": "initial"})
-                        session_manager.add_message(chat_id_local, "assistant", friendly_err)
-                        yield emit({"type": "assistant_final", "content": friendly_err, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                        return
-                    finally:
-                        try:
-                            cleanup()
-                        except Exception:
-                            pass
-
-                elif intent_data["intent"] == "rejection":
-                    session_manager.update_session(chat_id_local, {"state": "initial", "proposed_pipeline": None})
-                    friendly = "👍 No problem! Pipeline cancelled. What else would you like me to help you with?"
-                    session_manager.add_message(chat_id_local, "assistant", friendly)
-                    yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
-                    return
-
-                else:
-                    # Treat as edit/modify
-                    try:
-                        original_plan = session.get("proposed_pipeline", {})
-                        edit_context = f"Original: {original_plan.get('pipeline_name')}. User wants: {msg}"
-                        new_pipeline = generate_pipeline(
-                            user_input=edit_context,
-                            file_path=session.get("current_file"),
-                            prefer_bedrock=bool(prefer_bedrock_local)
+                    if final_payload:
+                        session_manager.update_session(chat_id, {"pipeline_result": final_payload, "state": "initial"})
+                        session_manager.save_pipeline_execution(
+                            session_id=chat_id,
+                            pipeline=plan,
+                            result=final_payload,
+                            file_path=file_ref,
+                            executor=executor_used
+                        )
+                        success_count = len([s for s in steps_completed if s.get("status") == "completed"])
+                        friendly_final = (
+                            f"🎉 Pipeline Completed Successfully!\n"
+                            f"- Pipeline: {plan.get('pipeline_name', 'Document Processing')}\n"
+                            f"- Total Steps: {len(steps_completed)}\n"
+                            f"- Successful: {success_count}\n"
+                            f"- Executor: {executor_used}\n"
+                            f"✅ All done! What else would you like me to help you with?"
                         )
-                        session_manager.update_session(chat_id_local, {"proposed_pipeline": new_pipeline, "state": "pipeline_proposed"})
-                        formatted = format_pipeline_for_display(new_pipeline)
-                        friendly = formatted + f"\n\n```json\n{json.dumps(new_pipeline, indent=2)}\n```"
-                        session_manager.add_message(chat_id_local, "assistant", friendly)
-                        yield emit({"type": "assistant_final", "content": friendly, "pipeline": new_pipeline, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+                        _add_and_mirror_message(chat_id, "assistant", friendly_final)
+                        yield emit({"type": "assistant_final", "content": friendly_final, "result": final_payload, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                         return
-                    except Exception as e:
-                        friendly = f"❌ Edit failed: {str(e)}"
-                        session_manager.add_message(chat_id_local, "assistant", friendly)
-                        yield emit({"type": "assistant_final", "content": friendly, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+                    else:
+                        done = f"✅ Pipeline Completed! Executed {len(steps_completed)} steps."
+                        session_manager.update_session(chat_id, {"state": "initial"})
+                        _add_and_mirror_message(chat_id, "assistant", done)
+                        yield emit({"type": "assistant_final", "content": done, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                         return
 
-            # Default
-            friendly = "Please upload a document and tell me what you'd like me to do (e.g., extract text, summarize, translate)."
-            session_manager.add_message(chat_id_local, "assistant", friendly)
-            yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id_local)]})
+                except Exception as e:
+                    friendly_err = f"❌ Pipeline Execution Failed\n\nError: {str(e)}"
+                    session_manager.update_session(chat_id, {"state": "initial"})
+                    _add_and_mirror_message(chat_id, "assistant", friendly_err)
+                    yield emit({"type": "assistant_final", "content": friendly_err, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+                    return
+                finally:
+                    try:
+                        cleanup()
+                    except Exception:
+                        pass
+
+            elif intent_data["intent"] == "rejection":
+                session_manager.update_session(chat_id, {"state": "initial", "proposed_pipeline": None})
+                friendly = "👍 No problem! Pipeline cancelled. What else would you like me to help you with?"
+                _add_and_mirror_message(chat_id, "assistant", friendly)
+                yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+                return
 
-        return gen()
+            else:
+                try:
+                    original_plan = session_local.get("proposed_pipeline", {})
+                    edit_context = f"Original: {original_plan.get('pipeline_name')}. User wants: {message}"
+                    new_pipeline = generate_pipeline(
+                        user_input=edit_context,
+                        file_path=session_local.get("current_file"),
+                        prefer_bedrock=bool(prefer_bedrock)
+                    )
+                    session_manager.update_session(chat_id, {"proposed_pipeline": new_pipeline, "state": "pipeline_proposed"})
+                    formatted = format_pipeline_for_display(new_pipeline)
+                    friendly = formatted + f"\n\n```json\n{json.dumps(new_pipeline, indent=2)}\n```"
+                    _add_and_mirror_message(chat_id, "assistant", friendly)
+                    yield emit({"type": "assistant_final", "content": friendly, "pipeline": new_pipeline, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+                    return
+                except Exception as e:
+                    friendly = f"❌ Edit failed: {str(e)}"
+                    _add_and_mirror_message(chat_id, "assistant", friendly)
+                    yield emit({"type": "assistant_final", "content": friendly, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
+                    return
 
-    async def stream_wrapper():
-        chat_id_local, msg_local, prefer_local, uploaded_info = await prepare()
-        return StreamingResponse(make_stream(chat_id_local, msg_local, bool(prefer_local), uploaded_info), media_type="application/x-ndjson")
+        # Default
+        friendly = "Please upload a document and tell me what you'd like me to do (e.g., extract text, summarize, translate)."
+        _add_and_mirror_message(chat_id, "assistant", friendly)
+        yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
 
-    # Note: for FastAPI sync path, we return the awaitable response
-    return stream_wrapper()
+    # Return a real StreamingResponse
+    return StreamingResponse(stream_gen(), media_type="application/x-ndjson")
 
 
 # ========================
@@ -1211,7 +1212,7 @@ async def smart_chat(request: ChatRequest):
         session_manager.update_session(chat_id, {"current_file": request.file_path})
         session = _get_session_or_init(chat_id)
 
-    session_manager.add_message(chat_id, "user", request.message)
+    _add_and_mirror_message(chat_id, "user", request.message)
 
     intent_data = intent_classifier.classify_intent(request.message)
     current_state = session.get("state", "initial")
@@ -1422,7 +1423,7 @@ def smart_chat_stream(request: ChatRequest):
             session_manager.update_session(chat_id, {"current_file": request.file_path})
             session = _get_session_or_init(chat_id)
 
-        session_manager.add_message(chat_id, "user", request.message)
+        _add_and_mirror_message(chat_id, "user", request.message)
 
         intent_data = intent_classifier.classify_intent(request.message)
         current_state = session.get("state", "initial")
@@ -1435,7 +1436,7 @@ def smart_chat_stream(request: ChatRequest):
 
         if intent_data["intent"] in {"casual_chat", "question", "unclear"} and current_state == "initial":
             friendly = intent_classifier.get_friendly_response(intent_data["intent"], request.message)
-            session_manager.add_message(chat_id, "assistant", friendly)
+            _add_and_mirror_message(chat_id, "assistant", friendly)
             yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
             return
 
@@ -1446,13 +1447,13 @@ def smart_chat_stream(request: ChatRequest):
                     "For example:\n- 'extract text and summarize'\n- 'get tables from pages 2-5'\n- 'translate to Spanish'\n\n"
                     "Type 'help' to see all capabilities!"
                 )
-                session_manager.add_message(chat_id, "assistant", friendly)
+                _add_and_mirror_message(chat_id, "assistant", friendly)
                 yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
             if not session.get("current_file"):
                 friendly = "📁 Please upload a document first before I can process it!\n\nClick 'Upload Document' to get started."
-                session_manager.add_message(chat_id, "assistant", friendly)
+                _add_and_mirror_message(chat_id, "assistant", friendly)
                 yield emit({"type": "assistant_final", "content": friendly, "intent": intent_data, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
@@ -1478,12 +1479,12 @@ def smart_chat_stream(request: ChatRequest):
                     f"- Describe changes to modify the plan"
                 )
 
-                session_manager.add_message(chat_id, "assistant", friendly)
+                _add_and_mirror_message(chat_id, "assistant", friendly)
                 yield emit({"type": "assistant_final", "content": friendly, "pipeline": pipeline, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
             except Exception as e:
                 friendly = f"❌ Error generating pipeline: {str(e)}"
-                session_manager.add_message(chat_id, "assistant", friendly)
+                _add_and_mirror_message(chat_id, "assistant", friendly)
                 yield emit({"type": "assistant_final", "content": friendly, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
@@ -1552,7 +1553,7 @@ def smart_chat_stream(request: ChatRequest):
                             err = event.get("error", "Unknown error")
                             friendly_err = f"❌ Pipeline Failed\n\nError: {err}\n\nCompleted {len(steps_completed)} step(s) before failure."
                             session_manager.update_session(chat_id, {"state": "initial"})
-                            session_manager.add_message(chat_id, "assistant", friendly_err)
+                            _add_and_mirror_message(chat_id, "assistant", friendly_err)
                             yield emit({"type": "assistant_final", "content": friendly_err, "error": err, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                             return
 
@@ -1574,20 +1575,20 @@ def smart_chat_stream(request: ChatRequest):
                             f"- Executor: {executor_used}\n"
                             f"✅ All done! What else would you like me to help you with?"
                         )
-                        session_manager.add_message(chat_id, "assistant", friendly_final)
+                        _add_and_mirror_message(chat_id, "assistant", friendly_final)
                         yield emit({"type": "assistant_final", "content": friendly_final, "result": final_payload, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                         return
                     else:
                         done = f"✅ Pipeline Completed! Executed {len(steps_completed)} steps using {executor_used}."
                         session_manager.update_session(chat_id, {"state": "initial"})
-                        session_manager.add_message(chat_id, "assistant", done)
+                        _add_and_mirror_message(chat_id, "assistant", done)
                         yield emit({"type": "assistant_final", "content": done, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                         return
 
                 except Exception as e:
                     friendly_err = f"❌ Pipeline Execution Failed\n\nError: {str(e)}\n\nCompleted {len(steps_completed)} step(s) before failure."
                     session_manager.update_session(chat_id, {"state": "initial"})
-                    session_manager.add_message(chat_id, "assistant", friendly_err)
+                    _add_and_mirror_message(chat_id, "assistant", friendly_err)
                     yield emit({"type": "assistant_final", "content": friendly_err, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                     return
                 finally:
@@ -1599,7 +1600,7 @@ def smart_chat_stream(request: ChatRequest):
             elif intent_data["intent"] == "rejection":
                 session_manager.update_session(chat_id, {"state": "initial", "proposed_pipeline": None})
                 friendly = "👍 No problem! Pipeline cancelled. What else would you like me to help you with?"
-                session_manager.add_message(chat_id, "assistant", friendly)
+                _add_and_mirror_message(chat_id, "assistant", friendly)
                 yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                 return
 
@@ -1615,17 +1616,17 @@ def smart_chat_stream(request: ChatRequest):
                     session_manager.update_session(chat_id, {"proposed_pipeline": new_pipeline, "state": "pipeline_proposed"})
                     formatted = format_pipeline_for_display(new_pipeline)
                     friendly = formatted + f"\n\n```json\n{json.dumps(new_pipeline, indent=2)}\n```"
-                    session_manager.add_message(chat_id, "assistant", friendly)
+                    _add_and_mirror_message(chat_id, "assistant", friendly)
                     yield emit({"type": "assistant_final", "content": friendly, "pipeline": new_pipeline, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                     return
                 except Exception as e:
                     friendly = f"❌ Edit failed: {str(e)}"
-                    session_manager.add_message(chat_id, "assistant", friendly)
+                    _add_and_mirror_message(chat_id, "assistant", friendly)
                     yield emit({"type": "assistant_final", "content": friendly, "error": str(e), "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
                     return
 
         friendly = "Please upload a document and tell me what you'd like me to do (e.g., extract text, summarize, translate)."
-        session_manager.add_message(chat_id, "assistant", friendly)
+        _add_and_mirror_message(chat_id, "assistant", friendly)
         yield emit({"type": "assistant_final", "content": friendly, "history": [m.dict() for m in _normalize_history_for_api(chat_id)]})
 
     return StreamingResponse(gen(), media_type="application/x-ndjson")
@@ -1799,6 +1800,7 @@ def download_chat_file(chat_id: str):
             yield chunk
 
     media_type = obj.get("ContentType", "application/octet-stream")
+    # Note: os.path.basename used below requires os import (already imported at top)
     return StreamingResponse(stream(), media_type=media_type, headers={
         "Content-Disposition": f'attachment; filename="{os.path.basename(key)}"'
     })