refactor(lead-stage): clean 2-agent pipeline - no hardcoded heuristics

- Classifier (1 LLM call): classify stage via STAGE_JSON + call tools
- Stylist (1 LLM call): format response with optional stage context
- Removed _TOOL_STAGE_MAP hardcoded heuristic (stage is optional)
- Added re import for STAGE_JSON regex parsing
- Cleaned up debug logging

backend/agent/lead_stage_agent/controller.py

+"""
+Lead Stage Agent Controller — Entry point with Langfuse tracing.
+
+Wraps LeadStageGraph (Classifier ⇄ Tools → Stylist) with:
+- Langfuse trace + span context
+- Conversation history loading
+- User insight injection
+- Background task for history persistence
+
+Every call creates a named Langfuse trace "lead-stage-chat" so the full
+Classifier → Tool → Stylist pipeline is visible in the Langfuse dashboard.
+"""
+
+import json
+import logging
+import time
+import uuid
+
+from fastapi import BackgroundTasks
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.runnables import RunnableConfig
+from langfuse import Langfuse, get_client as get_langfuse
+
+from common.conversation_manager import get_conversation_manager
+from common.langfuse_client import async_flush_langfuse, get_callback_handler
+
+from agent.controller_helpers import load_user_insight_from_redis
+from agent.helper import handle_post_chat_async
+
+from .graph import get_lead_stage_agent
+
+logger = logging.getLogger(__name__)
+
+
+async def lead_stage_chat_controller(
+    *,
+    query: str,
+    identity_key: str,
+    background_tasks: BackgroundTasks,
+    model_name: str | None = None,
+    conversation_id: str | None = None,
+    is_authenticated: bool = False,
+    device_id: str | None = None,
+) -> dict:
+    """
+    Controller cho Lead Stage Agent — full Langfuse tracing.
+
+    Flow:
+    1. Load user_insight từ Redis
+    2. Load chat history từ ConversationManager
+    3. Build Langfuse trace context
+    4. Run LeadStageGraph (Classifier ⇄ Tools → Stylist)
+    5. Save conversation background
+    6. Return response + pipeline + lead_stage metadata
+
+    Returns:
+        dict with: status, ai_response, products, lead_stage, pipeline, timing, trace_id
+    """
+    start_time = time.time()
+    run_id = str(uuid.uuid4())
+    session_id = conversation_id or f"{identity_key}-lead-{run_id[:8]}"
+
+    logger.info(f"📥 [Lead Controller] identity={identity_key} | query={query[:80]}")
+
+    # ═══ 1. LOAD USER INSIGHT ═══
+    user_insight = None
+    if identity_key:
+        user_insight = await load_user_insight_from_redis(str(identity_key))
+
+    # ═══ 2. LOAD CHAT HISTORY ═══
+    memory = await get_conversation_manager()
+    history_dicts = await memory.get_chat_history(
+        str(identity_key),
+        limit=10,
+        include_product_ids=False,
+        conversation_id=conversation_id,
+    )
+    history_msgs = [
+        HumanMessage(content=m["message"]) if m["is_human"] else AIMessage(content=m["message"])
+        for m in history_dicts
+    ][::-1]
+
+    # ═══ 3. BUILD LANGFUSE TRACE ═══
+    trace_id = Langfuse.create_trace_id()
+    tags = ["lead_stage", "experiment"]
+    if is_authenticated:
+        tags.append("user:authenticated")
+    else:
+        tags.append("user:anonymous")
+
+    langfuse = get_langfuse()
+    observation_ctx = None
+
+    if langfuse:
+        try:
+            observation_ctx = langfuse.start_as_current_observation(
+                as_type="span",
+                name="lead-stage-chat",
+                trace_context={"trace_id": trace_id},
+            )
+        except Exception as e:
+            logger.warning(f"⚠️ Langfuse span init failed: {e}")
+
+    # Enter observation context
+    span = None
+    if observation_ctx:
+        try:
+            span = observation_ctx.__enter__()
+            span.update_trace(
+                name="lead-stage-chat",
+                user_id=str(identity_key),
+                session_id=session_id,
+                tags=tags,
+                input={"query": query, "identity_key": identity_key},
+                metadata={
+                    "device_id": device_id,
+                    "customer_id": identity_key if is_authenticated else None,
+                    "model": model_name,
+                    "conversation_id": conversation_id,
+                    "history_turns": len(history_msgs),
+                    "has_user_insight": user_insight is not None,
+                },
+            )
+        except Exception as e:
+            logger.warning(f"⚠️ Langfuse trace update failed: {e}")
+
+    # Create CallbackHandler INSIDE observation context for proper nesting
+    langfuse_handler = get_callback_handler()
+
+    exec_config = RunnableConfig(
+        callbacks=[langfuse_handler] if langfuse_handler else [],
+        run_name="lead-stage-graph",
+        run_id=run_id,
+        metadata={
+            "langfuse_session_id": session_id,
+            "langfuse_user_id": str(identity_key),
+            "langfuse_tags": tags,
+            "trace_id": trace_id,
+            "conversation_id": session_id,
+            "customer_id": identity_key if is_authenticated else None,
+            "device_id": device_id,
+        },
+    )
+
+    # ═══ 4. RUN LEAD STAGE GRAPH ═══
+    try:
+        agent = get_lead_stage_agent(model_name)
+        chat_result = await agent.chat(
+            user_message=query,
+            user_insight=user_insight,
+            history=history_msgs,
+            config=exec_config,
+        )
+
+        ai_response = chat_result.get("response", "")
+        lead_stage = chat_result.get("lead_stage") or {}
+        pipeline = chat_result.get("pipeline", [])
+        products = chat_result.get("products", [])
+        total_elapsed = chat_result.get("elapsed_ms", 0)
+
+        # Timing breakdown from pipeline
+        classifier_ms = 0
+        stylist_ms = 0
+        tool_ms = 0
+        for p in pipeline:
+            step = p.get("step", "")
+            ms = p.get("elapsed_ms", 0)
+            if "classifier" in step:
+                classifier_ms += ms
+            elif step in ("stylist", "responder"):
+                stylist_ms += ms
+            elif step == "tool_result":
+                tool_ms += ms
+
+        # Update Langfuse trace output
+        if span:
+            try:
+                span.update_trace(
+                    output={
+                        "ai_response": (ai_response or "")[:500],
+                        "stage": lead_stage.get("stage_name", "UNKNOWN"),
+                        "stage_confidence": lead_stage.get("confidence"),
+                        "product_count": len(products),
+                        "elapsed_ms": round(total_elapsed),
+                        "classifier_ms": classifier_ms,
+                        "stylist_ms": stylist_ms,
+                    },
+                )
+            except Exception:
+                pass
+
+        logger.info(
+            f"✅ [Lead Controller] Stage: {lead_stage.get('stage_name', 'N/A')} | "
+            f"Products: {len(products)} | "
+            f"Time: {total_elapsed:.0f}ms (C:{classifier_ms}ms S:{stylist_ms}ms)"
+        )
+
+        # ═══ 5. SAVE CONVERSATION (background) ═══
+        response_payload = {
+            "ai_response": ai_response,
+            "product_ids": products,
+            "lead_stage": lead_stage,
+        }
+
+        background_tasks.add_task(
+            handle_post_chat_async,
+            memory=memory,
+            identity_key=str(identity_key),
+            human_query=query,
+            ai_response=response_payload,
+            conversation_id=conversation_id,
+        )
+
+        # ═══ 6. RETURN ═══
+        return {
+            "status": "success",
+            "ai_response": ai_response,
+            "products": products,
+            "trace_id": trace_id,
+            "lead_stage": lead_stage,
+            "pipeline": pipeline,
+            "timing": {
+                "classifier_ms": classifier_ms,
+                "stylist_ms": stylist_ms,
+                "tool_ms": tool_ms,
+                "total_ms": round(total_elapsed),
+            },
+        }
+
+    except Exception as e:
+        elapsed = round((time.time() - start_time) * 1000)
+        logger.error(f"❌ [Lead Controller] Error after {elapsed}ms: {e}", exc_info=True)
+
+        # Log error to Langfuse
+        if span:
+            try:
+                span.update_trace(
+                    output={"error": str(e)[:500], "elapsed_ms": elapsed},
+                    level="ERROR",
+                )
+            except Exception:
+                pass
+
+        return {
+            "status": "error",
+            "error_code": "LEAD_STAGE_ERROR",
+            "message": f"Lead Stage Error: {str(e)[:200]}",
+            "trace_id": trace_id,
+        }
+
+    finally:
+        # Close Langfuse observation context
+        if observation_ctx:
+            try:
+                observation_ctx.__exit__(None, None, None)
+            except Exception:
+                pass
+
+        # Async flush Langfuse in background
+        background_tasks.add_task(async_flush_langfuse)

backend/agent/lead_stage_agent/graph.py

 """
 Lead Stage Graph — 2-Agent LangGraph Architecture cho phân tích khách hàng.
 
-Flow: User query → Classifier AI → Stylist AI ⇄ Tools → answer
+Flow: User → Classifier ⇄ Tools → Stylist → END
 
 2 con AI:
-  - Classifier (Planner): Phân tích ý định, đánh giá lead stage (Awareness, Interest, vv) → output JSON.
-  - Stylist (Responder): Nhận stage injection từ Classifier, sinh ra câu trả lời tư vấn phù hợp.
-    Có thể gọi tag_search_tool để tìm kiếm sản phẩm khi cần.
+  - Classifier (Planner + Tool Caller): Xác định lead stage, gọi tools khi cần,
+    có thể trả lời trực tiếp nếu câu hỏi đơn giản.
+  - Stylist (Responder): KHÔNG gọi tools. Chỉ nhận kết quả + format câu trả lời đẹp.
 
 Trang thí nghiệm — KHÔNG ảnh hưởng production flow.
 """
 
 import json
 import logging
+import re
 import time
 from typing import Annotated, Any, TypedDict
 
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
 from langchain_core.runnables import RunnableConfig
-from langgraph.graph import END, StateGraph, START
+from langgraph.graph import END, StateGraph
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode
 
 from common.llm_factory import create_llm
 from config import DEFAULT_MODEL
 
-from .prompts import LEAD_STAGE_CLASSIFIER_PROMPT, format_stage_injection
+from .prompts import LEAD_STAGE_CLASSIFIER_PROMPT, STYLIST_SYSTEM_PROMPT, format_stage_injection
 from .lead_search_tool import tag_search_tool
 
 # --- Nhập thêm các tools từ agent/tools ---
@@ -38,6 +39,9 @@ from agent.tools.add_to_cart_tool import add_to_cart
 
 logger = logging.getLogger(__name__)
 
+# Regex to extract STAGE_JSON: {...} from classifier response
+_STAGE_JSON_RE = re.compile(r'STAGE_JSON:\s*(\{.*\})', re.DOTALL)
+
 
 def _extract_text(content) -> str:
     """Extract plain text from msg.content — handles both str and Gemini list format."""
@@ -53,30 +57,70 @@ def _extract_text(content) -> str:
         return "\n".join(parts) if parts else str(content)
     return str(content or "")
 
+
+def _parse_stage_from_text(text: str) -> dict | None:
+    """Try to parse lead stage JSON from classifier response text.
+    
+    Supports 2 formats:
+    1. STAGE_JSON: {...}  (inline marker)
+    2. Pure JSON response  (entire response is JSON)
+    """
+    if not text:
+        return None
+    
+    # Format 1: STAGE_JSON: {...}
+    m = _STAGE_JSON_RE.search(text)
+    if m:
+        try:
+            result = json.loads(m.group(1))
+            if "stage" in result and "stage_name" in result:
+                return result
+        except (json.JSONDecodeError, TypeError):
+            pass
+    
+    # Format 2: Pure JSON
+    stripped = text.strip()
+    if stripped.startswith("{"):
+        # Strip markdown code fences if present
+        if stripped.startswith("```"):
+            stripped = re.sub(r'^```\w*\n?', '', stripped)
+            stripped = re.sub(r'\n?```$', '', stripped).strip()
+        try:
+            result = json.loads(stripped)
+            if "stage" in result and "stage_name" in result:
+                return result
+        except (json.JSONDecodeError, TypeError):
+            pass
+    
+    return None
+
+
+
 # ═══════════════════════════════════════════════
 # State
 # ═══════════════════════════════════════════════
 class LeadStageState(TypedDict):
     messages: Annotated[list[BaseMessage], add_messages]
-    
+
     # Context data
     user_insight: str | None
     chat_history_summary: str | None
-    
+
     # Internal routing & stage data
     lead_stage: dict | None
     stage_injection: str | None
+    classifier_responded: bool  # True nếu Classifier trả lời trực tiếp (skip Stylist)
     diagnostics: Annotated[list[dict[str, Any]], lambda x, y: x + y]
 
 
 # ═══════════════════════════════════════════════
-# Graph Builder — 2-Agent Architecture + Tools
+# Graph Builder — 2-Agent Architecture
+# Classifier (1 call: classify + tools) → Stylist (1 call: format)
 # ═══════════════════════════════════════════════
 class LeadStageGraph:
-    """2-Agent LangGraph: Classifier → Stylist ⇄ Tools."""
+    """2-Agent LangGraph: Classifier ⇄ Tools → Stylist → END."""
 
     def __init__(self, model_name: str | None = None):
-        # Dùng model_name nếu có, không thì mặc định theo DEFAULT_MODEL (từ config)
         self.model_name = model_name or DEFAULT_MODEL
         self.tools = [
             tag_search_tool,
@@ -85,225 +129,255 @@ class LeadStageGraph:
             canifa_store_search,
             check_is_stock,
             collect_customer_info,
-            add_to_cart
+            add_to_cart,
         ]
 
-        # AI 1: Classifier (buộc dùng JSON mode để lấy state info)
+        # AI 1: Classifier + Tool Caller (streaming off, has tools)
         self.classifier_llm = create_llm(
-            model_name=self.model_name, 
-            streaming=False, 
-            json_mode=True
+            model_name=self.model_name,
+            streaming=False,
         )
-        
-        # AI 2: Stylist (streaming, có tool binding để tìm sản phẩm)
+        self.classifier_with_tools = self.classifier_llm.bind_tools(self.tools)
+
+        # AI 2: Stylist — response only (streaming, NO tools)
         self.stylist_llm = create_llm(
-            model_name=self.model_name, 
-            streaming=True
+            model_name=self.model_name,
+            streaming=True,
         )
-        self.stylist_with_tools = self.stylist_llm.bind_tools(self.tools)
 
         self._compiled = None
-        logger.info(f"✅ LeadStageGraph 2-Agent + Tools initialized with model: {self.model_name}")
+        logger.info(f"✅ LeadStageGraph initialized | model: {self.model_name}")
 
     # ─────────────────────────────────────
-    # Node 1: CLASSIFIER — Phân tích stage
+    # Node 1: CLASSIFIER — 1 LLM call: classify stage + call tools
     # ─────────────────────────────────────
     async def _classifier_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
-        """Classifier AI: phân tích câu hỏi, trả về JSON đại diện Lead Stage."""
+        """
+        Classifier AI — 1 LLM call duy nhất:
+        - Xác định lead stage (qua STAGE_JSON marker)
+        - Gọi tools nếu cần (search SP, check stock...)
+        - Hoặc trả lời trực tiếp nếu câu hỏi đơn giản
+
+        Lần 2+ (sau tools return): pass-through → route thẳng tới Stylist
+        """
         messages = state["messages"]
         user_query = _extract_text(messages[-1].content) if messages else ""
         user_insight = state.get("user_insight") or "Chưa có thông tin."
         history_summary = state.get("chat_history_summary") or "Chưa có lịch sử."
 
-        classifier_input = (
+        has_tool_results = any(isinstance(m, ToolMessage) for m in messages)
+
+        # ── PASS 2+: Tool results back → skip LLM, go to Stylist ──
+        if has_tool_results:
+            logger.info("🔄 Classifier pass-through → Stylist")
+            return {
+                "lead_stage": state.get("lead_stage"),
+                "stage_injection": state.get("stage_injection"),
+                "classifier_responded": False,
+                "diagnostics": [],
+            }
+
+        # ── PASS 1: Single LLM call — classify + tools ──
+        context_block = (
             f"=== TIN NHẮN HIỆN TẠI ===\n{user_query}\n\n"
             f"=== USER INSIGHT (tích lũy từ trước) ===\n{user_insight}\n\n"
             f"=== TÓM TẮT LỊCH SỬ CHAT ===\n{history_summary}"
         )
 
-        sys_msg = SystemMessage(content=LEAD_STAGE_CLASSIFIER_PROMPT)
-        human_msg = HumanMessage(content=classifier_input)
+        classifier_messages = [
+            SystemMessage(content=LEAD_STAGE_CLASSIFIER_PROMPT),
+        ]
+        for msg in messages:
+            if not isinstance(msg, SystemMessage):
+                classifier_messages.append(msg)
+        classifier_messages.append(HumanMessage(content=context_block))
 
         start = time.time()
-        try:
-            response = await self.classifier_llm.ainvoke([sys_msg, human_msg], config=config)
-            content = _extract_text(response.content)
-            
-            # Parse JSON
-            stage_result = json.loads(content)
-            elapsed_ms = (time.time() - start) * 1000
+        response = await self.classifier_with_tools.ainvoke(classifier_messages, config=config)
+        elapsed_ms = (time.time() - start) * 1000
+
+        response_text = _extract_text(response.content)
+        has_tool_calls = hasattr(response, "tool_calls") and response.tool_calls
 
+        # ── Parse lead stage from response text (STAGE_JSON: {...}) ──
+        lead_stage = None
+        stage_injection = None
+        stage_result = _parse_stage_from_text(response_text)
+        if stage_result:
+            lead_stage = stage_result
+            stage_injection = format_stage_injection(stage_result)
             logger.info(
-                f"🎯 Lead Stage: {stage_result.get('stage_name')} "
-                f"(stage={stage_result.get('stage')}, conf={stage_result.get('confidence')}) "
-                f"| {elapsed_ms:.0f}ms"
+                f"🎯 Stage: {stage_result.get('stage_name')} "
+                f"(conf={stage_result.get('confidence')}) | {elapsed_ms:.0f}ms"
             )
 
-            injection = format_stage_injection(stage_result)
-            
-            diag = {
-                "step": "classifier",
-                "label": "🧠 Classifier AI",
-                "content": f"Stage {stage_result.get('stage')}: {stage_result.get('stage_name')} (Conf: {stage_result.get('confidence')})\nLý do: {stage_result.get('reasoning')}",
-                "raw_json": json.dumps(stage_result, ensure_ascii=False, indent=2),
-                "elapsed_ms": round(elapsed_ms)
-            }
 
-            return {
-                "lead_stage": stage_result,
-                "stage_injection": injection,
-                "diagnostics": [diag]
-            }
+        # ── Build diagnostic ──
+        diag = {"elapsed_ms": round(elapsed_ms)}
 
-        except Exception as e:
-            elapsed_ms = (time.time() - start) * 1000
-            logger.error(f"❌ Lead Stage Classifier failed: {e}")
-            
-            fallback = {
-                "stage": 2,
-                "stage_name": "INTEREST",
-                "confidence": 0.3,
-                "reasoning": f"Classifier error: {str(e)[:100]}. Fallback to INTEREST.",
-                "tone_directive": "Consultant",
-                "behavioral_hints": ["Hỏi thêm thông tin để hiểu nhu cầu khách."],
-            }
-            injection = format_stage_injection(fallback)
-
-            diag = {
-                "step": "classifier",
-                "label": "🧠 Classifier AI (Error Fallback)",
-                "content": f"Lỗi parse JSON: {e}",
-                "raw_json": json.dumps(fallback, ensure_ascii=False, indent=2),
-                "elapsed_ms": round(elapsed_ms)
-            }
+        if has_tool_calls:
+            diag["step"] = "classifier_tool_call"
+            diag["label"] = "🔧 Classifier (Stage + Tool Call)"
+            tool_names = ', '.join(tc['name'] for tc in response.tool_calls)
+            stage_info = f"Stage {lead_stage['stage']}: {lead_stage['stage_name']}" if lead_stage else "Stage: unknown"
+            diag["content"] = f"{stage_info} | Tools: {tool_names}"
+            diag["tool_calls"] = [{"name": tc["name"], "args": tc["args"]} for tc in response.tool_calls]
+            if lead_stage:
+                diag["raw_json"] = json.dumps(lead_stage, ensure_ascii=False, indent=2)
+        elif lead_stage:
+            diag["step"] = "classifier_stage"
+            diag["label"] = "🧠 Classifier (Stage Only)"
+            diag["content"] = (
+                f"Stage {lead_stage.get('stage')}: {lead_stage.get('stage_name')} "
+                f"(Conf: {lead_stage.get('confidence')})\n"
+                f"Lý do: {lead_stage.get('reasoning', '')}"
+            )
+            diag["raw_json"] = json.dumps(lead_stage, ensure_ascii=False, indent=2)
+        else:
+            diag["step"] = "classifier_direct"
+            diag["label"] = "💬 Classifier (Direct Response)"
+            diag["content"] = response_text[:500]
 
-            return {
-                "lead_stage": fallback,
-                "stage_injection": injection,
-                "diagnostics": [diag]
-            }
+        # classifier_responded = True → skip Stylist (only for simple greetings)
+        is_direct = not has_tool_calls and not lead_stage
+
+        return {
+            "messages": [response],
+            "lead_stage": lead_stage,
+            "stage_injection": stage_injection,
+            "classifier_responded": is_direct,
+            "diagnostics": [diag],
+        }
 
     # ─────────────────────────────────────
-    # Node 2: STYLIST — Format câu trả lời (có tools)
+    # Router: Classifier → tools | stylist | end
+    # ─────────────────────────────────────
+    def _after_classifier(self, state: LeadStageState) -> str:
+        """
+        Sau Classifier:
+        - Có tool_calls → đi tools
+        - Classifier trả lời trực tiếp (greeting, no stage) → end
+        - Mọi trường hợp khác (có lead_stage, hoặc sau tools) → stylist
+        """
+        last = state["messages"][-1]
+        if hasattr(last, "tool_calls") and last.tool_calls:
+            return "tools"
+        if state.get("classifier_responded"):
+            return "end"
+        return "stylist"
+
+    # ─────────────────────────────────────
+    # Node 2: STYLIST — Format response (NO tools)
     # ─────────────────────────────────────
     async def _stylist_node(self, state: LeadStageState, config: RunnableConfig | None = None) -> dict:
-        """Stylist AI: nhận stage injection, đưa ra câu trả lời tư vấn. Có thể gọi tool."""
+        """
+        Stylist AI: Nhận lead stage + tool results, format câu trả lời tư vấn đẹp.
+        KHÔNG gọi tools — chỉ viết response.
+        """
         messages = state["messages"]
-        injection = state.get("stage_injection", "")
-
-        # Stylist System Prompt with FULL tool awareness
-        stylist_sys_prompt = f"""\
-Bạn là Trợ lý phong cách (AI Stylist) độc quyền của CANIFA — thương hiệu thời trang hàng đầu Việt Nam.
-Nhiệm vụ của bạn không chỉ là tìm đồ mà còn là KHƠI GỢI NHU CẦU và TƯ VẤN THỜI TRANG khéo léo.
-Quy tắc Tư Vấn "Khéo Léo":
-- Luôn khen ngợi hoặc đồng tình với gu thẩm mỹ/mong muốn của khách.
-- Khi giới thiệu sản phẩm: Gợi ý cách phối đồ (mix & match), mô tả cảm giác khi mặc (chất vải thoáng mát, form dáng tôn dáng...).
-- Không bao giờ trả lời cộc lốc một danh sách sản phẩm. Hãy nói như một người tư vấn tại shop đang trò chuyện cùng bạn của mình.
-
-=== HƯỚNG DẪN DÙNG CÔNG CỤ (TOOLS) BẮT BUỘC ===
-Bạn có các công cụ để truy xuất thông tin thực tế. Hãy chọn công cụ phù hợp với câu hỏi của khách:
-
-1. 🛍️ TÌM KIẾM SẢN PHẨM: Khi khách hỏi quần áo, váy vóc, đồ đi làm, mặc nhà...
-   → Dùng tool `tag_search_tool`. Đợi kết quả công cụ trả về mới bắt đầu viết câu tư vấn sản phẩm khéo léo.
-   → KHÔNG BAO GIỜ tự bịa ra sản phẩm.
+        injection = state.get("stage_injection") or ""
 
-2. 📦 KIỂM TRA TỒN KHO: Khi khách hỏi size M mẫu này còn không, sản phẩm SKU 12345 còn hàng ở đâu...
-   → Dùng tool `check_is_stock`. Đầu vào: mã SKU, color (mã), size.
+        # Build system prompt with stage context
+        stylist_sys = STYLIST_SYSTEM_PROMPT.format(stage_injection=injection)
 
-3. 🏬 TÌM CỬA HÀNG: Khi khách hỏi Cửa hàng ở Hà Đông, Ở Đà Nẵng có shop nào...
-   → Dùng tool `canifa_store_search`. Đầu vào: tỉnh/thành phố hoặc khu vực.
-
-4. ❓ CSKH & CHÍNH SÁCH: Khi khách hỏi phí ship, đổi trả đồ, đăng ký thẻ VIP, điều kiện lên hạng...
-   → Dùng tool `canifa_knowledge_search`. Câu trả lời sẽ có trong tài liệu quy định nội bộ của Canifa.
-
-5. 🎁 KHUYẾN MÃI: Khi khách hỏi nay có sale gì, chương trình sinh nhật...
-   → Dùng tool `canifa_get_promotions`. Trả về tin tức event/sale hiện tại.
-
-6. 📝 LẤY THÔNG TIN LEAD: Khi khách muốn đăng ký tư vấn, nhận ưu đãi đặc biệt qua SDT/Email, hoặc gửi khiếu nại cần gọi lại...
-   → Dùng tool `collect_customer_info`. Yêu cầu khách cung cấp tên và số điện thoại trước khi gọi if thiếu.
-
-7. 🛒 GIỎ HÀNG (ADD TO CART): Khi khách đòi mua, chốt đơn mua một sản phẩm cụ thể (đã biết mã, màu, size)...
-   → Dùng tool `add_to_cart`. Đảm bảo phải có mã sản phẩm, size, màu. Yêu cầu khách chọn size/màu nếu chưa ghi rõ.
+        # Collect all tool results to give stylist context
+        tool_context_parts = []
+        for msg in messages:
+            if isinstance(msg, ToolMessage):
+                tool_name = getattr(msg, "name", "tool")
+                tool_content = _extract_text(msg.content)
+                # Parse for cleaner summary
+                try:
+                    parsed = json.loads(tool_content)
+                    if parsed.get("status") == "success":
+                        products = parsed.get("products", [])
+                        if products:
+                            product_lines = []
+                            for p in products[:8]:
+                                line = f"- {p.get('name', '')} | {p.get('price', '')} | SKU: {p.get('sku_code', '')}"
+                                if p.get("url"):
+                                    line += f" | Link: {p['url']}"
+                                product_lines.append(line)
+                            tool_context_parts.append(
+                                f"[{tool_name}] Tìm thấy {len(products)} sản phẩm:\n" + "\n".join(product_lines)
+                            )
+                        else:
+                            tool_context_parts.append(f"[{tool_name}] {json.dumps(parsed, ensure_ascii=False)[:500]}")
+                    else:
+                        tool_context_parts.append(f"[{tool_name}] {parsed.get('message', tool_content[:300])}")
+                except Exception:
+                    tool_context_parts.append(f"[{tool_name}] {tool_content[:500]}")
 
-Nếu khách khen/chê hoặc chỉ chào hỏi bình thường → Trả lời trực tiếp bằng phong thái khéo léo, bỏ qua việc gọi tool nếu không cần tra cứu.
+        # Build stylist messages
+        stylist_messages = [SystemMessage(content=stylist_sys)]
 
-{injection}
-"""
-        
-        # Inject responder system prompt
-        stylist_messages = [SystemMessage(content=stylist_sys_prompt)]
+        # Add user's original question
         for msg in messages:
-            if not isinstance(msg, SystemMessage):
+            if isinstance(msg, HumanMessage):
                 stylist_messages.append(msg)
+                break  # Only first human message (the actual query)
+
+        # Add tool results as context
+        if tool_context_parts:
+            tool_summary = "=== KẾT QUẢ TRA CỨU TỪ TOOLS ===\n" + "\n\n".join(tool_context_parts)
+            stylist_messages.append(HumanMessage(content=tool_summary))
 
         start = time.time()
-        response = await self.stylist_with_tools.ainvoke(stylist_messages, config=config)
+        response = await self.stylist_llm.ainvoke(stylist_messages, config=config)
         elapsed_ms = (time.time() - start) * 1000
 
         response_text = _extract_text(response.content)
-        has_tool_calls = hasattr(response, "tool_calls") and response.tool_calls
 
         diag = {
-            "step": "stylist" if not has_tool_calls else "stylist_tool_call",
-            "label": "💬 Stylist AI" if not has_tool_calls else "🔧 Stylist AI → Tool Call",
-            "elapsed_ms": round(elapsed_ms)
+            "step": "stylist",
+            "label": "💬 Stylist AI (Response)",
+            "content": response_text[:500] + ("..." if len(response_text) > 500 else ""),
+            "elapsed_ms": round(elapsed_ms),
         }
 
-        if has_tool_calls:
-            diag["content"] = f"Gọi {len(response.tool_calls)} tool(s): {', '.join(tc['name'] for tc in response.tool_calls)}"
-            diag["tool_calls"] = [
-                {"name": tc["name"], "args": tc["args"]}
-                for tc in response.tool_calls
-            ]
-        else:
-            diag["content"] = response_text[:500] + ("..." if len(response_text) > 500 else "")
-
-        # Trả về message để add vào list, và diagnostics để track timeline
         return {
             "messages": [response],
-            "diagnostics": [diag]
+            "diagnostics": [diag],
         }
 
     # ─────────────────────────────────────
-    # Router: Stylist → tool hoặc end
+    # Build Graph
     # ─────────────────────────────────────
-    def _after_stylist(self, state: LeadStageState) -> str:
-        """Sau Stylist: nếu có tool_calls → đi tools, nếu không → end."""
-        last = state["messages"][-1]
-        if hasattr(last, "tool_calls") and last.tool_calls:
-            return "tools"
-        return "end"
-
     def build(self):
-        """Build 2-agent + tools graph."""
+        """Build 2-agent graph: Classifier ⇄ Tools → Stylist → END."""
         if self._compiled:
             return self._compiled
 
         workflow = StateGraph(LeadStageState)
 
-        # 3 Nodes: Classifier, Stylist, Tools
+        # 3 Nodes
         workflow.add_node("classifier", self._classifier_node)
-        workflow.add_node("stylist", self._stylist_node)
         workflow.add_node("tools", ToolNode(self.tools))
+        workflow.add_node("stylist", self._stylist_node)
 
-        # Flow: Entry → Classifier → Stylist
+        # Flow: Entry → Classifier
         workflow.set_entry_point("classifier")
-        workflow.add_edge("classifier", "stylist")
 
-        # Stylist → (tools | end)
+        # Classifier → (tools | stylist | end)
         workflow.add_conditional_edges(
-            "stylist",
-            self._after_stylist,
-            {"tools": "tools", "end": END},
+            "classifier",
+            self._after_classifier,
+            {"tools": "tools", "stylist": "stylist", "end": END},
         )
 
-        # Tools → Stylist (loop lại để Stylist format kết quả)
-        workflow.add_edge("tools", "stylist")
+        # Tools → Classifier (loop back so classifier sees tool results)
+        workflow.add_edge("tools", "classifier")
+
+        # Stylist → END
+        workflow.add_edge("stylist", END)
 
         self._compiled = workflow.compile()
-        logger.info("✅ LeadStageGraph 2-Agent + Tools compiled")
+        logger.info("✅ LeadStageGraph compiled: Classifier ⇄ Tools → Stylist → END")
         return self._compiled
 
+    # ─────────────────────────────────────
+    # Main Entry Point
+    # ─────────────────────────────────────
     async def chat(self, user_message: str, user_insight: str | None = None, history: list[BaseMessage] | None = None, config: RunnableConfig | None = None) -> dict:
         """
         Main entry point — gửi message và nhận response.
@@ -323,7 +397,8 @@ Nếu khách khen/chê hoặc chỉ chào hỏi bình thường → Trả lời
             "chat_history_summary": None,
             "lead_stage": None,
             "stage_injection": None,
-            "diagnostics": []
+            "classifier_responded": False,
+            "diagnostics": [],
         }
 
         result = await graph.ainvoke(initial_state, config=config)
@@ -333,26 +408,20 @@ Nếu khách khen/chê hoặc chỉ chào hỏi bình thường → Trả lời
 
         # Extract AI response (last AI message without tool_calls)
         ai_response = ""
-        if all_messages:
-            for msg in reversed(all_messages):
-                if isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
-                    ai_response = _extract_text(msg.content)
-                    break
-
-        # Build pipeline trace from diagnostics + messages
-        pipeline = []
-        pipeline.append({
-            "step": "user",
-            "label": "👤 User",
-            "content": user_message,
-        })
-        
-        # Append diagnostic tracking (classifier -> stylist -> tool -> stylist)
+        for msg in reversed(all_messages):
+            if isinstance(msg, AIMessage) and not (hasattr(msg, "tool_calls") and msg.tool_calls):
+                ai_response = _extract_text(msg.content)
+                break
+
+
+        # Build pipeline trace
+        pipeline = [{"step": "user", "label": "👤 User", "content": user_message}]
+
         if result.get("diagnostics"):
             for diag in result["diagnostics"]:
                 pipeline.append(diag)
 
-        # Parse tool results from messages for product cards + pipeline enrichment
+        # Parse tool results for product cards
         all_products = []
         for msg in all_messages:
             if isinstance(msg, ToolMessage):
@@ -386,13 +455,12 @@ Nếu khách khen/chê hoặc chỉ chào hỏi bình thường → Trả lời
                     "raw_json": raw_json,
                 })
 
-        # Add final responder step if there was a tool cycle
-        # (diagnostics already captures the final stylist response)
-
         logger.info(
             "🏷️ [LEAD GRAPH] query='%s' | stage=%s | products=%d | time=%.2fms",
-            user_message[:50], result.get("lead_stage", {}).get("stage_name", "UNKNOWN"),
-            len(all_products), elapsed_ms
+            user_message[:50],
+            result.get("lead_stage", {}).get("stage_name", "UNKNOWN") if result.get("lead_stage") else "DIRECT",
+            len(all_products),
+            elapsed_ms,
         )
 
         return {
@@ -410,6 +478,7 @@ Nếu khách khen/chê hoặc chỉ chào hỏi bình thường → Trả lời
 
 _instance: LeadStageGraph | None = None
 
+
 def get_lead_stage_agent(model_name: str | None = None) -> LeadStageGraph:
     """Get or create LeadStageGraph singleton."""
     global _instance
@@ -430,18 +499,18 @@ async def run_lead_stage_classifier(
     Được dùng bởi controller chính khi Controller đóng vai trò Stylist (chứa tools, logic cứng).
     """
     agent = get_lead_stage_agent()
-    # Mock state để chạy method _classifier_node độc lập
     state = {
         "messages": [HumanMessage(content=user_query)],
         "user_insight": user_insight,
         "chat_history_summary": chat_history_summary,
         "lead_stage": None,
         "stage_injection": None,
-        "diagnostics": []
+        "classifier_responded": False,
+        "diagnostics": [],
     }
     res = await agent._classifier_node(state)
     return {
         "lead_stage": res.get("lead_stage"),
         "classifier_latency_ms": res.get("diagnostics", [{}])[0].get("elapsed_ms", 0),
-        "stage_injection": res.get("stage_injection")
+        "stage_injection": res.get("stage_injection"),
     }

backend/agent/lead_stage_agent/prompts.py

 """
-Prompts cho Lead Stage Agent.
+Prompts cho Lead Stage Agent — 2-Agent Architecture.
 
-AI #1 — Lightweight classifier:
-  Input: user_query + user_insight + history summary
-  Output: JSON { stage, stage_name, confidence, reasoning, tone_directive, behavioral_hints }
+CHỈ 2 PROMPTS:
+  AI #1 — Classifier: Xác định lead stage + gọi tools (1 LLM call duy nhất)
+  AI #2 — Stylist: Nhận kết quả → format câu trả lời đẹp (1 LLM call)
 """
 
+# ═══════════════════════════════════════════════
+# AI #1: Classifier — Classify stage + Call tools (1 call)
+# ═══════════════════════════════════════════════
 LEAD_STAGE_CLASSIFIER_PROMPT = """\
-Bạn là Lead Stage Classifier cho chatbot bán hàng CANIFA.
-Nhiệm vụ DUY NHẤT: Phân tích tin nhắn khách hàng và xác định họ đang ở GIAI ĐOẠN MUA HÀNG nào.
+Bạn là AI Classifier & Planner cho chatbot bán hàng CANIFA — thương hiệu thời trang hàng đầu Việt Nam.
 
-## 5 GIAI ĐOẠN MUA HÀNG:
+## 2 NHIỆM VỤ TRONG 1 LẦN TRẢ LỜI:
+
+### NHIỆM VỤ 1: XÁC ĐỊNH LEAD STAGE
+Phân tích tin nhắn và xác định khách đang ở giai đoạn nào:
 
 | Stage | Tên            | Trigger Signals                                                          |
 |-------|----------------|--------------------------------------------------------------------------|
@@ -20,32 +25,56 @@ Nhiệm vụ DUY NHẤT: Phân tích tin nhắn khách hàng và xác định h
 | 4     | DECISION       | "Lấy cái này", "chốt đơn", "mua luôn", "order", hỏi cách thanh toán   |
 | 5     | RETENTION      | Hỏi đổi trả, tracking, mua thêm, quay lại sau thời gian                |
 
-## CÁCH PHÂN TÍCH:
-1. Đọc tin nhắn HIỆN TẠI (quan trọng nhất)
-2. Đọc USER_INSIGHT (nếu có) để hiểu context tích lũy
-3. Đọc CHAT_HISTORY_SUMMARY (nếu có) để thấy hành trình
-4. Quyết định stage dựa trên tổng hợp 3 nguồn trên
-
-## QUY TẮC:
-- Stage CÓ THỂ GIẢM (VD: khách đã xem SP nhưng quay lại hỏi "còn gì khác?" → stage 2)
-- Stage CÓ THỂ NHẢY (VD: khách vào thẳng "mua cái áo polo size L" → stage 4)
-- Khi KHÔNG CHẮC → chọn stage THẤP hơn (ít aggressive hơn)
-- confidence < 0.5 → đặt stage = stage hiện tại hoặc thấp hơn
-
-## OUTPUT — JSON DUY NHẤT, KHÔNG có text thêm:
-{
-  "stage": <1-5>,
-  "stage_name": "<AWARENESS|INTEREST|CONSIDERATION|DECISION|RETENTION>",
-  "confidence": <0.0-1.0>,
-  "reasoning": "<Giải thích NGẮN GỌN tại sao chọn stage này, 1-2 câu>",
-  "tone_directive": "<Welcomer|Consultant|Expert Stylist|Closer|Personal Shopper>",
-  "behavioral_hints": [
-    "<Gợi ý hành vi cụ thể cho AI Stylist, tối đa 3 hints>"
-  ]
-}
+Quy tắc: Stage có thể nhảy/giảm. Không chắc → chọn thấp hơn.
+
+### NHIỆM VỤ 2: GỌI TOOLS HOẶC TRẢ LỜI
+- Cần tra cứu → gọi tool phù hợp
+- Câu hỏi đơn giản (chào hỏi, cảm ơn) → trả lời ngắn gọn luôn
+
+Tools:
+- 🛍️ `tag_search_tool` — Tìm quần áo, váy, đồ đi làm, mặc nhà...
+- 📦 `check_is_stock` — Check size/màu còn hàng (cần SKU, color, size)
+- 🏬 `canifa_store_search` — Tìm cửa hàng (tỉnh/thành phố)
+- ❓ `canifa_knowledge_search` — Chính sách (ship, đổi trả, VIP...)
+- 🎁 `canifa_get_promotions` — Khuyến mãi, sale, event
+- 📝 `collect_customer_info` — Đăng ký tư vấn, gửi SĐT/Email
+- 🛒 `add_to_cart` — Chốt mua SP cụ thể (đã biết mã, màu, size)
+
+## CÁCH OUTPUT:
+- Nếu CẦN gọi tool → gọi tool (hệ thống sẽ tự xử lý)
+- Nếu KHÔNG cần tool → trả lời text ngắn gọn cho khách
+- Dù gọi tool hay trả lời text, HÃY bao gồm stage JSON ở CUỐI response:
+
+STAGE_JSON: {"stage": <1-5>, "stage_name": "<tên>", "confidence": <0-1>, "reasoning": "<lý do>", "tone_directive": "<Friendly|Consultant|Expert|Closer|CareAgent>", "behavioral_hints": ["hint1"]}
+"""
+
+# ═══════════════════════════════════════════════
+# AI #2: Stylist — Format response (NO tools)
+# ═══════════════════════════════════════════════
+STYLIST_SYSTEM_PROMPT = """\
+Bạn là Trợ lý phong cách (AI Stylist) độc quyền của CANIFA — thương hiệu thời trang hàng đầu Việt Nam.
+
+## VAI TRÒ DUY NHẤT CỦA BẠN:
+Nhận kết quả từ AI Classifier (lead stage + tool results) và FORMAT thành câu trả lời tư vấn
+đẹp, khéo léo, phù hợp với giai đoạn mua hàng của khách.
+
+## BẠN KHÔNG ĐƯỢC:
+- Gọi thêm tools
+- Tự bịa sản phẩm
+- Thay đổi lead stage
+
+## QUY TẮC TƯ VẤN:
+- Luôn khen ngợi hoặc đồng tình với gu thẩm mỹ của khách.
+- Khi giới thiệu SP: Gợi ý cách phối đồ (mix & match), mô tả cảm giác khi mặc.
+- Không trả lời cộc lốc danh sách SP. Nói như người tư vấn tại shop đang trò chuyện vui vẻ.
+- Sản phẩm nào có link → bắt buộc đính kèm link để khách click xem.
+
+{stage_injection}
 """
 
-# Template inject vào system prompt của AI #2 (Stylist)
+# ═══════════════════════════════════════════════
+# Template inject stage context vào Stylist prompt
+# ═══════════════════════════════════════════════
 LEAD_STAGE_INJECTION_TEMPLATE = """\
 ══════ LEAD STAGE CONTEXT ══════
 🎯 Giai đoạn khách hàng: Stage {stage} — {stage_name}