# Runs Router - Handles agent run execution, listing, and management
import asyncio
import ast
import json
from pathlib import Path
from datetime import datetime
from fastapi import APIRouter, BackgroundTasks, HTTPException, Request, Body, Depends
from pydantic import BaseModel, Field
from typing import Any, Dict, Optional

from shared.state import (
    active_loops,
    get_multi_mcp,
    get_remme_store,
    get_remme_extractor,
    PROJECT_ROOT,
)
from core.graph_adapter import nx_to_reactflow
from core.supabase_auth import require_supabase_user
from core.supabase_logging import (
    build_idempotency_key,
    compute_payload_hash,
    log_clinical_result,
    log_inbound_request,
    update_request_status,
)
from core.event_bus import event_bus
from core.run_store import get_run_store
from core.run_executor import execute_resume, execute_run, is_celery_enabled
from remme.utils import get_embedding
from config.settings_loader import settings, get_run_poll_timeout
from integrations.contracts import CanonicalRunRequest
from integrations.tenancy import (
    can_route_to_growth,
    resolve_tenant_context,
    storage_namespace_for_tenant,
)

router = APIRouter(tags=["Runs"])
RUNS_INDEX_FILE = PROJECT_ROOT / "memory" / "runs.index.json"

class _LazyShared:
    def __init__(self, factory):
        self._factory = factory
        self._value = None

    def _get(self):
        if self._value is None:
            self._value = self._factory()
        return self._value

    def __getattr__(self, name):
        return getattr(self._get(), name)


# Keep router imports lightweight; heavy MCP/FAISS objects are created on first use.
multi_mcp = _LazyShared(get_multi_mcp)
remme_store = _LazyShared(get_remme_store)
remme_extractor = _LazyShared(get_remme_extractor)
run_store = get_run_store()


# === Pydantic Models ===

class RunRequest(BaseModel):
    query: str = Field(min_length=1)
    model: str = None  # Will use settings default if not provided
    contract_version: Optional[str] = "v1"
    integration_id: Optional[str] = None
    workflow_id: Optional[str] = None
    source_system: Optional[str] = "s18"
    tenant_id: Optional[str] = None
    tenant_tier: Optional[str] = None
    data_region: Optional[str] = None
    external_event_id: Optional[str] = None
    consent_ref: Optional[str] = None
    raw_payload: Optional[Dict[str, Any]] = None
    idempotency_key: Optional[str] = None
    
    def __init__(self, **data):
        super().__init__(**data)
        if self.model is None:
            self.model = settings.get("agent", {}).get("default_model", "gemini-2.5-flash")


class RunResponse(BaseModel):
    id: str
    status: str
    created_at: str
    query: str


class UserInputRequest(BaseModel):
    node_id: str
    response: str


# === Background Tasks ===

def _infer_query_type(query: str) -> str:
    q = (query or "").lower()
    if "cbc" in q:
        return "cbc"
    if "rac" in q:
        return "rac"
    if "abdm" in q or "fhir" in q:
        return "abdm_fhir"
    if "mental_health" in q or "phq9" in q or "gad7" in q:
        return "mental_health"
    return "generic"


def _infer_query_type_from_canonical(canonical_request: CanonicalRunRequest) -> str:
    workflow_id = (canonical_request.workflow_id or "").strip().lower()
    if workflow_id and workflow_id != "generic":
        return workflow_id
    return _infer_query_type(canonical_request.query)


def _session_file_candidates(run_id: str, summaries_dir: Path):
    """
    Build likely session file locations for a run id.
    Run ids are usually epoch timestamps (ms precision), so we can map to the
    date-sharded path directly instead of scanning the whole tree.
    """
    filename = f"session_{run_id}.json"
    candidates = []

    # Fast path for default run ids generated by this service.
    if run_id.isdigit():
        ts_value = int(run_id)
        # 13+ digits => milliseconds, otherwise seconds.
        ts_seconds = ts_value / 1000.0 if len(run_id) >= 13 else float(ts_value)
        try:
            dt = datetime.fromtimestamp(ts_seconds)
            candidates.append(
                summaries_dir
                / f"{dt.year}"
                / f"{dt.month:02d}"
                / f"{dt.day:02d}"
                / filename
            )
        except (OverflowError, OSError, ValueError):
            pass

    # Compatibility fallback for custom/non-timestamp run ids.
    for path in summaries_dir.glob(f"**/{filename}"):
        candidates.append(path)

    return candidates


def _find_session_file(run_id: str, summaries_dir: Path) -> Optional[Path]:
    for candidate in _session_file_candidates(run_id, summaries_dir):
        if candidate.exists():
            return candidate
    return None


def _read_run_entry_sync(session_file: Path) -> Optional[Dict[str, Any]]:
    try:
        data = json.loads(session_file.read_text(encoding="utf-8", errors="ignore"))
        graph_details = data.get("graph", {})
        query = graph_details.get("original_query") or graph_details.get("globals", {}).get("original_query", "Unknown Query")
        created_at = graph_details.get("created_at") or datetime.fromtimestamp(session_file.stat().st_ctime).isoformat()
        nodes = data.get("nodes", [])
        node_statuses = [n.get("status", "pending") for n in nodes if n.get("id") != "ROOT"]
        if any(s == "running" for s in node_statuses):
            computed_status = "running"
        elif any(s == "failed" for s in node_statuses):
            computed_status = "failed"
        elif all(s == "completed" for s in node_statuses) and node_statuses:
            computed_status = "completed"
        else:
            computed_status = graph_details.get("status", "completed")
        total_tokens = sum((n.get("total_tokens", 0) or 0) for n in nodes)
        return {
            "id": session_file.stem.replace("session_", ""),
            "query": query,
            "created_at": created_at,
            "status": computed_status,
            "total_tokens": total_tokens,
            "path": str(session_file),
        }
    except Exception:
        return None


def _write_runs_index_sync(index_rows: list[Dict[str, Any]]) -> None:
    RUNS_INDEX_FILE.parent.mkdir(parents=True, exist_ok=True)
    tmp = RUNS_INDEX_FILE.with_suffix(".json.tmp")
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(index_rows, fh, indent=2)
    tmp.replace(RUNS_INDEX_FILE)


async def _refresh_runs_index() -> list[Dict[str, Any]]:
    summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
    if not summaries_dir.exists():
        return []
    session_files = list(summaries_dir.glob("**/session_*.json"))
    rows = []
    for entry in await asyncio.gather(*[asyncio.to_thread(_read_run_entry_sync, path) for path in session_files]):
        if entry:
            rows.append(entry)
    rows.sort(key=lambda x: x["id"], reverse=True)
    await asyncio.to_thread(_write_runs_index_sync, rows)
    return rows


async def _load_runs_index() -> list[Dict[str, Any]]:
    if RUNS_INDEX_FILE.exists():
        try:
            data = await asyncio.to_thread(lambda: json.loads(RUNS_INDEX_FILE.read_text(encoding="utf-8")))
            if isinstance(data, list):
                return data
        except Exception:
            pass
    return await _refresh_runs_index()


def _normalize_run_status(raw_status: Optional[str]) -> str:
    status = (raw_status or "failed").strip().lower()
    if status == "success":
        return "completed"
    if status == "paused":
        return "stopped"
    if status not in {
        "accepted",
        "starting",
        "running",
        "waiting_input",
        "completed",
        "failed",
        "stopped",
        "interrupted",
    }:
        return "failed"
    return status


def _retrieve_memories_sync(query: str):
    return remme_store.search_text(query, limit=3, requester="run_context")


async def _build_memory_context(run_id: str, query: str):
    memory_context = ""
    results = []

    try:
        # Keep sync embedding/search work off the event loop so /runs polling
        # stays responsive under load.
        results = await asyncio.to_thread(_retrieve_memories_sync, query)
        if results:
            memory_str = "\n".join([f"- {r['text']} (Confidence: {r.get('score', 0):.2f})" for r in results])
            memory_context = f"PREVIOUS MEMORIES ABOUT USER:\n{memory_str}\n"
            print(f" Remme: Injected {len(results)} memories into run {run_id}")
    except Exception as e:
        print(f"⚠️ Remme Retrieval Failed: {e}")

    try:
        rag_result = await multi_mcp.call_tool("rag", "search_stored_documents_rag", {"query": query})
        if rag_result and hasattr(rag_result, "content") and isinstance(rag_result.content, list):
            snippets = []
            for item in rag_result.content:
                if hasattr(item, "text"):
                    try:
                        parsed = ast.literal_eval(item.text)
                        if isinstance(parsed, list):
                            snippets.extend(parsed[:5])
                        else:
                            snippets.append(item.text[:500])
                    except Exception:
                        snippets.append((item.text or "")[:500])
            if snippets:
                rag_str = "\n".join(f"- {s[:400]}..." if len(s) > 400 else f"- {s}" for s in snippets[:8])
                memory_context += f"\nRELEVANT DOCUMENT SNIPPETS (RAG):\n{rag_str}\n"
                print(f"[{run_id}] Injected {len(snippets)} RAG snippets into run context")
    except Exception as e:
        print(f"⚠️ RAG context injection failed (non-fatal): {e}")

    return memory_context, results


async def _mirror_run_status_from_events(run_id: str, stop_signal: asyncio.Event):
    queue = await event_bus.subscribe(max_queue_size=500)
    try:
        while not stop_signal.is_set():
            try:
                event = await asyncio.wait_for(queue.get(), timeout=0.5)
            except asyncio.TimeoutError:
                continue

            data = event.get("data", {})
            if data.get("session_id") != run_id:
                continue

            event_type = event.get("type")
            if event_type == "run_started":
                await asyncio.to_thread(run_store.update_status, run_id, "running")
            elif event_type == "step_update":
                status = data.get("status")
                if status == "waiting_input":
                    await asyncio.to_thread(run_store.update_status, run_id, "waiting_input")
                elif status == "running":
                    await asyncio.to_thread(run_store.update_status, run_id, "running")
            elif event_type == "run_finished":
                status = _normalize_run_status(data.get("status"))
                await asyncio.to_thread(run_store.update_status, run_id, status)
                return
    finally:
        event_bus.unsubscribe(queue)


async def process_run(
    run_id: str,
    canonical_request: CanonicalRunRequest,
    audit_context: Optional[Dict[str, Any]] = None,
    tenant_context: Optional[Dict[str, str]] = None,
):
    """Background task to execute the agent loop"""
    tenant_context = tenant_context or {
        "tenant_id": canonical_request.tenant_id,
        "tenant_tier": canonical_request.tenant_tier,
        "data_region": canonical_request.data_region or "in",
    }
    storage_namespace = storage_namespace_for_tenant(tenant_context, settings.get("tenancy", {}))
    final_result: Dict[str, Any] = {"status": "failed", "summary": "Run did not complete", "run_id": run_id}
    context = None
    results = []
    from core.loop import AgentLoop4

    loop = AgentLoop4(multi_mcp=multi_mcp)
    active_loops[run_id] = loop
    await asyncio.to_thread(
        run_store.update_status,
        run_id,
        "starting",
        query=canonical_request.query,
        integration_id=canonical_request.integration_id,
        workflow_id=canonical_request.workflow_id,
        tenant_id=canonical_request.tenant_id,
        tenant_tier=canonical_request.tenant_tier,
        data_region=canonical_request.data_region,
    )
    mirror_stop = asyncio.Event()
    mirror_task = asyncio.create_task(_mirror_run_status_from_events(run_id, mirror_stop))
    trace_token = multi_mcp.set_trace_context(
        {
            "integration_id": canonical_request.integration_id,
            "workflow_id": canonical_request.workflow_id,
            "contract_version": canonical_request.contract_version,
            "tenant_id": canonical_request.tenant_id,
            "tenant_tier": canonical_request.tenant_tier,
            "data_region": canonical_request.data_region,
            "storage_namespace": storage_namespace,
        }
    )
    try:
        # 1. RETRIEVE MEMORIES (Remme)
        # Search for past relevant facts to injecting into this run
        query = canonical_request.query
        memory_context, results = await _build_memory_context(run_id, query)

        # Execute the loop
        # The loop will maintain its own internal context
        print(f"[{run_id}] MEMORY CONTEXT READY ({len(memory_context)} chars)")
        try:
             context = await loop.run(
                 query,
                 [],
                 {
                     "_integration_meta": {
                         "integration_id": canonical_request.integration_id,
                         "workflow_id": canonical_request.workflow_id,
                         "contract_version": canonical_request.contract_version,
                     }
                 },
                 [],
                 session_id=run_id,
                 memory_context=memory_context,
                 storage_namespace=storage_namespace,
             )
        except asyncio.CancelledError:
             print(f"[{run_id}] Run cancelled.")
             context = loop.context # Recovery context from loop if possible
        
        # 2. EXTRACT NEW MEMORIES (Remme)
        # We put this in a finally block? No, because we want it only on success/completion of meaningful work.
        # But if user stops it, we might want to extract partials.
        # For now, let's leave it after run() but handle the stop case explicitly if context is returned.
        
    except Exception as e:
        print(f"Run {run_id} failed: {e}")
    finally:
        mirror_stop.set()
        if not mirror_task.done():
            mirror_task.cancel()
            try:
                await mirror_task
            except asyncio.CancelledError:
                pass
        multi_mcp.reset_trace_context(trace_token)
        # Clean up
        if run_id in active_loops:
            del active_loops[run_id]
            
        # Attempt extraction if we have context (even if stopped)
        # Note: 'context' variable needs to be accessible here.
        pass 
        # After run completes, extract new facts
        try:
            # Get the history from context (Plan Graph or Session Summary)
            # For now, we don't return the full conversation history from loop.run directly
            # But context has plan_graph... 
            # Ideally we extract from the "Summary" generated by the ReportingAgent if available
            # OR we can pass the query and the FINAL output.
            
            # Simple V1: Extract from Query + Final Answer (if available)
            final_output = ""
            # Try to find final output from graph
            if context and context.plan_graph:
                # Find nodes with output
                for node_id in context.plan_graph.nodes:
                    node = context.plan_graph.nodes[node_id]
                    if node.get("status") == "completed" and node.get("output"):
                        final_output += f"{node_id} Output: {str(node['output'])}\n"

            history = [{"role": "assistant", "content": final_output}]
            
            print(f" Remme: Extracting facts from run {run_id}...")
            # Pass existing memories from earlier search to context-aware extractor
            # ⚡ RUN IN THREAD TO AVOID BLOCKING EVENT LOOP
            commands, preferences = await asyncio.to_thread(
                remme_extractor.extract, 
                query, 
                history, 
                existing_memories=results
            )
            
            if commands:
                # Skip embedding work entirely once Ollama has failed on this
                # run - each failing call can block ~12 minutes (two endpoints
                # x OLLAMA_TIMEOUT). Fail fast and keep the event loop healthy.
                embedding_available = True
                for cmd in commands:
                    if not isinstance(cmd, dict):
                        print(f"⚠️ Remme: Skipping invalid command format: {cmd}")
                        continue

                    action = cmd.get("action")
                    text = cmd.get("text")
                    target_id = cmd.get("id")

                    try:
                        if action == "add" and text:
                            if not embedding_available:
                                continue
                            emb = await asyncio.to_thread(get_embedding, text, "search_document")
                            remme_store.add(text, emb, category="derived", source=f"run_{run_id}")
                            print(f"✅ Remme: Added new fact: {text}")
                        elif action == "update" and target_id and text:
                            if not embedding_available:
                                continue
                            emb = await asyncio.to_thread(get_embedding, text, "search_document")
                            remme_store.update_text(target_id, text, emb, source=f"run_{run_id}")
                            print(f"🔄 Remme: Updated fact {target_id}: {text}")
                        elif action == "delete" and target_id:
                            remme_store.delete(target_id, source=f"run_{run_id}")
                            print(f"🗑️ Remme: Deleted fact {target_id}")
                    except RuntimeError as e:
                        # core.embedding raises RuntimeError when Ollama is
                        # unreachable. Trip the breaker so we don't pay the
                        # timeout again for the remaining commands.
                        embedding_available = False
                        print(f"❌ Remme Action Failed (embedding unavailable): {e}")
                    except Exception as e:
                        print(f"❌ Remme Action Failed: {e}")
            
            # Apply preferences to hubs
            if preferences:
                from remme.extractor import apply_preferences_to_hubs
                apply_preferences_to_hubs(preferences)
                print(f"✅ Remme: Processed {len(preferences)} preference updates.")
            
            if commands:
                print(f"✅ Remme: Processed {len(commands)} memory updates.")
            elif not preferences:
                print(f"ℹ️ Remme: No new facts extracted from run {run_id}.")

        except Exception as e:
            print(f"⚠️ Remme Extraction Failed: {e}")
            import traceback
            traceback.print_exc()

        # 3. AUTO-SAVE REPORTS TO NOTES
            # 3. AUTO-SAVE REPORTS TO NOTES
        try:
             if context and context.plan_graph:
                import re
                
                notes_dir = PROJECT_ROOT / "data" / "Notes" / "Arcturus"
                notes_dir.mkdir(parents=True, exist_ok=True)
                
                def sanitize_filename(title):
                    title = re.sub(r'[\\/*?:"<>|#]', "", title)
                    title = title.replace("\n", " ").strip()
                    return title[:60].strip()

                def extract_title(content):
                    match = re.search(r'^#+\s+(.+)$', content, re.MULTILINE)
                    if match:
                        return match.group(1).strip()
                    # Fallback to first non-empty line
                    lines = [l.strip() for l in content.split('\n') if l.strip()]
                    if lines:
                        return lines[0]
                    return "Untitled Report"

                for node_id in context.plan_graph.nodes:
                    node = context.plan_graph.nodes[node_id]
                    agent_type = node.get("agent", "")
                    output = node.get("output", {})
                    if not output: continue

                    # Check for Formatter output keys
                    markdown = output.get("markdown_report")
                    if not markdown:
                         for k, v in output.items():
                             if k.startswith("formatted_report") and isinstance(v, str):
                                 markdown = v
                                 break
                    
                    if markdown and len(markdown) > 100:
                         title = extract_title(markdown)
                         filename = sanitize_filename(title) + ".md"
                         target_path = notes_dir / filename
                         
                         if target_path.exists() and len(target_path.read_text(encoding='utf-8')) >= len(markdown):
                             continue

                         with open(target_path, 'w', encoding='utf-8') as f:
                             f.write(markdown)
                         print(f"✅ Auto-Saved Report to Notes: {filename}")

        except Exception as e:
            print(f"⚠️ Failed to auto-save report: {e}")
            
        if context and context.plan_graph:
            # Derive status from the plan graph instead of optimistically flipping
            # to "completed". Mirrors the logic used in process_resume so a run
            # that was cancelled or crashed mid-flight is not reported as done.
            failed_node_id = None
            failed_error = None
            for node_id in context.plan_graph.nodes:
                node = context.plan_graph.nodes[node_id]
                if node.get("status") == "failed":
                    failed_node_id = node_id
                    failed_error = node.get("error")
                    break
            if failed_node_id is not None:
                final_result["status"] = "failed"
                final_result["run_id"] = run_id
                final_result["error"] = failed_error
            else:
                graph_status = context.plan_graph.graph.get("status")
                node_statuses = [
                    context.plan_graph.nodes[n].get("status", "pending")
                    for n in context.plan_graph.nodes
                    if n != "ROOT"
                ]
                if graph_status in {"completed", "failed", "paused", "stopped"}:
                    derived_status = graph_status
                elif node_statuses and all(s == "completed" for s in node_statuses):
                    derived_status = "completed"
                else:
                    # Partial / running / interrupted - keep as failed so the run
                    # does not masquerade as successful.
                    derived_status = "failed"
                final_result["status"] = derived_status
                final_result["run_id"] = run_id
        
        if context:
             try:
                 output_str = ""
                 if context.plan_graph:
                     # 1. Look for FormatterAgent output first (The Final Report)
                     for node_id in context.plan_graph.nodes:
                         node = context.plan_graph.nodes[node_id]
                         node_agent = node.get("agent", "")
                         out = node.get("output", {})
                         
                         if node_agent == "FormatterAgent" or "Format" in node_agent:
                             if isinstance(out, dict):
                                 md = out.get("markdown_report")
                                 if not md:
                                     # Try all keys for something that looks like a report
                                     for k, v in out.items():
                                         if (k.startswith("formatted_report") or k == "report") and isinstance(v, str):
                                             md = v
                                             break
                                 
                                 if md:
                                     output_str = md
                                     break
                                 
                                 # Fallback: if 'output' is a long string
                                 if isinstance(out.get("output"), str) and len(out["output"]) > 100:
                                     output_str = out["output"]
                                     break
                     
                     # 2. Fallback: Find any node with a substantial string output
                     if not output_str:
                         for node_id in reversed(list(context.plan_graph.nodes)):
                             if node_id == "ROOT": continue
                             node = context.plan_graph.nodes[node_id]
                             out = node.get("output", {})
                             
                             if isinstance(out, dict):
                                 # Try to find the largest string value in the dict (recursive search)
                                 def find_largest_string(d):
                                     largest = ""
                                     for v in d.values():
                                         if isinstance(v, str):
                                             if len(v) > len(largest):
                                                 largest = v
                                         elif isinstance(v, dict):
                                             sub = find_largest_string(v)
                                             if len(sub) > len(largest):
                                                 largest = sub
                                     return largest
                                 
                                 largest_str = find_largest_string(out)
                                 if len(largest_str) > 50:
                                     output_str = largest_str
                                     break
                                     
                             elif isinstance(out, str) and len(out) > 50:
                                 output_str = out
                                 break

                 # 3. RUTHLESS CLEANING: Remove typical JSON leakage if content is actually Markdown
                 if output_str:
                     import re
                     # If the output starts and ends with {} or [], it might be a JSON dump 
                     # that contains a markdown_report field.
                     if (output_str.startswith("{") and output_str.endswith("}")) or (output_str.startswith("[") and output_str.endswith("]")):
                         try:
                             # Try to parse it and extract the report field
                             data = json.loads(output_str)
                             if isinstance(data, dict):
                                 # Look for report-like keys
                                 for k in ["markdown_report", "formatted_report", "output", "summary", "report"]:
                                     if data.get(k) and isinstance(data[k], str) and len(data[k]) > 50:
                                         output_str = data[k]
                                         break
                         except:
                             pass
                     
                     # Final check: remove block delimiters if LLM wrapped them in ```markdown
                     output_str = re.sub(r'^```(?:markdown)?\n', '', output_str)
                     output_str = re.sub(r'\n```$', '', output_str)

                 final_result["output"] = output_str.strip() if output_str else "No substantial output found."
                 if final_result.get("status") == "failed":
                     final_result["summary"] = f"Failed: {final_result.get('error', 'Unknown error')}"
                 else:
                     final_result["summary"] = output_str.strip() if output_str else "Completed."
             except Exception as e:
                 print(f"⚠️ Extraction Error: {e}")
        
        if audit_context:
            try:
                final_status = final_result.get("status", "failed")
                error_code = final_result.get("error") if final_status == "failed" else None
                await update_request_status(
                    idempotency_key=audit_context["idempotency_key"],
                    run_id=run_id,
                    status=final_status,
                    error_code=error_code,
                )
                await log_clinical_result(
                    {
                        "run_id": run_id,
                        "request_id": audit_context.get("request_id"),
                        "idempotency_key": audit_context["idempotency_key"],
                        "query_type": _infer_query_type_from_canonical(canonical_request),
                        "normalized_result": {
                            "summary": final_result.get("summary"),
                            "output": final_result.get("output"),
                        },
                        "integration_id": canonical_request.integration_id,
                        "workflow_id": canonical_request.workflow_id,
                        "contract_version": canonical_request.contract_version,
                        "tenant_id": canonical_request.tenant_id,
                        "tenant_tier": canonical_request.tenant_tier,
                        "data_region": canonical_request.data_region,
                        "summary": final_result.get("summary"),
                        "triage_flag": "high" if final_status == "failed" else "normal",
                        "status": final_status,
                        "error_code": error_code,
                    }
                )
            except Exception as e:
                print(f"⚠️ Supabase result logging failed for run {run_id}: {e}")
        try:
            await _refresh_runs_index()
        except Exception:
            pass

        await asyncio.to_thread(
            run_store.update_status,
            run_id,
            _normalize_run_status(final_result.get("status")),
            summary=final_result.get("summary"),
            error=final_result.get("error"),
        )

        return final_result


async def process_resume(run_id: str, audit_context: Optional[Dict[str, Any]] = None):
    """Background task to resume a saved run from disk."""
    final_result: Dict[str, Any] = {"status": "failed", "summary": "Run did not complete", "run_id": run_id}
    context = None
    results = []
    query = ""
    from core.loop import AgentLoop4

    loop = AgentLoop4(multi_mcp=multi_mcp)
    active_loops[run_id] = loop
    await asyncio.to_thread(run_store.update_status, run_id, "starting")
    mirror_stop = asyncio.Event()
    mirror_task = asyncio.create_task(_mirror_run_status_from_events(run_id, mirror_stop))
    try:
        summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
        found_file = _find_session_file(run_id, summaries_dir)
        if not found_file:
            raise FileNotFoundError(f"Run {run_id} not found")

        data = await asyncio.to_thread(lambda: json.loads(found_file.read_text(encoding="utf-8", errors="ignore")))
        query = data.get("graph", {}).get("original_query") or ""
        memory_context = data.get("graph", {}).get("memory_context")
        if memory_context is None and query:
            memory_context, results = await _build_memory_context(run_id, query)

        context = await loop.resume(found_file, memory_context=memory_context)
    except Exception as e:
        print(f"Resume {run_id} failed: {e}")
    finally:
        mirror_stop.set()
        if not mirror_task.done():
            mirror_task.cancel()
            try:
                await mirror_task
            except asyncio.CancelledError:
                pass
        if run_id in active_loops:
            del active_loops[run_id]

        if context and context.plan_graph:
            for node_id in context.plan_graph.nodes:
                node = context.plan_graph.nodes[node_id]
                if node.get("status") == "failed":
                    final_result["status"] = "failed"
                    final_result["error"] = node.get("error")
                    break
            else:
                final_result["status"] = context.plan_graph.graph.get("status", "completed")

        if audit_context:
            try:
                await update_request_status(
                    idempotency_key=audit_context["idempotency_key"],
                    run_id=run_id,
                    status=final_result.get("status", "failed"),
                    error_code=final_result.get("error"),
                )
                await log_clinical_result(
                    {
                        "run_id": run_id,
                        "request_id": audit_context.get("request_id"),
                        "idempotency_key": audit_context["idempotency_key"],
                        "query_type": _infer_query_type(query),
                        "normalized_result": {
                            "summary": final_result.get("summary"),
                            "output": final_result.get("output"),
                        },
                        "summary": final_result.get("summary"),
                        "triage_flag": "high" if final_result.get("status") == "failed" else "normal",
                        "status": final_result.get("status", "failed"),
                        "error_code": final_result.get("error"),
                        "event": "resume",
                    }
                )
            except Exception as e:
                print(f"⚠️ Supabase resume logging failed for run {run_id}: {e}")
        try:
            await _refresh_runs_index()
        except Exception:
            pass

        await asyncio.to_thread(
            run_store.update_status,
            run_id,
            _normalize_run_status(final_result.get("status")),
            summary=final_result.get("summary"),
            error=final_result.get("error"),
        )

        return final_result


# === Endpoints ===

@router.post("/runs")
async def create_run(
    request: RunRequest,
    background_tasks: BackgroundTasks,
    user: Dict[str, Any] = Depends(require_supabase_user),
):
    # Millisecond precision avoids collisions when multiple runs start in the same second.
    run_id = str(int(datetime.now().timestamp() * 1000))
    request_payload = request.model_dump()
    tenant_context = resolve_tenant_context(
        request_payload=request_payload,
        user=user,
        tenancy_settings=settings.get("tenancy", {}),
    )
    request_payload.update(
        {
            "tenant_id": tenant_context["tenant_id"],
            "tenant_tier": tenant_context["tenant_tier"],
            "data_region": tenant_context["data_region"],
        }
    )
    from integrations.registry import get_integration_adapter

    adapter = get_integration_adapter(
        integration_id=request.integration_id,
        source_system=request.source_system,
        tenant_context=tenant_context,
    )
    try:
        canonical_request = adapter.to_canonical(request_payload)
    except ValueError as exc:
        raise HTTPException(status_code=422, detail=str(exc))

    payload_hash = compute_payload_hash(canonical_request.query, canonical_request.raw_payload)
    source_system = (canonical_request.source_system or "s18").strip().lower()
    idempotency_key = canonical_request.idempotency_key or build_idempotency_key(
        source_system,
        canonical_request.external_event_id or "",
        payload_hash,
    )
    request_id = f"req_{run_id}"

    audit_context = {
        "request_id": request_id,
        "idempotency_key": idempotency_key,
    }
    if can_route_to_growth(tenant_context, settings.get("tenancy", {})):
        print(
            f"[{run_id}] Growth routing hook active for tenant={tenant_context['tenant_id']} "
            f"tier={tenant_context['tenant_tier']} region={tenant_context['data_region']}"
        )

    try:
        await log_inbound_request(
            {
                "run_id": run_id,
                "request_id": request_id,
                "source_system": source_system,
                "external_event_id": canonical_request.external_event_id,
                "idempotency_key": idempotency_key,
                "payload_hash": payload_hash,
                "query": canonical_request.query,
                "raw_payload": canonical_request.raw_payload,
                "integration_id": canonical_request.integration_id,
                "workflow_id": canonical_request.workflow_id,
                "contract_version": canonical_request.contract_version,
                "auth_sub": user.get("sub"),
                "auth_email": user.get("email"),
                "consent_ref": canonical_request.consent_ref,
                "tenant_id": canonical_request.tenant_id,
                "tenant_tier": canonical_request.tenant_tier,
                "data_region": canonical_request.data_region,
                "status": "accepted",
            }
        )
    except Exception as e:
        print(f"⚠️ Supabase inbound logging failed for run {run_id}: {e}")

    await asyncio.to_thread(
        run_store.upsert_run,
        run_id=run_id,
        status="accepted",
        query=canonical_request.query,
        created_at=datetime.now().isoformat(),
        request_id=request_id,
        idempotency_key=idempotency_key,
        integration_id=canonical_request.integration_id,
        workflow_id=canonical_request.workflow_id,
        tenant_id=canonical_request.tenant_id,
        tenant_tier=canonical_request.tenant_tier,
        data_region=canonical_request.data_region,
        metadata={
            "source_system": source_system,
            "contract_version": canonical_request.contract_version,
        },
    )

    # Celery mode enqueues immediately; local mode keeps the existing in-process background path.
    if is_celery_enabled():
        await execute_run(run_id, canonical_request, audit_context, tenant_context)
    else:
        background_tasks.add_task(execute_run, run_id, canonical_request, audit_context, tenant_context)
    
    return adapter.from_canonical(
        {
        "id": run_id,
        "request_id": request_id,
        "status": "starting",
        "created_at": datetime.now().isoformat(),
        "query": canonical_request.query,
        "idempotency_key": idempotency_key,
        "tenant_id": canonical_request.tenant_id,
        "tenant_tier": canonical_request.tenant_tier,
        "data_region": canonical_request.data_region,
        "poll_timeout_seconds": get_run_poll_timeout(),
        },
        canonical_request,
    )


@router.post("/runs/{run_id}/resume")
async def resume_run(
    run_id: str,
    background_tasks: BackgroundTasks,
    user: Dict[str, Any] = Depends(require_supabase_user),
):
    if run_id in active_loops:
        return {
            "id": run_id,
            "status": "running",
            "poll_timeout_seconds": get_run_poll_timeout(),
        }

    summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
    found_file = _find_session_file(run_id, summaries_dir)
    if not found_file:
        raise HTTPException(status_code=404, detail="Run not found")

    request_id = f"resume_{run_id}"
    idempotency_key = f"resume:{run_id}"
    audit_context = {"request_id": request_id, "idempotency_key": idempotency_key}
    try:
        await log_inbound_request(
            {
                "run_id": run_id,
                "request_id": request_id,
                "source_system": "s18",
                "idempotency_key": idempotency_key,
                "payload_hash": f"resume:{run_id}",
                "query": "resume_run",
                "raw_payload": {"event": "resume", "run_id": run_id},
                "auth_sub": user.get("sub"),
                "auth_email": user.get("email"),
                "status": "accepted",
            }
        )
    except Exception as e:
        print(f"⚠️ Supabase inbound resume logging failed for run {run_id}: {e}")

    await asyncio.to_thread(
        run_store.update_status,
        run_id,
        "starting",
        metadata={"resume": True},
    )
    if is_celery_enabled():
        await execute_resume(run_id, audit_context)
    else:
        background_tasks.add_task(execute_resume, run_id, audit_context)
    return {
        "id": run_id,
        "status": "resuming",
        "created_at": datetime.now().isoformat(),
        "poll_timeout_seconds": get_run_poll_timeout(),
    }


@router.get("/runs")
async def list_runs(user: Dict[str, Any] = Depends(require_supabase_user)):
    """List runs from durable registry, with disk fallback for legacy sessions."""
    stored_runs = await asyncio.to_thread(run_store.list_runs, 300)
    if stored_runs:
        return stored_runs
    runs = await _load_runs_index()
    return [{k: v for k, v in row.items() if k != "path"} for row in runs]


@router.get("/runs/{run_id}")
async def get_run(run_id: str, user: Dict[str, Any] = Depends(require_supabase_user)):
    """Get graph state for a run"""
    stored = await asyncio.to_thread(run_store.get_run, run_id)

    # Check memory first (if running), then disk fallback.
    if run_id in active_loops:
        loop = active_loops[run_id]
        context = getattr(loop, "context", None)
        plan_graph = getattr(context, "plan_graph", None) if context else None

        if plan_graph is not None:
            react_flow = nx_to_reactflow(plan_graph)
            node_statuses = [
                data.get("status", "pending")
                for node_id, data in plan_graph.nodes(data=True)
                if node_id != "ROOT"
            ]
            if any(s == "waiting_input" for s in node_statuses):
                status = "waiting_input"
            elif any(s == "running" for s in node_statuses):
                status = "running"
            elif any(s == "failed" for s in node_statuses):
                status = "failed"
            elif node_statuses and all(s == "completed" for s in node_statuses):
                status = "completed"
            else:
                status = plan_graph.graph.get("status", "running")
            await asyncio.to_thread(run_store.update_status, run_id, _normalize_run_status(status))

            return {
                "id": run_id,
                "status": status,
                "graph": react_flow,
                "poll_timeout_seconds": get_run_poll_timeout(),
            }

        # Run exists but graph has not been initialized yet; avoid transient 404.
        return {
            "id": run_id,
            "status": "running",
            "graph": {"nodes": [], "edges": []},
            "poll_timeout_seconds": get_run_poll_timeout(),
        }
    
    # Search disk
    summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
    found_file = _find_session_file(run_id, summaries_dir)
        
    if found_file:
        data = await asyncio.to_thread(lambda: json.loads(found_file.read_text(encoding="utf-8", errors="ignore")))
        # Reconstruct Graph to use adapter
        import networkx as nx
        if "edges" in data:
            G = nx.node_link_graph(data, edges="edges")
        elif "links" in data:
            G = nx.node_link_graph(data, edges="links")
        elif "link" in data:
            # networkx default uses 'link' (singular)
            G = nx.node_link_graph(data, edges="link")
        else:
            data["edges"] = []
            G = nx.node_link_graph(data, edges="edges")
        react_flow = nx_to_reactflow(G)
        
        # Determine status: Running if in memory, else use file status
        status = "running" if run_id in active_loops else data.get("graph", {}).get("status", "completed")
        session_file = str(found_file)
        await asyncio.to_thread(
            run_store.update_status,
            run_id,
            _normalize_run_status(status),
            session_file=session_file,
        )

        return {
            "id": run_id,
            "status": status,
            "graph": react_flow,
            "poll_timeout_seconds": get_run_poll_timeout(),
        }

    # Bridge short startup races before background task registers active loop/session file.
    if run_id.isdigit():
        try:
            ts_value = int(run_id)
            created_ms = ts_value if len(run_id) >= 13 else ts_value * 1000
            age_ms = int(datetime.now().timestamp() * 1000) - created_ms
            if 0 <= age_ms <= 30000:
                return {
                    "id": run_id,
                    "status": "starting",
                    "graph": {"nodes": [], "edges": []},
                    "poll_timeout_seconds": get_run_poll_timeout(),
                }
        except (OverflowError, ValueError):
            pass

    if stored:
        return {
            "id": run_id,
            "status": _normalize_run_status(stored.get("status")),
            "graph": {"nodes": [], "edges": []},
            "poll_timeout_seconds": get_run_poll_timeout(),
            "summary": stored.get("summary"),
            "error": stored.get("error"),
            "created_at": stored.get("created_at"),
            "updated_at": stored.get("updated_at"),
        }

    raise HTTPException(status_code=404, detail="Run not found")


@router.post("/runs/{run_id}/input")
async def provide_input(
    run_id: str,
    request: UserInputRequest,
    user: Dict[str, Any] = Depends(require_supabase_user),
):
    """Provide user input to a running agent (e.g., ClarificationAgent response)"""
    if run_id in active_loops:
        loop = active_loops[run_id]
        if loop.context:
            try:
                def _is_clarification_node(node_id: str, node_data: Dict[str, Any]) -> bool:
                    status = node_data.get("status")
                    if status not in ["running", "waiting_input"]:
                        return False
                    tags = " ".join(
                        str(node_data.get(k, "")).lower()
                        for k in ["agent", "type", "label", "description"]
                    )
                    if "clarification" in tags or "clarify" in str(node_id).lower():
                        return True
                    writes = node_data.get("writes") or []
                    return any("user_clarification" in str(w).lower() for w in writes)

                selected_node_id: Optional[str] = None
                selected_node_data: Optional[Dict[str, Any]] = None

                # Prefer explicit node id from client, then fall back to first active clarification node.
                if request.node_id in loop.context.plan_graph.nodes:
                    candidate = loop.context.plan_graph.nodes[request.node_id]
                    if _is_clarification_node(request.node_id, candidate):
                        selected_node_id = request.node_id
                        selected_node_data = candidate

                if selected_node_id is None:
                    for node_id, node_data in loop.context.plan_graph.nodes(data=True):
                        if _is_clarification_node(node_id, node_data):
                            selected_node_id = node_id
                            selected_node_data = node_data
                            break

                if selected_node_id is None or selected_node_data is None:
                    raise HTTPException(status_code=400, detail="No clarification node is currently waiting for input")

                # Get the writes key for this clarification
                writes = selected_node_data.get("writes") or []
                write_key = writes[0] if writes else f"user_clarification_{selected_node_id}"

                # Store user input in globals_schema
                loop.context.plan_graph.graph.setdefault("globals_schema", {})[write_key] = request.response

                # Mark the clarification step as completed with user's response as output
                loop.context.plan_graph.nodes[selected_node_id]["output"] = {
                    "clarificationMessage": "User provided clarification",
                    "node_id": selected_node_id,
                    write_key: request.response,
                }
                loop.context.plan_graph.nodes[selected_node_id]["status"] = "completed"

                # Save the session
                await asyncio.to_thread(loop.context._save_session)
                await asyncio.to_thread(run_store.update_status, run_id, "running")

                return {
                    "id": run_id,
                    "status": "input_received",
                    "node_id": selected_node_id,
                    "stored_as": write_key,
                }
            except HTTPException:
                raise
            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Error processing input: {str(e)}")
        else:
            raise HTTPException(status_code=400, detail="Context not initialized")
    
    raise HTTPException(status_code=404, detail="Active run not found")


@router.post("/runs/{run_id}/stop")
async def stop_run(run_id: str, user: Dict[str, Any] = Depends(require_supabase_user)):
    """Stop a running agent execution"""
    if run_id in active_loops:
        loop = active_loops[run_id]
        loop.stop()
        await asyncio.to_thread(run_store.update_status, run_id, "stopped")
        return {"id": run_id, "status": "stopping"}

    stored = await asyncio.to_thread(run_store.get_run, run_id)
    if stored:
        await asyncio.to_thread(run_store.update_status, run_id, "stopped")
        return {"id": run_id, "status": "stopped"}

    raise HTTPException(status_code=404, detail="Active run not found")


@router.delete("/runs/{run_id}")
async def delete_run(run_id: str, user: Dict[str, Any] = Depends(require_supabase_user)):
    """Delete a run from disk and memory"""
    # 1. Stop if running
    if run_id in active_loops:
        loop = active_loops[run_id]
        loop.stop()
        del active_loops[run_id]
        
    # 2. Delete file
    summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
    deleted = False
    found_file = _find_session_file(run_id, summaries_dir)
    if found_file:
        try:
            found_file.unlink()
            deleted = True
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Failed to delete file: {str(e)}")
            
    if not deleted and run_id not in active_loops: # If wasn't running and file not found
        # Might be okay if it was just in memory? But we are memory-less persistence mostly
        # Let's return success if we stopped it at least, or warn
        pass

    return {"id": run_id, "status": "deleted"}


# === AGENT TESTING ENDPOINTS ===

class AgentTestRequest(BaseModel):
    input: Optional[str] = None

@router.post("/runs/{run_id}/agent/{node_id}/test")
async def test_agent(
    run_id: str,
    node_id: str,
    request: AgentTestRequest = Body(None),
    user: Dict[str, Any] = Depends(require_supabase_user),
):
    """
    Re-run a single agent in TEST MODE (sandbox).
    - Loads the session
    - Extracts the node's inputs from globals_schema
    - Runs the agent with those inputs
    - Returns the NEW output WITHOUT saving to session
    """
    try:
        # 1. Find the session file
        summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
        found_file = _find_session_file(run_id, summaries_dir)
        
        if not found_file:
            raise HTTPException(status_code=404, detail="Session not found")
        
        # 2. Load session data
        import networkx as nx
        session_data = json.loads(found_file.read_text(encoding="utf-8", errors="ignore"))
        if "edges" in session_data:
            G = nx.node_link_graph(session_data, edges="edges")
        elif "links" in session_data:
            G = nx.node_link_graph(session_data, edges="links")
        elif "link" in session_data:
            G = nx.node_link_graph(session_data, edges="link")
        else:
            session_data["edges"] = []
            G = nx.node_link_graph(session_data, edges="edges")
        
        # 3. Find the node
        if node_id not in G.nodes:
            raise HTTPException(status_code=404, detail=f"Node {node_id} not found in session")
        
        node_data = G.nodes[node_id]
        agent_type = node_data.get("agent")
        
        if not agent_type:
            raise HTTPException(status_code=400, detail="Node has no agent type")
        
        # 4. Collect inputs from globals_schema based on 'reads'
        globals_schema = G.graph.get("globals_schema", {})
        reads = node_data.get("reads", [])
        inputs = {key: globals_schema.get(key) for key in reads if key in globals_schema}
        
        # 5. Build the input payload helper
        def build_agent_input(instruction=None, previous_output=None, iteration_context=None):
            # Determine base values
            prompt_to_use = instruction or node_data.get("agent_prompt", node_data.get("description", ""))
            query_to_use = G.graph.get("original_query", "")

            # Apply overrides from request if present
            if request and request.input:
                if agent_type == "PlannerAgent":
                    query_to_use = request.input
                else:
                    # For downstream agents, the input overrides the instruction/goal
                    prompt_to_use = request.input

            payload = {
                "step_id": node_id,
                "agent_prompt": prompt_to_use,
                "reads": reads,
                "writes": node_data.get("writes", []),
                "inputs": inputs,
                "original_query": query_to_use,
                "session_context": {
                    "session_id": run_id,
                    "created_at": G.graph.get("created_at", ""),
                    "file_manifest": G.graph.get("file_manifest", []),
                },
                **({"previous_output": previous_output} if previous_output else {}),
                **({"iteration_context": iteration_context} if iteration_context else {})
            }
             # Formatter-specific additions
            if agent_type == "FormatterAgent":
                global_data = G.graph.get('globals_schema', {}).copy()
                print(f"🕵️‍♂️ DEBUG FORMATTER: run_id={run_id}")
                print(f"🕵️‍♂️ DEBUG FORMATTER: file={found_file}")
                print(f"🕵️‍♂️ DEBUG FORMATTER: globals_keys={list(global_data.keys())}")
                if 'formatted_report_T010' in global_data:
                    print(f"🕵️‍♂️ DEBUG FORMATTER: FOUND STALE KEY 'formatted_report_T010'!")
                payload["all_globals_schema"] = global_data
            return payload

        # 6. Execute with ReAct Loop (Max 15 turns)
        from agents.base_agent import AgentRunner
        from memory.context import ExecutionContextManager
        
        agent_runner = AgentRunner(multi_mcp)
        temp_context = ExecutionContextManager.__new__(ExecutionContextManager)
        temp_context.plan_graph = G
        temp_context.multi_mcp = multi_mcp

        max_turns = 15
        current_input = build_agent_input()
        iterations_data = []
        final_output = {}
        final_execution_result = None

        for turn in range(1, max_turns + 1):
            print(f"🔄 Test Mode: {agent_type} Iteration {turn}/{max_turns}")
            
            # Run Agent
            result = await agent_runner.run_agent(agent_type, current_input)
            
            if not result["success"]:
                return {
                    "status": "error",
                    "error": result.get("error", "Agent execution failed"),
                    "node_id": node_id,
                    "agent_type": agent_type
                }
            
            output = result["output"]
            final_output = output # Update final output
            iterations_data.append({"iteration": turn, "output": output})
            
            # 1. Check for 'call_tool' (ReAct)
            if output.get("call_tool"):
                tool_call = output["call_tool"]
                tool_name = tool_call.get("name")
                tool_args = tool_call.get("arguments", {})
                
                print(f"🛠️ Test Mode: Executing Tool: {tool_name}")
                
                try:
                    # Execute tool via MultiMCP
                    tool_result = await multi_mcp.route_tool_call(tool_name, tool_args)
                    
                    # Serialize result content
                    if isinstance(tool_result.content, list):
                        result_str = "\n".join([str(item.text) for item in tool_result.content if hasattr(item, "text")])
                    else:
                        result_str = str(tool_result.content)

                    # Save result to history
                    iterations_data[-1]["tool_result"] = result_str
                    
                    # Prepare input for next iteration
                    instruction = output.get("thought", "Use the tool result to generate the final output.")
                    if turn == max_turns - 1:
                         instruction += " \n\n⚠️ WARNING: This is your FINAL turn. You MUST provide the final 'output' now. Do not call any more tools. Summarize what you have."

                    current_input = build_agent_input(
                        instruction=instruction,
                        previous_output=output,
                        iteration_context={"tool_result": result_str}
                    )
                    continue # Loop to next turn

                except Exception as e:
                    print(f"Test Mode: Tool Execution Failed: {e}")
                    current_input = build_agent_input(
                        instruction="The tool execution failed. Try a different approach or tool.",
                        previous_output=output,
                        iteration_context={"tool_result": f"Error: {str(e)}"}
                    )
                    continue

            # 2. Check for call_self (Legacy/Advanced recursion)
            elif output.get("call_self"):
                # Handle code execution if needed
                if temp_context._has_executable_code(output):
                     # Pass 'inputs' as overrides so variables from prev iterations (like ipl_urls_1A) are available
                    execution_result = await temp_context._auto_execute_code(node_id, output, input_overrides=inputs)
                    final_execution_result = execution_result
                    
                    # Save result to history
                    iterations_data[-1]["execution_result"] = execution_result

                    if execution_result.get("status") == "success":
                        execution_data = execution_result.get("result", {})
                        inputs = {**inputs, **execution_data}  # Update inputs for next iteration
                
                # Prepare input for next iteration
                current_input = build_agent_input(
                    instruction=output.get("next_instruction", "Continue the task"),
                    previous_output=output,
                    iteration_context=output.get("iteration_context", {})
                )
                continue

            # 3. Success (No tool call, just output)
            else:
                 # Execute code if present (Final Iteration)
                if temp_context._has_executable_code(output):
                     # Pass 'inputs' as overrides here too
                    final_execution_result = await temp_context._auto_execute_code(node_id, output, input_overrides=inputs)
                    iterations_data[-1]["execution_result"] = final_execution_result
                    if final_execution_result:
                         final_output = temp_context._merge_execution_results(output, final_execution_result)
                break # Exit loop
        
        # 8. Get the original output for comparison
        original_output = node_data.get("output", {})
        
        # Ensure final_execution_result is passed even if loop broke early
        if not final_execution_result and iterations_data:
             final_execution_result = iterations_data[-1].get("execution_result")

        return {
            "status": "success",
            "node_id": node_id,
            "agent_type": agent_type,
            "original_output": original_output,
            "test_output": final_output,
            "execution_result": final_execution_result,
            "inputs_used": inputs,
            "iterations": iterations_data # Optional: Pass full iterations if needed by UI
        }
        
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/runs/{run_id}/agent/{node_id}/save")
async def save_agent_test(
    run_id: str,
    node_id: str,
    request: Request,
    user: Dict[str, Any] = Depends(require_supabase_user),
):
    """
    Save test results back to the session file.
    - Updates the node's output
    - Updates globals_schema with new writes
    """
    try:
        body = await request.json()
        new_output = body.get("output")
        
        if not new_output:
            raise HTTPException(status_code=400, detail="Missing 'output' in request body")
        
        # 1. Find the session file
        summaries_dir = PROJECT_ROOT / "memory" / "session_summaries_index"
        found_file = _find_session_file(run_id, summaries_dir)
        
        if not found_file:
            raise HTTPException(status_code=404, detail="Session not found")
        
        # 2. Load and update session
        import networkx as nx
        session_data = json.loads(found_file.read_text(encoding="utf-8", errors="ignore"))
        if "edges" in session_data:
            G = nx.node_link_graph(session_data, edges="edges")
        elif "links" in session_data:
            G = nx.node_link_graph(session_data, edges="links")
        elif "link" in session_data:
            G = nx.node_link_graph(session_data, edges="link")
        else:
            session_data["edges"] = []
            G = nx.node_link_graph(session_data, edges="edges")
        
        if node_id not in G.nodes:
            raise HTTPException(status_code=404, detail=f"Node {node_id} not found")

        # 2.5 SPECIAL HANDLING: PlannerAgent Graph Update
        # If this is a PlannerAgent (has plan_graph in output), we must REBUILD the graph structure.
        if "plan_graph" in new_output:
            print(f"🔄 Planner Update Detected for {node_id}. Rebuilding graph...")
            plan_graph = new_output["plan_graph"]
            
            # 1. Keep crucial nodes (ROOT and the Planner/Query node itself)
            # We assume node_id is the Planner node. 
            nodes_to_keep = ["ROOT", node_id] 
            
            # 2. Identify nodes to remove (all existing nodes except kept ones)
            nodes_to_remove = [n for n in G.nodes if n not in nodes_to_keep]
            for n in nodes_to_remove:
                G.remove_node(n)
                
            # 3. Add NEW nodes from plan
            # plan_graph['nodes'] is a list of dicts
            new_nodes = plan_graph.get("nodes", [])
            for n_data in new_nodes:
                nid = n_data["id"] 
                # Ensure we don't overwrite the planner if it's in the list for some reason (unlikely but safe)
                if nid not in G.nodes:
                    G.add_node(nid, **n_data)
                    # Initialize status for new nodes
                    G.nodes[nid]["status"] = "idle"
            
            # 4. Add NEW edges from plan
            # plan_graph['edges'] or 'links'
            new_edges = plan_graph.get("edges", plan_graph.get("links", []))
            
            # clear existing edges? We already removed nodes, so connected edges are gone.
            # But we need to ensure ROOT -> Planner connection exists if not implicitly handled.
            # In our schema, ROOT->Query (Planner). 
            # The plan_graph usually defines edges from "ROOT" to the first new node.
            # We must REMAP "ROOT" in the plan to be "node_id" (The Planner Node) 
            # so that the flow is ROOT -> Planner -> FirstNode
            
            for edge in new_edges:
                src = edge["source"]
                tgt = edge["target"]
                
                # REMAP ROOT -> Current Planner Node
                if src == "ROOT":
                    src = node_id
                    
                G.add_edge(src, tgt)
                
            # Ensure ROOT is connected to Planner (node_id)
            if not G.has_edge("ROOT", node_id):
                G.add_edge("ROOT", node_id)
                
            print(f"✅ Graph Rebuilt. Nodes: {len(G.nodes)}, Edges: {len(G.edges)}")

        node_data = G.nodes[node_id]
        writes = node_data.get("writes", [])
        
        # 3. Update node output
        node_data["output"] = new_output
        node_data["last_tested"] = datetime.now().isoformat()
        node_data["status"] = "completed"
        
        # 4. Update globals_schema with execution results if available
        # 4. Update globals_schema
        # CRITICAL FIX: Prioritize the 'merged' output (new_output) which contains the actual results
        # Execution result is less reliable as it might be raw or unmerged
        
        globals_schema = G.graph.get("globals_schema", {})
        
        # 1. Try extracting from new_output (which is test_output from frontend = merged result)
        if isinstance(new_output, dict):
             for key in writes:
                if key in new_output:
                     # Validate it's not just an empty placeholder if possible, but trust the save
                     val = new_output[key]
                     # If it's a list and not empty, or dict and not empty, update
                     if val or val == 0 or val is False: 
                         globals_schema[key] = val
                         
        # 2. Fallback to execution_result only if new_output didn't have it
        exec_result = body.get("execution_result")
        if exec_result and isinstance(exec_result, dict):
             result_data = exec_result.get("result", exec_result) # Handle {status:..., result:...} or direct
             
             for key in writes:
                 if key not in globals_schema or not globals_schema[key]: # Only if missing/empty
                     if isinstance(result_data, dict) and key in result_data:
                         globals_schema[key] = result_data[key]
        
        G.graph["globals_schema"] = globals_schema
        
        # 5. Update iterations array with execution_result if provided
        execution_result = body.get("execution_result")
        if execution_result:
            # Check if iterations exist, if not create a default one
            if not node_data.get("iterations"):
                 node_data["iterations"] = []
            
            iterations = node_data["iterations"]
            
            if iterations:
                # Update the last iteration with the new execution result
                iterations[-1]["execution_result"] = execution_result
            else:
                # Create a pseudo-iteration if none exist (e.g. single-shot agents)
                iterations.append({
                    "iteration": 1, 
                    "output": new_output,
                    "execution_result": execution_result
                })

        # 5.5. Cascading Invalidation: Mark downstream nodes as 'stale'
        # This gives visual feedback (muted opacity) in the frontend that these nodes need re-running
        try:
             descendants = nx.descendants(G, node_id)
             for desc_id in descendants:
                 if desc_id in G.nodes:
                     # Only mark as stale if they were previously completed or failed
                     # If they are 'pending', they stay pending.
                     current_status = G.nodes[desc_id].get('status')
                     if current_status in ['completed', 'failed', 'running']:
                        G.nodes[desc_id]['status'] = 'stale'
        except Exception as e:
             print(f"Warning: Failed to invalidate downstream nodes: {e}")
        
        # 6. Save back to file
        # Use edges="edges" to match our expected format (not default "link")
        graph_data = nx.node_link_data(G, edges="edges")
        with open(found_file, 'w', encoding='utf-8') as f:
            json.dump(graph_data, f, indent=2, default=str, ensure_ascii=False)
        
        # 7. AUTO-SAVE TO NOTES (for FormatterAgent "Run Again" saves)
        agent_type = G.nodes[node_id].get("agent", "")
        if agent_type == "FormatterAgent" and isinstance(new_output, dict):
            try:
                import re
                notes_dir = PROJECT_ROOT / "data" / "Notes" / "Arcturus"
                notes_dir.mkdir(parents=True, exist_ok=True)
                
                def sanitize_filename(title):
                    title = re.sub(r'[\\/*?:"<>|#]', "", title)
                    title = title.replace("\n", " ").strip()
                    return title[:60].strip()

                def extract_title(content):
                    match = re.search(r'^#+\s+(.+)$', content, re.MULTILINE)
                    if match:
                        return match.group(1).strip()
                    lines = [l.strip() for l in content.split('\n') if l.strip()]
                    return lines[0] if lines else "Untitled Report"
                
                markdown = new_output.get("markdown_report")
                if not markdown:
                    for k, v in new_output.items():
                        if k.startswith("formatted_report") and isinstance(v, str):
                            markdown = v
                            break
                
                if markdown and len(markdown) > 100:
                    title = extract_title(markdown)
                    filename = sanitize_filename(title) + ".md"
                    target_path = notes_dir / filename
                    
                    # Write or overwrite (Run Again means user explicitly wants new version)
                    with open(target_path, 'w', encoding='utf-8') as f:
                        f.write(markdown)
                    print(f"✅ Auto-Saved (Run Again) to Notes: {filename}")
                    
            except Exception as e:
                print(f"⚠️ Failed to auto-save to Notes: {e}")
        
        return {
            "status": "success",
            "node_id": node_id,
            "message": "Test results saved to session"
        }
        
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))
