app/websockets/handlers.py

"""WebSocket event dispatch table.

Each event branch is self-contained. Adding a new event type means adding
one elif block here -> nothing else needs to change.

Singletons (agents, services, db) are module-level so they survive across
requests within one process lifetime.
"""
import asyncio
import base64
import contextvars
import json
import logging
import os
import re
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from app.agents.brain_agent import BrainAgent, _fallback_topic_name, _is_generic_label
from app.services.output_cache import OutputCache
from app.agents.evaluator_agent import EvaluatorAgent
from app.agents.infinity_wiki_agent import InfinityWikiAgent
from app.agents.modality_router import VisualModalityRouter
from app.agents.formula_engine import FormulaEngine
from app.agents.text_ref_engine import TextRefEngine
from app.agents.shell_visual_engine import ShellVisualEngine
from app.agents.report_agent import ReportAgent
from app.agents.wiki_agent import WikiAgent
from app.agents.tutor_agent import TutorAgent
from app.agents.study_buddy_agent import StudyBuddyAgent
from app.agents.net_research_agent import NetResearchAgent, SubQuery
from app.observability.operation import observe_operation
from app.rag.chromadb_client import ChromaDBClient
from app.rag.ingestion import LIBRARY_COLLECTION, ingest_text
from app.rag.models import BoundingBox as EvidenceBox, ConsumerType, EvidenceRequest, SelectionAnchor
from app.rag.retrieval_service import PaperEvidenceService
from app.rag.context_composer import ContextComposer
from app.schemas.journal import JournalEntry, JournalEventType
from app.schemas.graph import NodeData
from app.schemas.visual_lesson import VisualizationRequest, VisualizationReadyPayload
from app.services.graph_state import GraphStateManager
from app.services.journal_service import JournalService
from app.services.student_memory import StudentMemoryService
from app.services.agent_control import AgentControlService
from app.services.commit_adaptation import run_commit_adaptation
from app.services.explicit_profile_capture import capture_explicit_profile_signal
from app.services.student_interaction_log import append_student_interaction
from app.services.context_pack import ContextPackService
from app.services.chat_tts_stream import ChatTTSStream
from app.services.research_context_tools import plan_research_context_tools
from app.services.context_access_policy import ContextAccessPolicy
from app.services.tts_service import get_tts_manager
from app.services.tts_text import normalize_for_speech
from app.services.transcript_diagnostics import repetition_diagnostics
from app.services.visualization_service import VisualizationService
from app.services.visual_lesson_service import VisualLessonService
from app.services.voice_transcription_service import (
    VoiceAudioInput,
    VoiceTranscriptionService,
    transcribe_voice,
    validate_voice_target,
)
from app.services.voice_telemetry import record_voice_event
from app.services.voice_benchmark_logger import (
    BROWSER_PIPELINE_LOG_PATH,
    ERROR_LOG_PATH,
    SCHEMA_VERSION,
    app_version,
    default_run_id,
    git_commit,
    log_jsonl,
    new_request_id,
    now_iso,
)
from app.services.wikipedia_service import WikipediaService
from app.services.summary_writer import build_summary_markdown
from app.services.scholar_service import fetch_top_papers
from app.websockets.connection_manager import ConnectionManager

logger = logging.getLogger(__name__)


def _log_pair_buddy_turn_phase(
    project_id: str,
    phase: str,
    started_at: float,
    *,
    now: float | None = None,
) -> None:
    elapsed_ms = max(0, round(((time.monotonic() if now is None else now) - started_at) * 1000))
    logger.info(
        "[PAIR_BUDDY] turn phase=%s project=%s elapsed_ms=%d",
        phase,
        project_id[:8],
        elapsed_ms,
    )

# All singletons are lazy -> nothing loads at import time so uvicorn binds
# to the port immediately. Models and clients initialise on first request.
_cm = ConnectionManager()
_graph_mgr = GraphStateManager()
_journal = JournalService()
_memory = StudentMemoryService()
_agent_control = AgentControlService()
_context_pack = ContextPackService()
_wikipedia = WikipediaService()

# Ephemeral in-memory cache to bridge the gap while Cognee embeds summaries in the background
_recent_session_summaries: Dict[str, List[str]] = {}

# Cache of "all document ids currently loaded in this session" -> used to scope
# RAG retrieval fairly across every uploaded paper for requests that have no
# single graph node to anchor to (Infinite Wiki's free-text selection flow).
# Populated lazily, refreshed whenever _build_graph_streaming (re)builds a graph.
_session_doc_ids_cache: Dict[str, List[str]] = {}
_stt_audio_chunk_uploads: dict[str, dict[str, Any]] = {}

STT_AUDIO_CHUNK_TTL_SECONDS = 10 * 60


def _is_demo_mode() -> bool:
    return os.getenv("DEPLOYMENT_ENV", "desktop") == "demo"


def _voice_log_transcripts() -> bool:
    return os.getenv("VOICE_LOG_TRANSCRIPTS", "true").lower() not in {"0", "false", "no", "off"}


def _base_voice_payload(event: str, request_id: str, run_id: str) -> dict[str, Any]:
    return {
        "schema_version": SCHEMA_VERSION,
        "event": event,
        "request_id": request_id,
        "run_id": run_id,
        "timestamp": now_iso(),
        "app_version": app_version(),
        "git_commit": git_commit(),
    }


def _log_frontend_timing(request_id: str, run_id: str, frontend_timing: Any) -> None:
    if not isinstance(frontend_timing, dict):
        return
    payload = _base_voice_payload("stt_frontend_timing", request_id, run_id)
    payload.update(_browser_benchmark_fields())
    payload.update({
        "client_vad_enabled": False,
        "client_vad_engine": None,
        "client_vad_segments_count": 0,
        "client_speech_duration_ms": None,
        "client_speech_start_ms": None,
        "client_speech_end_ms": None,
        "client_vad_auto_stop_used": False,
        "client_vad_auto_stop_reason": None,
    })
    payload.update(frontend_timing)
    log_jsonl(BROWSER_PIPELINE_LOG_PATH, payload)


def _cleanup_stale_stt_audio_chunks() -> None:
    now = time.time()
    stale_keys = [
        key for key, upload in _stt_audio_chunk_uploads.items()
        if now - float(upload.get("created_at", now)) > STT_AUDIO_CHUNK_TTL_SECONDS
    ]
    for key in stale_keys:
        _stt_audio_chunk_uploads.pop(key, None)


def _store_stt_audio_chunk(project_id: str, data: dict[str, Any]) -> dict[str, Any] | None:
    _cleanup_stale_stt_audio_chunks()
    upload_id = str(data.get("upload_id") or data.get("request_id") or "").strip()
    if not upload_id:
        raise ValueError("missing_upload_id")
    chunk_index = int(data.get("chunk_index", -1))
    chunk_count = int(data.get("chunk_count", 0))
    if chunk_count <= 0 or chunk_index < 0 or chunk_index >= chunk_count:
        raise ValueError("invalid_audio_chunk_index")
    chunk_text = str(data.get("audio_base64_chunk") or "")
    if not chunk_text:
        raise ValueError("missing_audio_base64_chunk")

    key = f"{project_id}:{upload_id}"
    upload = _stt_audio_chunk_uploads.get(key)
    if upload is None:
        upload = {
            "created_at": time.time(),
            "chunk_count": chunk_count,
            "chunks": {},
            "metadata": {
                "request_id": data.get("request_id"),
                "run_id": data.get("run_id"),
                "target": data.get("target"),
                "audio_mime_type": data.get("audio_mime_type"),
                "reference_max_seconds": data.get("reference_max_seconds"),
                "frontend_timing": data.get("frontend_timing"),
            },
        }
        _stt_audio_chunk_uploads[key] = upload
    if int(upload.get("chunk_count", chunk_count)) != chunk_count:
        raise ValueError("audio_chunk_count_changed")

    chunks = upload["chunks"]
    chunks[chunk_index] = chunk_text
    if len(chunks) < chunk_count:
        return None

    assembled = "".join(chunks[index] for index in range(chunk_count))
    metadata = dict(upload.get("metadata") or {})
    metadata.update({
        "audio_base64": assembled,
        "audio_chunked": True,
        "audio_chunk_count": chunk_count,
        "audio_base64_char_count": len(assembled),
    })
    _stt_audio_chunk_uploads.pop(key, None)
    return metadata


def _browser_benchmark_fields() -> dict[str, Any]:
    return {
        "benchmark_type": "browser_pipeline",
        "clip_id": None,
        "expected_text": None,
        "manual_score": None,
        "wer": None,
    }


async def _session_all_doc_ids(project_id: str) -> List[str]:
    if project_id not in _session_doc_ids_cache:
        loop = asyncio.get_event_loop()
        docs = await loop.run_in_executor(None, _session_documents, project_id)
        _session_doc_ids_cache[project_id] = [d["document_id"] for d in docs]
    return _session_doc_ids_cache[project_id]


async def _session_paper_inventory(project_id: str) -> str:
    loop = asyncio.get_event_loop()
    docs = await loop.run_in_executor(None, _session_documents, project_id)
    if not docs:
        return "No uploaded papers are registered for this project."
    lines = [
        f"{index}. {doc.get('filename', 'Untitled paper')} [document_id: {doc.get('document_id', '')}]"
        for index, doc in enumerate(docs, start=1)
    ]
    return "\n".join(lines)


async def _run_optional_web_context(query: str, selection_text: str = "") -> tuple[str, list[str]]:
    search_query = (query or selection_text or "").strip()
    if not search_query:
        return "", []
    loop = asyncio.get_event_loop()
    try:
        plan = await loop.run_in_executor(None, lambda: _get_net_research().plan(search_query, selection_text))
        sub_queries = plan.sub_queries if plan.needs_web else []
    except Exception:
        logger.exception("Lazy web context planner failed")
        sub_queries = []

    if not sub_queries:
        sub_queries = [SubQuery(entity_label=search_query[:80] or "web context", search_query=search_query)]

    try:
        findings = await asyncio.gather(*(
            _get_net_research().research_subquery(sq, _wiki.search_tavily)
            for sq in sub_queries
        ))
        web_text = "\n\n".join(
            f"--- Independently researched: {finding.entity_label} ---\n{finding.summary}"
            for finding in findings
        )
        return web_text, [sq.entity_label for sq in sub_queries]
    except Exception:
        logger.exception("Lazy web context summarization failed")
        try:
            results = await _wiki.search_tavily(search_query)
        except Exception:
            logger.exception("Lazy web context search failed")
            return "", [sq.entity_label for sq in sub_queries]
        web_text = "\n\n".join(
            f"[{item.get('title', '?')}]({item.get('url', '')})\n{item.get('content', '')}"
            for item in results
        )
        return web_text, [sq.entity_label for sq in sub_queries]


def _with_agent_control(project_id: str, agent_id: str, prior_context: str) -> str:
    """Append the editable ResearchMate operating profile to agent memory context."""
    try:
        control_context = _agent_control.compose_for_agent(project_id, agent_id)
    except Exception:
        control_context = ""
    if not control_context:
        return prior_context or ""
    if prior_context:
        return f"{prior_context}\n\nResearchMate operating profile:\n{control_context}"
    return f"ResearchMate operating profile:\n{control_context}"


async def _build_agent_memory_context(
    *,
    project_id: str,
    agent_id: str,
    query: str,
    paper_chunks: list[dict[str, Any]] | None = None,
    history: list[dict[str, Any]] | None = None,
    recent_summaries: list[str] | None = None,
    max_chars: int = 9000,
) -> str:
    try:
        context_pack = await _context_pack.build(
            agent_id=agent_id,
            query=query,
            project_id=project_id,
            paper_chunks=paper_chunks or [],
            history=history or [],
            recent_summaries=recent_summaries or [],
            max_chars=max_chars,
        )
        memory_context = context_pack.render_memory_context()
    except Exception:  # noqa: BLE001 - memory/context must not break study flow
        logger.exception("Context pack build failed for %s", agent_id)
        memory_context = ""
    return _with_agent_control(project_id, agent_id, memory_context)


async def _recall_memory_channels(
    project_id: str,
    query: str,
    consumer: ConsumerType,
) -> tuple[str, str]:
    project_items, student_items = await ContextComposer().recall_memory(
        EvidenceRequest(
            query=query or "relevant project and student context",
            consumer=consumer,
            project_id=project_id,
            token_budget=512,
            project_memory_policy="relevant",
            student_memory_policy="cached",
        )
    )
    return (
        "\n".join(item.statement for item in project_items),
        "\n".join(item.statement for item in student_items),
    )


_db: ChromaDBClient | None = None
_brain: BrainAgent | None = None
_tutor: TutorAgent | None = None
_router: VisualModalityRouter | None = None
_visuals: VisualizationService | None = None
_report: ReportAgent | None = None
_cache = OutputCache()
_wiki = WikiAgent()
_study_buddy: StudyBuddyAgent | None = None
_net_research: NetResearchAgent | None = None
_visual_lesson_service: VisualLessonService | None = None

from app.services.annotation_service import get_annotation_service
from app.services.memory_service import MemoryService
_annotations = get_annotation_service()
# Local memory: ephemeral per-PDF report clusters + persistent learning trajectory.
_local_mem = MemoryService()


def _get_db() -> ChromaDBClient:
    global _db
    if _db is None:
        _db = ChromaDBClient()
    return _db


def _get_brain() -> BrainAgent:
    global _brain
    if _brain is None:
        _brain = BrainAgent()
    return _brain


def _get_tutor() -> TutorAgent:
    global _tutor
    if _tutor is None:
        _tutor = TutorAgent()
    return _tutor


def _get_study_buddy() -> StudyBuddyAgent:
    global _study_buddy
    if _study_buddy is None:
        _study_buddy = StudyBuddyAgent()
    return _study_buddy


_VISUAL_GENERATE_STAGE_LABELS = {
    "formula": "deriving the formula",
    "graph": "writing the chart",
    "2d_text": "writing the explainer",
    "3d": "building the 3D scene",
    "2d_anim": "writing the animation",
}


def _get_router() -> VisualModalityRouter:
    global _router
    if _router is None:
        _router = VisualModalityRouter()
    return _router

def _get_visuals() -> VisualizationService:
    global _visuals
    if _visuals is None:
        client = _get_tutor()._client
        _visuals = VisualizationService(
            formula=FormulaEngine(client),
            text_ref=TextRefEngine(client),
            shell=ShellVisualEngine(client),
        )
    return _visuals


def _get_visual_lesson_service() -> VisualLessonService:
    global _visual_lesson_service
    if _visual_lesson_service is None:
        _visual_lesson_service = VisualLessonService()
    return _visual_lesson_service


_MODALITY_LABELS = {
    "formula": "Formula",
    "graph": "Chart",
    "2d_text": "Explainer",
    "3d": "3D Model",
    "2d_anim": "Animation",
}


def _get_net_research() -> NetResearchAgent:
    global _net_research
    if _net_research is None:
        _net_research = NetResearchAgent()
    return _net_research


def _get_report() -> ReportAgent:
    global _report
    if _report is None:
        _report = ReportAgent()
    return _report


_infinity: InfinityWikiAgent | None = None


def _get_infinity() -> InfinityWikiAgent:
    global _infinity
    if _infinity is None:
        _infinity = InfinityWikiAgent()
    return _infinity


async def _summarize_and_ingest_video(project_id: str, video_id: str, term: str, title: str, intention: str, target: str = "wiki"):
    """Summarize a video and ingest the summary so Quiz/Flashcards/revision can use it."""
    loop = asyncio.get_event_loop()
    summ = await loop.run_in_executor(None, _get_infinity().summarize_video, video_id, term, intention)
    text = summ.summary + ("\n" + "\n".join(f"- {p}" for p in summ.key_points) if summ.key_points else "")
    try:
        await loop.run_in_executor(
            None,
            lambda: ingest_text(text, f"YouTube: {title}", LIBRARY_COLLECTION, "content",
                                _get_db(), project_id, f"yt_{video_id}"),
        )
    except Exception as e:
        print("video summary ingest error:", e)
    await _cm.send(project_id, "WIKI_DEEPDIVE_SUMMARY", {
        "term": term, "video_id": video_id, "summary": summ.summary, "key_points": summ.key_points, "target": target,
    })


def get_connection_manager() -> ConnectionManager:
    return _cm


def get_db() -> ChromaDBClient:
    return _get_db()


def get_graph_manager() -> GraphStateManager:
    return _graph_mgr


# ------------------------------------------------------------------ #
# Helpers                                                             #
# ------------------------------------------------------------------ #

async def _get_chunks(project_id: str, query: str, n: int = 5, chunk_type: str | None = None,
                      document_ids: list[str] | None = None, consumer: str | None = None,
                      pipeline_version: str | None = None,
                      anchor_evidence_ids: list[str] | None = None,
                      selection_anchors: list[SelectionAnchor] | None = None):
    """Compatibility adapter over the structured hybrid retrieval boundary."""

    del chunk_type, pipeline_version
    try:
        consumer_type = ConsumerType(consumer or "chat")
    except ValueError:
        consumer_type = ConsumerType.CHAT
    service = ContextComposer()
    loop = asyncio.get_running_loop()
    request_context = contextvars.copy_context()
    result = await loop.run_in_executor(
        None,
        lambda: request_context.run(
            lambda: asyncio.run(service.compose(
                EvidenceRequest(
                    query=query,
                    consumer=consumer_type,
                    project_id=project_id,
                    document_ids=document_ids or [],
                    anchor_evidence_ids=anchor_evidence_ids or [],
                    selection_anchors=selection_anchors or [],
                    token_budget=max(512, n * 450),
                    project_memory_policy="none",
                    student_memory_policy="none",
                )
            ))
        ),
    )
    rows: list[dict[str, Any]] = []
    citation_by_id = {citation.evidence_id: citation for citation in result.citations}
    for index, item in enumerate(result.source_evidence[:n]):
        unit = item.evidence
        citation = citation_by_id.get(unit.evidence_id)
        filename = citation.filename if citation else ""
        page_label = (
            f"p.{unit.page_start}"
            if unit.page_start == unit.page_end
            else f"pp.{unit.page_start}-{unit.page_end}"
        )
        rows.append({
            "text": unit.index_text,
            "source": f"{filename or unit.metadata.get('source') or 'uploaded paper'}, {page_label}, {unit.evidence_id}",
            "filename": filename,
            "chunk_index": index,
            "document_id": unit.document_id,
            "evidence_id": unit.evidence_id,
            "page_start": unit.page_start,
            "page_end": unit.page_end,
            "section_path": unit.section_path,
            "element_type": unit.element_type.value,
            "bbox_norm": unit.bbox_norm.rounded() if unit.bbox_norm else None,
            "score": item.fused_score,
        })
    return rows


def _request_selection_anchors(data: Dict[str, Any]) -> list[SelectionAnchor]:
    document_id = str(data.get("selection_document_id") or "")
    region_id = str(data.get("selection_region_id") or "") or None
    if not document_id:
        return []
    anchors: list[SelectionAnchor] = []
    for snippet in data.get("selection_snippets") or []:
        try:
            boxes = [
                EvidenceBox(x=float(box["x"]), y=float(box["y"]), w=float(box["w"]), h=float(box["h"]))
                for box in snippet.get("boxes") or []
            ]
            anchors.append(SelectionAnchor(
                document_id=document_id,
                page_number=max(1, int(snippet.get("page_number") or 1)),
                boxes=boxes,
                region_id=region_id,
                text=str(snippet.get("text") or ""),
            ))
        except (KeyError, TypeError, ValueError):
            continue
    return anchors


def _session_documents(project_id: str) -> list[dict]:
    """Return per-paper canonical structural evidence for graph construction."""

    return PaperEvidenceService().project_outlines(project_id)


def _is_valid_graph(nodes: list[NodeData]) -> bool:
    """A well-formed curriculum tree has exactly one root (depth=0, parent_id=None) and
    every other node's parent_id resolves to another node in the same set. Guards against
    replaying/serving a malformed graph -> e.g. a stale cache from an older extraction bug,
    or an LLM fallback response that didn't follow the parent_id instructions -> which would
    otherwise render as disconnected, unlinked boxes with no tree structure.
    """
    if not nodes:
        return False
    ids = {n.id for n in nodes}
    if not any(n.depth == 0 and not n.parent_id for n in nodes):
        return False
    return all(n.depth == 0 or n.parent_id in ids for n in nodes)


def _ground_graph_nodes(project_id: str, nodes: list[NodeData]) -> None:
    """Attach canonical evidence and memory IDs with bounded batch reads."""

    service = PaperEvidenceService()
    evidence, memories = service.ground_graph_nodes(
        project_id,
        [
            (node.id, f"{node.label} {node.description}", node.document_ids)
            for node in nodes
        ],
    )
    for node in nodes:
        node.evidence_ids = evidence.get(node.id, [])
        node.project_memory_ids = memories.get(node.id, [])


def _safe_get_node(project_id: str, node_id: str) -> NodeData:
    try:
        return _graph_mgr.get_node(project_id, node_id)
    except KeyError:
        return NodeData(id=node_id, label=node_id, status="ongoing")


def _resolve_chunk_location(document_id: str | None, chunk_texts: list[str]) -> dict | None:
    if not document_id: return None
    import os
    pdf_path = os.path.expanduser(f"~/.studybuddy/pdfs/{document_id}.pdf")
    if not os.path.exists(pdf_path): return None
    
    try:
        import fitz
        doc = fitz.open(pdf_path)
        for chunk in chunk_texts:
            import re
            # Find a solid block of text without newlines to improve exact match chances
            blocks = [b.strip() for b in re.split(r'\n+', chunk) if len(b.strip()) >= 15]
            if not blocks: continue
            
            search_text = blocks[0][:40] # take up to 40 chars
            
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                rects = page.search_for(search_text)
                if rects:
                    pr = page.rect
                    boxes = [{"x": r.x0/pr.width, "y": r.y0/pr.height, "w": r.width/pr.width, "h": r.height/pr.height} for r in rects]
                    return {"page": page_num + 1, "boxes": boxes}
    except Exception as e:
        print("PDF search error:", e)
    return None


async def _replay_doc_graph(project_id: str, nodes: list, edges: list) -> None:
    """Stream a previously-built graph from cache (keeps the pop-in animation, no LLM calls)."""
    _graph_mgr.set_graph(project_id, nodes)
    # Order root → sections → deeper so children always attach to an existing parent.
    for n in sorted(nodes, key=lambda x: x.depth):
        await _cm.send(project_id, "GRAPH_NODE_ADDED", n.model_dump())
        await asyncio.sleep(0.05)  # small stagger so the cascade still reads as "fireworks"
    for e in edges:
        await _cm.send(project_id, "GRAPH_EDGE_ADDED", e)
    await _cm.send(project_id, "GRAPH_BUILD_DONE", {"count": len(nodes), "reused": True})


async def _compile_report(project_id: str, data: Dict[str, Any]) -> None:
    """Per-PDF report: process each note statelessly in parallel, pool the insights, then
    stream a synthesized, reworded report. A report edit re-synthesizes from the cached pool.
    """
    from app.agents.report_agent import NoteInsight

    document_id = data.get("document_id", "")
    topic = data.get("topic", "")
    intention = data.get("intention", "learn")
    knowledge_mode = "net_support" if os.getenv("TAVILY_API_KEY") else "content_only"
    edit_instruction = data.get("edit_instruction", "")
    loop = asyncio.get_event_loop()

    # Edit path: re-synthesize from the per-PDF memory cluster (no note reprocessing).
    insights: list = []
    if edit_instruction and _local_mem.read_cluster(document_id):
        insights = [NoteInsight(**d) for d in _local_mem.read_cluster(document_id)]
    else:
        # Fresh compile -> clear any stale cluster for this PDF first.
        _local_mem.flush_cluster(document_id)
        notes = _annotations.get_for_document(document_id) if document_id else []
        if not notes:
            await _cm.send(project_id, "REPORT_TOKEN", {"token":
                "_No notes yet. Highlight passages and add margin notes, or pin figures/formulas, "
                "then compile the report._"})
            await _cm.send(project_id, "REPORT_DONE", {})
            return

        await _cm.send(project_id, "REPORT_PROGRESS", {"done": 0, "total": len(notes), "stage": "reading notes"})

        async def proc(idx: int, note) -> Optional[NoteInsight]:
            snippet_text = " ".join(s.text for s in note.target_snippets)
            query = snippet_text or note.note_text or topic or "overview"
            chunks = await _get_chunks(project_id, query, n=3, consumer="report")
            ctx = "\n".join(c["text"] for c in chunks)
            extracted = note.note_text if note.image_base64 else ""
            try:
                ins = await loop.run_in_executor(
                    None, _get_report().process_note,
                    note.note_text or "", snippet_text, ctx, intention, extracted,
                )
            except Exception as e:
                print("process_note error:", e)
                ins = None
            await _cm.send(project_id, "REPORT_PROGRESS",
                           {"done": idx + 1, "total": len(notes), "stage": "reading notes"})
            return ins

        results = await asyncio.gather(*(proc(i, n) for i, n in enumerate(notes)))
        insights = [r for r in results if r]
        # Persist each processed note-insight to this PDF's ephemeral memory cluster.
        _local_mem.push_insights(document_id, [r.model_dump() for r in insights])

    web_context = ""
    report_web_sources = []
    if knowledge_mode == "net_support" and topic:
        wr = await _wiki.search_tavily(topic)
        if wr:
            web_context = "\n\n".join(f"[Web: {r.get('title')}]\n{r.get('content')}" for r in wr)
            seen_urls: set = set()
            for r in wr:
                u = r.get("url")
                if u and u not in seen_urls:
                    seen_urls.add(u)
                    report_web_sources.append({"title": r.get("title", u), "url": u})

    toc_labels = [n.label for n in _graph_mgr.list_nodes(project_id)]
    await _cm.send(project_id, "REPORT_PROGRESS", {"done": 1, "total": 1, "stage": "writing report"})

    full = ""
    async for token in _get_report().synthesize_report(
        insights, topic, toc_labels, intention,
        knowledge_mode=knowledge_mode, edit_instruction=edit_instruction, web_context=web_context,
    ):
        full += token
        await _cm.send(project_id, "REPORT_TOKEN", {"token": token})
    await _cm.send(project_id, "REPORT_DONE", {"web_sources": report_web_sources})

    # Attach a grounded visual to the report if warranted.
    try:
        ctx_chunks = [{"source": "report", "text": full}]
        decision = await loop.run_in_executor(None, _get_router().classify, topic or "report", full, ctx_chunks, intention)
        visual = await loop.run_in_executor(
            None, _get_visuals().generate, topic or "report", decision.modality, intention, ctx_chunks
        )
        await _cm.send(project_id, "REPORT_SECTION_VISUAL", {"section_id": "compiled", "visual": visual.model_dump()})
    except Exception as e:
        logger.warning("Report visual error: %s", e)


async def _build_graph_streaming(project_id: str, intention: str, topic: str, document_id: str = "") -> None:
    """Root-first, then parallel section expansion. Streams each node/edge as it resolves.

    Deliberately NOT a single call over all documents' full structure: with many uploaded
    PDFs that concatenated context would blow past the prompt budget and produce a shallow,
    lossy tree. Each section's expand call instead gets its own document's full structure,
    scaling with document count. To stop that parallelism from inventing the same subtopic
    under two different sections, each section is told what its siblings already cover.

    Falls back to a single-call tree if the streaming path fails.
    """
    loop = asyncio.get_event_loop()

    # Reuse: a graph already exists for this PDF → replay it instead of regenerating.
    cached = _graph_mgr.load_doc_graph(document_id)
    if cached and _is_valid_graph(cached[0]):
        await _replay_doc_graph(project_id, cached[0], cached[1])
        return

    all_docs = await loop.run_in_executor(None, _session_documents, project_id)
    all_doc_ids = [d["document_id"] for d in all_docs]
    _session_doc_ids_cache[project_id] = all_doc_ids  # refresh the no-node RAG-scoping cache

    # Per-paper graph tabs pass a single file hash as document_id. The older project
    # graph path passes a combined project hash (or nothing), which intentionally keeps
    # the cross-paper build behavior.
    docs = all_docs
    if document_id:
        matching_docs = [d for d in all_docs if d["document_id"] == document_id]
        if matching_docs:
            docs = matching_docs

    doc_ids = [d["document_id"] for d in docs]
    doc_names = [d["filename"] for d in docs]
    structure = "\n\n---\n\n".join(f"Document: {d['filename']}\n{d['structure']}" for d in docs)[:10000]
    if not structure:
        await _cm.send(project_id, "GRAPH_BUILD_DONE", {"count": 0})
        return

    try:
        # Cognee, load-bearing: fetch what this student already knows and feed it into
        # the curriculum build so the ROOT + SECTION shape branches on prior sessions
        # (skip/collapse what's mastered, scaffold what's weak) -> memory changes the
        # structure of the graph, not just an agent's chat context. Returns "" on a fresh
        # install or any Cognee error (the context packer is defensive),
        # so this is a no-op on day 1 and increasingly shapes the tree from day 2 on.
        prior_knowledge = await _build_agent_memory_context(
            project_id=project_id,
            agent_id="brain",
            query=topic,
            max_chars=7000,
        )
        rs = await loop.run_in_executor(
            None, _get_brain().derive_root_and_sections, structure, intention, topic, prior_knowledge, doc_names
        )
        nodes: list[NodeData] = []
        edges: list = []

        smart_fallback = ""
        if _is_generic_label(rs.root_label or topic):
            try:
                smart_fallback = await loop.run_in_executor(
                    None, _get_brain().generate_session_title, topic, doc_names, intention
                )
            except Exception:
                pass
        root = NodeData(
            id="n0", label=_fallback_topic_name(rs.root_label or topic, doc_names, smart_fallback),
            description=rs.root_description, depth=0, complexity=3,
            parent_id=None, status="ongoing",
            document_ids=doc_ids,  # root spans all papers
        )
        nodes.append(root)
        _graph_mgr.add_node(project_id, root)
        await _cm.send(project_id, "GRAPH_NODE_ADDED", root.model_dump())

        section_nodes: list[NodeData] = []
        section_doc_indices: dict[str, list[int]] = {}  # section id -> its source doc indices (for expansion)
        for i, s in enumerate(rs.sections):
            sid = f"s{i + 1}"
            valid_indices = sorted({
                idx for idx in (s.source_docs or [0]) if docs and 0 <= idx < len(docs)
            }) if docs else []
            if not valid_indices and docs:
                valid_indices = [0]
            sec_doc_ids = [doc_ids[idx] for idx in valid_indices]
            sn = NodeData(
                id=sid, label=s.label, description=s.description, depth=1,
                complexity=s.complexity, parent_id="n0", status="ongoing",
                document_ids=sec_doc_ids, memory_tag=getattr(s, "memory_tag", "new"),
            )
            section_doc_indices[sid] = valid_indices
            section_nodes.append(sn)
            nodes.append(sn)
            _graph_mgr.add_node(project_id, sn)
            await _cm.send(project_id, "GRAPH_NODE_ADDED", sn.model_dump())
            edge = {"source": "n0", "target": sid, "relationship": "prerequisite"}
            edges.append(edge)
            await _cm.send(project_id, "GRAPH_EDGE_ADDED", edge)

        # All section labels, so expand_section can tell the model what its siblings
        # already own (semantic overlap, e.g. "AdaMax" vs "AdaMax Variant" -> not just
        # exact-string dupes) -> WITHOUT concatenating every section's full document
        # context into one call, which is what makes this still scale to many PDFs.
        all_section_labels = [s.label for s in rs.sections]

        # Sections still expand concurrently (asyncio.gather) -> sibling-awareness reduces
        # overlap but can't fully eliminate it, since a section can't see its siblings'
        # ACTUAL children (not yet generated) while it runs. Exact-label dedup below is a
        # last-resort safety net for the case sibling-awareness misses.
        seen_labels: set[str] = {root.label.strip().lower()}
        seen_labels.update(s.label.strip().lower() for s in rs.sections)

        async def expand(sn: NodeData) -> tuple:
            siblings = [lbl for lbl in all_section_labels if lbl != sn.label]
            indices = section_doc_indices.get(sn.id, [])
            doc_excerpts = [
                {"index": idx, "filename": docs[idx]["filename"], "structure_text": docs[idx]["structure"]}
                for idx in indices
            ] if docs else [{"index": 0, "filename": "content", "structure_text": structure}]
            try:
                exp = await loop.run_in_executor(
                    None, _get_brain().expand_section, sn.label, doc_excerpts,
                    intention, siblings, prior_knowledge,
                )
                children = exp.children
            except Exception as e:
                print("expand_section error:", e)
                children = []
            out_nodes: list[NodeData] = []
            out_edges: list = []
            next_idx = 1
            for ch in children:
                norm_label = ch.label.strip().lower()
                if norm_label in seen_labels:
                    continue  # already covered under another section -> don't duplicate
                seen_labels.add(norm_label)
                cid = f"{sn.id}c{next_idx}"
                next_idx += 1
                child_doc_ids = [doc_ids[i] for i in (ch.source_docs or []) if docs and 0 <= i < len(docs)]
                cn = NodeData(
                    id=cid, label=ch.label, description=ch.description, depth=sn.depth + 1,
                    complexity=ch.complexity, parent_id=sn.id, status="ongoing",
                    document_ids=child_doc_ids or sn.document_ids,  # fall back to the section's papers
                    memory_tag=getattr(ch, "memory_tag", "new"),
                )
                out_nodes.append(cn)
                _graph_mgr.add_node(project_id, cn)
                await _cm.send(project_id, "GRAPH_NODE_ADDED", cn.model_dump())
                e = {"source": sn.id, "target": cid, "relationship": ch.relationship}
                out_edges.append(e)
                await _cm.send(project_id, "GRAPH_EDGE_ADDED", e)
            return out_nodes, out_edges

        results = await asyncio.gather(*(expand(sn) for sn in section_nodes))
        for child_nodes, child_edges in results:
            nodes.extend(child_nodes)
            edges.extend(child_edges)

        # Run post-generation cleanup pass to merge semantic duplicates. The cleanup LLM
        # call can silently drop nodes (truncation, or over-aggressive merging) with no
        # exception raised -> validate its output before trusting it. Reject and keep the
        # already-known-good pre-cleanup tree if the result is structurally broken (orphaned
        # parent_id) OR if it silently dropped an entire source document's content -> losing
        # a whole paper is worse than leaving a few semantic duplicates un-merged.
        doc_name_lookup = {d["document_id"]: d["filename"] for d in docs}
        try:
            cleaned_nodes, cleaned_edges = await loop.run_in_executor(
                None,
                lambda: _get_brain().cleanup_curriculum(nodes, edges, doc_name_lookup),
            )
            # Root node deliberately spans every document (document_ids=doc_ids) -> checking
            # coverage against non-root nodes only, so a document with NO real section/child
            # content left can't hide behind the root's blanket tag.
            cleaned_doc_ids = {d for n in cleaned_nodes if n.depth > 0 for d in (n.document_ids or [])}
            if not _is_valid_graph(cleaned_nodes):
                raise ValueError("cleanup produced a structurally invalid graph (orphaned parent_id)")
            if doc_ids and not set(doc_ids).issubset(cleaned_doc_ids):
                dropped = [doc_name_lookup.get(d, d) for d in doc_ids if d not in cleaned_doc_ids]
                raise ValueError(f"cleanup silently dropped source document(s): {dropped}")
            nodes, edges = cleaned_nodes, cleaned_edges
            await _cm.send(project_id, "GRAPH_CLEANUP_DONE", {
                "nodes": [n.model_dump() for n in nodes],
                "edges": edges
            })
        except Exception as cleanup_err:
            print("GRAPH_CLEANUP rejected, keeping raw streamed graph:", cleanup_err)

        _ground_graph_nodes(project_id, nodes)
        _graph_mgr.set_graph(project_id, nodes)
        _graph_mgr.save_doc_graph(document_id, nodes, edges)  # persist for reuse across sessions
        await _cm.send(project_id, "GRAPH_BUILD_DONE", {"count": len(nodes)})
    except Exception as e:
        print("BUILD_GRAPH streaming failed, falling back to single-call:", e)
        try:
            overviews = [{"filename": "content", "structure_text": structure}]
            nodes, edges = await loop.run_in_executor(
                None, _get_brain().extract_curriculum_from_documents, overviews, intention, topic, ""
            )
            if not _is_valid_graph(nodes):
                print("BUILD_GRAPH fallback produced a malformed tree (bad parent_id refs), discarding:",
                      [(n.id, n.depth, n.parent_id) for n in nodes])
                await _cm.send(project_id, "GRAPH_BUILD_DONE", {"count": 0})
                return
            _ground_graph_nodes(project_id, nodes)
            _graph_mgr.set_graph(project_id, nodes)
            _graph_mgr.save_doc_graph(document_id, nodes, edges)
            for n in nodes:
                await _cm.send(project_id, "GRAPH_NODE_ADDED", n.model_dump())
            for e2 in edges:
                await _cm.send(project_id, "GRAPH_EDGE_ADDED", e2)
            await _cm.send(project_id, "GRAPH_BUILD_DONE", {"count": len(nodes)})
        except Exception as e3:
            print("BUILD_GRAPH fallback also failed:", e3)
            await _cm.send(project_id, "GRAPH_BUILD_DONE", {"count": 0})



_UNROOTED_RAG_EVENTS: dict[str, str] = {
    "LEARN_NODE": "tutor",
    "REPORT_COMPILE": "report",
    "FLASHCARDS_REQUEST": "flashcards",
    "QUIZ_REQUEST": "quiz",
    "STUDY_BUDDY_AUDIO": "pair_buddy",
    "COMMIT_PROJECT": "project_graph",
    "CLOSE_PROJECT": "project_graph",
}


async def handle_event(project_id: str, event_type: str, data: Dict[str, Any]) -> None:
    """Dispatch one websocket event with feature-level RAG trace roots where needed."""

    consumer = _UNROOTED_RAG_EVENTS.get(event_type)
    if consumer is None:
        await _handle_event_inner(project_id, event_type, data)
        return

    with observe_operation(
        "rag.request",
        subsystem="rag",
        consumer=consumer,
        attributes={"feature": event_type.lower()},
    ):
        await _handle_event_inner(project_id, event_type, data)


async def _handle_event_inner(project_id: str, event_type: str, data: Dict[str, Any]) -> None:
    node_id: str = data.get("node_id", "")
    intention: str = data.get("intention", "learn")

    # ---- LEARN_NODE -----------------------------------------------
    if event_type == "LEARN_NODE":
        node = _safe_get_node(project_id, node_id)
        query = _get_brain().build_rag_query(data.get("node_label", node.label), intention)
        # Multi-paper: scope retrieval to this node's source paper(s).
        chunks = await _get_chunks(
            project_id,
            query,
            n=5,
            document_ids=node.document_ids or None,
            anchor_evidence_ids=node.evidence_ids or None,
            consumer="tutor",
        )

        # If chunks are empty the background indexer may still be running -> poll while it's
        # genuinely in progress, but stop immediately once we know it finished (success or error)
        # rather than repeating a "still indexing" message that will never become true.
        if not chunks:
            from app.routers.session import _ingest_status
            import asyncio as _asyncio

            waited = 0.0
            while not chunks and waited < 15.0:
                status = _ingest_status.get(project_id, "unknown")
                if status == "error":
                    break
                if status == "ready" and waited > 0:
                    break
                await _asyncio.sleep(3)
                waited += 3.0
                chunks = await _get_chunks(
                    project_id,
                    query,
                    n=5,
                    document_ids=node.document_ids or None,
                    anchor_evidence_ids=node.evidence_ids or None,
                    consumer="tutor",
                )

        if not chunks:
            status = _ingest_status.get(project_id, "unknown")
            if status == "error":
                msg = (
                    "_We couldn't index your document -> it may be a scanned/image-only PDF "
                    "with no readable text, or an unsupported format. Try re-uploading a "
                    "text-based PDF, DOCX, or TXT file._"
                )
            elif status == "ready":
                msg = (
                    "_We finished indexing your document, but couldn't find content relevant "
                    "to this topic. Try a different node, or re-upload the document if this "
                    "seems wrong._"
                )
            else:
                msg = "_Your documents are still being indexed. Please wait a moment and click the node again._"
            await _cm.send(project_id, "LESSON_TOKEN", {"token": msg})
            await _cm.send(project_id, "LESSON_DONE", {"visual_suggestion": "none"})
            return

        selection_text = data.get("selection_text", "")
        anchor_id = data.get("anchor_id") or node_id
        knowledge_mode = "net_support" if os.getenv("TAVILY_API_KEY") else "content_only"

        # Per-concept memory: adapt the lesson to what this student has shown across
        # prior sessions (Tier 1A load-bearing tenet -> Tutor reads memory). LEARN_NODE
        # is a deliberate, node-scoped action (not the fast chat path), so it can absorb
        # one Cognee read; keyed into the cache so a memory change re-generates.
        lesson_prior = await _build_agent_memory_context(
            project_id=project_id,
            agent_id="tutor",
            query=data.get("node_label", node.label),
            paper_chunks=chunks,
            max_chars=7000,
        )

        # Fetch Tavily web results if Net Support is enabled
        web_context = ""
        tavily_results: list = []
        if knowledge_mode == "net_support":
            tavily_results = await _wiki.search_tavily(data.get("node_label", node.label))
            if tavily_results:
                web_context = "\n\n".join(
                    f"[Web: {r.get('title')} -> {r.get('url')}]\n{r.get('content')}"
                    for r in tavily_results
                )
        web_sources = []
        seen_urls: set = set()
        for r in tavily_results:
            u = r.get("url")
            if u and u not in seen_urls:
                seen_urls.add(u)
                web_sources.append({"title": r.get("title", u), "url": u})

        cache_key = _cache.make_key(
            "LEARN_NODE", intention, anchor_id, [c["text"] for c in chunks],
            f"{selection_text}|{knowledge_mode}|{web_context[:100]}|{lesson_prior[:100]}"
        )
        cached_lesson = _cache.get(cache_key)
        _journal.append(
            JournalEntry(
                project_id=project_id,
                node_id=node_id,
                event_type=JournalEventType.NODE_OPENED,
                data={"node_label": data.get("node_label"), "cache_hit": cached_lesson is not None},
            )
        )
        if cached_lesson:
            # Replay cached lesson as tokens to preserve streaming UX
            for word in cached_lesson.split(" "):
                await _cm.send(project_id, "LESSON_TOKEN", {"token": word + " "})
            await _cm.send(project_id, "LESSON_DONE", {"visual_suggestion": "canvas", "web_sources": web_sources})
        else:
            full_lesson = ""
            async for token in _get_tutor().stream_lesson(
                node, chunks, intention, knowledge_mode=knowledge_mode,
                web_context=web_context, prior_knowledge=lesson_prior,
            ):
                full_lesson += token
                await _cm.send(project_id, "LESSON_TOKEN", {"token": token})
            _cache.put(cache_key, full_lesson)
            await _cm.send(project_id, "LESSON_DONE", {"visual_suggestion": "canvas", "web_sources": web_sources})

    # ---- BUILD_GRAPH (parallel curriculum streaming -> "fireworks") ------
    elif event_type == "BUILD_GRAPH":
        await _build_graph_streaming(
            project_id,
            intention,
            data.get("topic", ""),
            data.get("document_id", ""),
        )

    # ---- REPORT_COMPILE (notes → per-PDF report, launched from the graph window) ----
    elif event_type == "REPORT_COMPILE":
        await _compile_report(project_id, data)

    # ---- REPORT_CLOSE -> flush the ephemeral per-PDF report memory cluster ----
    elif event_type == "REPORT_CLOSE":
        _local_mem.flush_cluster(data.get("document_id", ""))

    # ---- VISUALIZATION_REQUEST ------------------------------------
    elif event_type in {"VISUALIZATION_REQUEST", "CHAT_VISUAL_REQUEST"}:
        prompt = (data.get("prompt") or data.get("topic") or data.get("selection_text") or "").strip()
        request_id = str(data.get("request_id") or new_request_id("visualization"))
        if not prompt:
            await _cm.send(project_id, "VISUALIZATION_ERROR", {
                "request_id": request_id,
                "code": "missing_prompt",
                "message": "Select paper context or name what you want to visualize.",
                "retryable": False,
            })
            return

        project_context_enabled = bool(data.get("project_context_enabled", True))
        policy = ContextAccessPolicy(
            project_context_enabled=project_context_enabled,
            explicit_context_available=bool(data.get("selection_text") or data.get("selection_image_base64")),
        )
        policy.log(request_id=request_id, event_type="visualization")
        visual_doc_ids: list[str] = []
        if policy.allow_project_retrieval:
            visual_doc_ids = list(data.get("active_document_ids") or []) or await _session_all_doc_ids(project_id)
        try:
            request = VisualizationRequest(
                request_id=request_id,
                prompt=prompt,
                intention=str(data.get("intention") or "learn"),
                selection_text=str(data.get("selection_text") or ""),
                surrounding_context=str(data.get("surrounding_context") or ""),
                selection_snippets=list(data.get("selection_snippets") or []),
                selection_image_base64=str(data.get("selection_image_base64") or ""),
                active_document_ids=visual_doc_ids,
                project_context_enabled=project_context_enabled,
                familiarity=str(data.get("familiarity") or "graduate"),
                resolution_token=str(data.get("resolution_token") or ""),
                resolution_candidate_id=str(data.get("resolution_candidate_id") or ""),
            )

            async def progress(stage: str, message: str) -> None:
                await _cm.send(project_id, "VISUALIZATION_PROGRESS", {
                    "request_id": request_id,
                    "stage": stage,
                    "message": message,
                })

            await progress("routing", "Choosing the verified visualization capability")
            lesson_service = _get_visual_lesson_service()
            request = lesson_service.prepare_request(request)
            capability = (
                "molecular" if request.resolution_token.startswith("molecular_")
                else "protein" if request.resolution_token
                else await lesson_service.route(request)
            )
            if capability == "clarification":
                ready = lesson_service.generate_clarification(request)
                await _cm.send(project_id, "VISUALIZATION_READY", ready.model_dump(mode="json"))
                return
            if capability in {"attention", "protein", "nucleic_acid", "molecular", "thermodynamics", "differential_equation", "graph_algorithm", "array_algorithm", "composition"}:
                if hasattr(lesson_service, "generate_verified"):
                    ready = await lesson_service.generate_verified(capability, project_id, request, progress=progress)
                elif capability == "attention":
                    ready = await lesson_service.generate_attention(project_id, request, progress=progress)
                else:
                    ready = await lesson_service.generate_protein(project_id, request, progress=progress)
                await _cm.send(project_id, "VISUALIZATION_READY", ready.model_dump(mode="json"))
                return
            raise ValueError("No verified visualization handler accepted this request")
        except Exception as exc:
            logger.exception("Visualization request failed")
            error_code = str(getattr(exc, "code", "generation_failed"))
            await _cm.send(project_id, "VISUALIZATION_ERROR", {
                "request_id": request_id,
                "code": error_code,
                "message": str(exc) or "The visualization could not be generated.",
                "retryable": error_code not in {
                    "identifier_not_found", "protein_name_missing", "structure_not_found", "unsupported_biomolecule",
                    "invalid_resolution_choice", "invalid_nucleic_sequence", "mixed_nucleic_alphabet", "sequence_too_long",
                },
            })

    # ---- CHAT_TURN -----------------------------------------------
    elif event_type == "CHAT_TURN":
        query = data.get("content", "")
        message_id = str(data.get("message_id") or new_request_id("chat"))
        chat_tts = await ChatTTSStream.create(
            manager=get_tts_manager(),
            project_id=project_id,
            message_id=message_id,
            send=lambda event, payload: _cm.send(project_id, event, payload),
            input_mode=str(data.get("input_mode") or "text"),
            auto_read_reply=bool(data.get("auto_read_reply", False)) and not _is_demo_mode(),
            model_id=str(data.get("tts_model_id") or "pocket-tts"),
            voice_id=str(data.get("tts_voice_id") or "anshuman-normal-custom"),
        )

        async def emit_chat_token(token: str) -> None:
            await _cm.send(project_id, "CHAT_TOKEN", {"token": token, "message_id": message_id})
            await chat_tts.feed(token)

        selection_text = data.get("selection_text", "")
        surrounding_context = data.get("surrounding_context", "")
        project_context_enabled = bool(data.get("project_context_enabled", True))
        chat_policy = ContextAccessPolicy(
            project_context_enabled=project_context_enabled,
            explicit_context_available=bool(selection_text or data.get("selection_image_base64")),
        )
        chat_policy.log(request_id=message_id, event_type="chat")
        knowledge_mode = "net_support" if chat_policy.web_available else "content_only"
        context_tool_plan = plan_research_context_tools(
            agent_id="chat",
            query=query,
            selection_text=selection_text,
            project_context_enabled=project_context_enabled,
            allow_web=chat_policy.web_available,
        )
        if context_tool_plan.needs_memory_clarification:
            await emit_chat_token(context_tool_plan.clarification_prompt)
            await chat_tts.finish()
            await _cm.send(project_id, "CHAT_DONE", {"message_id": message_id})
            return
        with observe_operation(
            "rag.request", subsystem="rag", consumer="chat",
        ) as root_op:
            _rag_request_start = time.monotonic()
            _first_token_time: float | None = None
            chunks: list[dict[str, Any]] = []
            paper_inventory = ""
            if context_tool_plan.use_project_context:
                all_doc_ids = await _session_all_doc_ids(project_id)
                paper_inventory = await _session_paper_inventory(project_id)
                chunks = await _get_chunks(
                    project_id, context_tool_plan.context_query or "project papers", n=5,
                    document_ids=all_doc_ids or None, consumer="chat",
                    selection_anchors=_request_selection_anchors(data),
                )
            merge_note = "\n\nThe active project may contain multiple papers. Keep paper-specific claims distinct and cite the relevant source label.\n"
            selection_img = data.get("selection_image_base64")
            selection_prefix = ""
            if selection_text or selection_img:
                selection_prefix = "The student has provided a specific selection (image and/or text) from the document.\n"
                selection_prefix += "CRITICAL INSTRUCTION: Your primary task is to explain and address THIS specific selection directly. Do NOT just summarize the general SOURCE MATERIAL. Use the SOURCE MATERIAL only to provide supporting context for the selection.\n\n"
                if selection_text:
                    selection_prefix += f"Student selected text:\n\"{selection_text}\"\n"
                if surrounding_context:
                    selection_prefix += f"Surrounding context:\n{surrounding_context[:4000]}\n"
                selection_prefix += "\n"

            # The model's training data has a stale cutoff and will confidently hallucinate
            # "today's date" (and reason about relative time -> "this year", "recently") from
            # it unless told the truth directly. The backend knows the real date for free ->
            # no need to burn a web_search call or leave it to the model's guesswork.
            import datetime as _dt
            _today_note = f"Today's date is {_dt.date.today().strftime('%A, %B %d, %Y')}.\n\n"
            _chat_interaction = append_student_interaction(
                project_id=project_id, source="chat_turn", agent_id="tutor", text=query,
            )
            await capture_explicit_profile_signal(
                project_id=project_id, interaction_id=_chat_interaction.interaction_id, user_text=query,
            )
            history_msgs = data.get("history", [])
            memory_context = ""
            if context_tool_plan.use_memory_context:
                try:
                    memory_query = context_tool_plan.memory_query or query or selection_text or "student profile"
                    project_memory, student_memory = await _recall_memory_channels(
                        project_id, memory_query, ConsumerType.CHAT
                    )
                    memory_context = (
                        f"PROJECT MEMORY (Chroma; project-specific research state):\n{project_memory or 'None'}\n\n"
                        f"STUDENT MEMORY (Cognee; cross-project learner profile):\n{student_memory or 'None'}"
                    )
                except Exception:  # noqa: BLE001 - memory must not break chat
                    logger.exception("Chat narrow memory recall failed")
            with observe_operation("context.assembly", subsystem="rag", consumer="chat") as assembly_op:
                chunk_ctx = "\n\n".join(
                    f"[{c.get('source') or c.get('filename') or '?'}]\n{str(c.get('text', ''))[:1800]}"
                    for c in chunks
                )
                assembly_op.add_count("selected_candidate_count", len(chunks))
                assembly_op.add_count("context_selected_chars", len(chunk_ctx))
            memory_prefix = (
                "MEMORY CONTEXT (two explicitly separated scopes):\n"
                f"{memory_context}\n\n"
                "Use MEMORY CONTEXT for questions about the student, their preferences, their name, "
                "past study activity, or this project's remembered state. Do not use it to invent "
                "facts about uploaded papers. MEMORY CONTEXT is not authoritative for the current uploaded "
                "paper list; CURRENT UPLOADED PAPERS below is authoritative. Paper facts still need SOURCE "
                "MATERIAL or WEB RESEARCH.\n\n"
                if memory_context else ""
            )
            paper_inventory_prefix = (
                "CURRENT UPLOADED PAPERS (authoritative project registry; never add papers that are not listed here):\n"
                f"{paper_inventory}\n\n"
                if paper_inventory else ""
            )
            tool_plan_prefix = (
                "OPTIONAL CONTEXT TOOLS USED THIS TURN: "
                f"{context_tool_plan.reason}.\n"
                "If SOURCE MATERIAL, MEMORY CONTEXT, CURRENT UPLOADED PAPERS, or WEB RESEARCH is absent, "
                "that tool was not fetched for this turn; do not pretend to have it.\n\n"
            )

            _CHAT_FORMATTING = (
                "Formatting:\n"
                "- Write math in LaTeX: inline $...$ and display $$...$$.\n"
                "- Use GitHub-style Markdown tables (a header row, then a |---|---| separator row).\n"
                "- For structural concepts you MAY add a ```mermaid fenced diagram. IMPORTANT: Mermaid "
                "does NOT render math -> use PLAIN-TEXT node labels only, with no $...$ and no backslash "
                "commands (write 'theta', 'gradient of loss', 'nabla L' -> never '$\\nabla \\ell$'). "
                "ALWAYS wrap a node label in double quotes if it contains parentheses or symbols, "
                "e.g. B[\"Low Memory Cost: O(n)\"]. Use vertical student-readable flowcharts: "
                "prefer `graph TD` or `flowchart TD`, never `graph LR`/`flowchart LR`. Keep diagrams "
                "to 4-6 nodes, one idea per node, and labels under 32 characters so they remain readable "
                "inside the chat panel.\n"
                "- For concrete numeric data, explain the values in prose. Interactive charts are authored only through "
                "the verified visualization composition pipeline, never as raw Plotly JSON.\n\n"
            )
            source_material_section = f"SOURCE MATERIAL:\n{chunk_ctx}" if chunk_ctx else ""

            loop = asyncio.get_event_loop()
            full_response = ""

            if knowledge_mode == "net_support":
                # Hybrid tutor: grounded in the source material OR the live web -> never raw model
                # weights for factual claims. Net Support widens the grounding to the web; it does
                # NOT license answering from memorized/parametric knowledge.
                system_msg = (
                    "You are ResearchMate, an expert research companion helping a student understand selected context "
                    "or the papers they uploaded. Be genuinely helpful and substantive -> explain, define, give intuition, "
                    "analogies, and worked reasoning. Tailor the depth to the "
                    f"{intention} goal.{merge_note}\n\n"
                    f"{_today_note}"
                    f"{tool_plan_prefix}"
                    f"{memory_prefix}"
                    f"{paper_inventory_prefix}"
                    f"{selection_prefix}"
                    "If the student supplied a selection, answer that selection first. Otherwise use the optional context tool results: "
                    "paper/project answers from SOURCE MATERIAL, current/external answers from WEB RESEARCH, and personal/project-history "
                    "answers from MEMORY CONTEXT. "
                    "If the student asks which papers are in the project, answer only from CURRENT UPLOADED PAPERS. "
                    "Use SOURCE MATERIAL below as your primary anchor when it is present and relevant. If the "
                    "student's question involves a fact, definition, dataset, statistic, or claim that is "
                    "NOT covered by SOURCE MATERIAL, WEB RESEARCH may appear below only if the web tool "
                    "was selected for this turn; rely on that section for such facts. You have NO tool-calling capability "
                    "in this conversation -> do NOT write out a tool call, function call, or any "
                    "'calling X' syntax yourself under any circumstance. If a fact is not covered by the "
                    "SOURCE MATERIAL or the WEB RESEARCH section, say plainly that you don't have a "
                    "grounded source for it instead of guessing or inventing one. The only things you may "
                    "explain directly without a source are pure reasoning, intuition, analogies, or "
                    "restating/clarifying material that IS already in the SOURCE MATERIAL.\n\n"
                    "Attribution rules:\n"
                    "- A fact from the student's material → cite inline like [Source: <label>].\n"
                    "- A fact from the web → cite with a Markdown link copied from WEB RESEARCH, like [Title](URL).\n"
                    "- Pure intuition/analogy/reasoning (no external fact) → no citation needed.\n"
                    "- NEVER invent a citation or attribute a claim to the source if it is not there.\n"
                    "- NEVER state a fact, figure, or claim you have not grounded in SOURCE MATERIAL or "
                    "the WEB RESEARCH section.\n"
                    "- If WEB RESEARCH is present but has no usable source link for a claim, omit that claim.\n\n"
                    f"{_CHAT_FORMATTING}"
                    f"{source_material_section}"
                )
                user_content = query or f"Explain: {selection_text[:200] if selection_text else 'this image'}"
                if selection_img:
                    user_msg = {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": user_content},
                            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{selection_img}"}}
                        ]
                    }
                else:
                    user_msg = {"role": "user", "content": user_content}

                messages = [{"role": "system", "content": system_msg}]
                for h in history_msgs[-10:]:
                    role = "user" if h.get("role") in ("student", "user") else "assistant"
                    messages.append({"role": role, "content": h.get("content", "")})
                messages.append(user_msg)
                if context_tool_plan.use_web_search:
                    web_text, entities = await _run_optional_web_context(query or selection_text, selection_text)
                    if entities:
                        await _cm.send(project_id, "CHAT_TOOL", {
                            "tool": "web_search", "status": "running",
                            "entities": entities,
                        })
                    if web_text:
                        messages.append({
                            "role": "system",
                            "content": (
                                "WEB RESEARCH (each section below was researched independently when possible -> "
                                f"do NOT blend facts between sections):\n{web_text}"
                            ),
                        })
                async for token in _get_tutor()._client.stream_complete(messages):
                    if _first_token_time is None:
                        _first_token_time = time.monotonic()
                        ttft_ms = round((_first_token_time - _rag_request_start) * 1000.0, 3)
                        root_op.set("backend_ttft_ms", ttft_ms)
                        root_op.event("first_token", {"backend_ttft_ms": ttft_ms})
                    full_response += token
                    await emit_chat_token(token)
            else:
                # content_only mode: strictly grounded in the uploaded material, no web, no outside facts.
                system_msg = (
                    "You are ResearchMate, helping a student understand selected context or uploaded project papers, at the "
                    f"{intention} goal.{merge_note}\n\n"
                    f"{_today_note}"
                    f"{tool_plan_prefix}"
                    f"{memory_prefix}"
                    f"{paper_inventory_prefix}"
                    f"{selection_prefix}"
                    "If the student supplied a selection, answer that selection first. Otherwise use SOURCE MATERIAL only when the "
                    "project-context tool was fetched for this turn. "
                    "If the student asks which papers are in the project, answer only from CURRENT UPLOADED PAPERS. "
                    "Answer paper/content questions using ONLY SOURCE MATERIAL when it is present. Cite inline like [Source: <label>]. "
                    "For questions about the student or project memory, use MEMORY CONTEXT if present. "
                    "If neither memory nor the material contains the answer, say so plainly and point to the nearest "
                    "relevant part -> never fabricate facts or pull from outside knowledge.\n\n"
                    f"{_CHAT_FORMATTING}"
                    f"{source_material_section}"
                )
                user_content = query or f"Explain: {selection_text[:200] if selection_text else 'this image'}"
                if selection_img:
                    user_msg = {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": user_content},
                            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{selection_img}"}}
                        ]
                    }
                else:
                    user_msg = {"role": "user", "content": user_content}

                messages = [{"role": "system", "content": system_msg}]
                for h in history_msgs[-10:]:
                    role = "user" if h.get("role") in ("student", "user") else "assistant"
                    messages.append({"role": role, "content": h.get("content", "")})
                messages.append(user_msg)
                async for token in _get_tutor()._client.stream_complete(messages):
                    if _first_token_time is None:
                        _first_token_time = time.monotonic()
                        ttft_ms = round((_first_token_time - _rag_request_start) * 1000.0, 3)
                        root_op.set("backend_ttft_ms", ttft_ms)
                        root_op.event("first_token", {"backend_ttft_ms": ttft_ms})
                    full_response += token
                    await emit_chat_token(token)

            with observe_operation("citation.attach", subsystem="rag", consumer="chat") as cite_op:
                cite_op.add_count("citation_count", full_response.count("[Source:"))

            _journal.append(
                JournalEntry(
                    project_id=project_id,
                    node_id=node_id,
                    event_type=JournalEventType.CHAT_TURN,
                    data={"role": "student", "content": query, "selection_text": selection_text, "response": full_response},
                )
            )
            await chat_tts.finish()
            await _cm.send(project_id, "CHAT_DONE", {"message_id": message_id})

    # ---- CHAT_COMPACT --------------------------------------------
    # Distills a long chat's history into a short brief so the frontend can drop the
    # verbatim messages -> keeps localStorage and the per-turn WS payload from growing
    # unbounded as a single conversation runs long (see MAX_COMPACT_TRANSCRIPT_CHARS).
    elif event_type == "CHAT_COMPACT":
        message_id = str(data.get("message_id") or new_request_id("compact"))
        history_msgs = data.get("history", [])
        transcript = "\n\n".join(
            f"{'Student' if h.get('role') in ('student', 'user') else 'Assistant'}: {h.get('content', '')}"
            for h in history_msgs if h.get("content")
        )
        if not transcript.strip():
            await _cm.send(project_id, "CHAT_COMPACTED", {"message_id": message_id, "summary": ""})
            return
        MAX_COMPACT_TRANSCRIPT_CHARS = 20000
        if len(transcript) > MAX_COMPACT_TRANSCRIPT_CHARS:
            transcript = transcript[-MAX_COMPACT_TRANSCRIPT_CHARS:]
        system_msg = (
            "Compact the following research-assistant chat transcript into a short context brief so the "
            "conversation can continue without the full history. Preserve concrete facts, decisions, paper or "
            "topic names, and any open questions the student may return to. Do not invent anything not present "
            "in the transcript. Write 150-300 words of plain prose, no headers or bullet lists.\n\n"
            f"Transcript:\n{transcript}"
        )
        summary = ""
        async for token in _get_tutor()._client.stream_complete([{"role": "system", "content": system_msg}]):
            summary += token
        await _cm.send(project_id, "CHAT_COMPACTED", {"message_id": message_id, "summary": summary.strip()})

    # ---- TTS playback controls ----------------------------------
    elif event_type == "TTS_START":
        message_id = str(data.get("message_id") or new_request_id("chat"))
        if _is_demo_mode():
            await _cm.send(project_id, "TTS_ERROR", {
                "project_id": project_id,
                "message_id": message_id,
                "code": "demo_mode",
                "retryable": False,
                "message": "Voice playback is disabled in demo mode",
            })
            return
        text = str(data.get("text") or "")
        manager = get_tts_manager()
        generation_id = await manager.begin(
            project_id,
            message_id,
            lambda event, payload: _cm.send(project_id, event, payload),
            model_id=str(data.get("model_id") or "pocket-tts"),
            voice_id=str(data.get("voice_id") or "anshuman-normal-custom"),
        )
        utterance = normalize_for_speech(text)
        if utterance:
            await manager.enqueue(generation_id, 0, utterance)
        await manager.finish(generation_id)

    elif event_type == "TTS_CANCEL":
        generation_id = str(data.get("generation_id") or "")
        if generation_id:
            await get_tts_manager().cancel(generation_id)

    elif event_type == "TTS_PLAYBACK_STATUS":
        generation_id = str(data.get("generation_id") or "")
        if generation_id:
            await get_tts_manager().report_playback_status(
                generation_id,
                str(data.get("outcome") or ""),
                float(data.get("audio_played_ms") or 0.0),
            )

    # ---- FLASHCARDS_REQUEST --------------------------------------
    elif event_type == "FLASHCARDS_REQUEST":
        label = (data.get("topic") or "").strip()
        if not label:
            await _cm.send(project_id, "ERROR", {
                "event_type": "FLASHCARDS_REQUEST",
                "message": "Flashcards require an explicit topic.",
            })
            return
        anchor_id = data.get("anchor_id") or f"flashcards:{label}"
        selection_text = data.get("selection_text", "")
        chunks: list[dict[str, Any]] = []
        if bool(data.get("project_context_enabled", True)):
            chunks = await _get_chunks(
                project_id, label, n=8, chunk_type="question", consumer="flashcards"
            )
            if not chunks:
                chunks = await _get_chunks(project_id, label, n=8, consumer="flashcards")

        cache_key = _cache.make_key("FLASHCARDS_REQUEST", intention, anchor_id, [c["text"] for c in chunks], selection_text)
        cached = _cache.get(cache_key)
        if cached:
            doc_id = chunks[0].get("document_id") if chunks else None
            for card in cached.get("cards", []):
                if "source_chunk_text" not in card:
                    idxs = card.get("source_chunk_indexes", [])
                    c_texts = [chunks[i]["text"] for i in idxs if i < len(chunks)]
                    card["source_chunk_text"] = "\n\n".join(c_texts)
            await _cm.send(project_id, "FLASHCARDS_READY", cached)
        else:
            result = _get_tutor().generate_flashcards(label, chunks, intention)
            payload = result.model_dump()

            doc_id = chunks[0].get("document_id") if chunks else None
            for card in payload["cards"]:
                idxs = card.get("source_chunk_indexes", [])
                c_texts = [chunks[i]["text"] for i in idxs if i < len(chunks)]
                card["source_location"] = _resolve_chunk_location(doc_id, c_texts)
                card["source_chunk_text"] = "\n\n".join(c_texts)

            _cache.put(cache_key, payload)
            await _cm.send(project_id, "FLASHCARDS_READY", payload)

    # ---- QUIZ_REQUEST --------------------------------------------
    elif event_type == "QUIZ_REQUEST":
        label = (data.get("topic") or "").strip()
        if not label:
            await _cm.send(project_id, "ERROR", {
                "event_type": "QUIZ_REQUEST",
                "message": "Quiz requires an explicit topic.",
            })
            return
        anchor_id = data.get("anchor_id") or f"quiz:{label}"
        selection_text = data.get("selection_text", "")
        chunks: list[dict[str, Any]] = []
        if bool(data.get("project_context_enabled", True)):
            chunks = await _get_chunks(
                project_id, label, n=8, chunk_type="question", consumer="quiz"
            )
            if not chunks:
                chunks = await _get_chunks(project_id, label, n=8, consumer="quiz")

        cache_key = _cache.make_key("QUIZ_REQUEST", intention, anchor_id, [c["text"] for c in chunks], selection_text)
        cached = _cache.get(cache_key)
        if cached:
            doc_id = chunks[0].get("document_id") if chunks else None
            for q in cached.get("questions", []):
                if "source_chunk_text" not in q:
                    idxs = q.get("source_chunk_indexes", [])
                    c_texts = [chunks[i]["text"] for i in idxs if i < len(chunks)]
                    q["source_chunk_text"] = "\n\n".join(c_texts)
            await _cm.send(project_id, "QUIZ_READY", cached)
        else:
            result = _get_tutor().generate_quiz(label, chunks, intention)
            payload = result.model_dump()

            doc_id = chunks[0].get("document_id") if chunks else None
            for q in payload["questions"]:
                idxs = q.get("source_chunk_indexes", [])
                c_texts = [chunks[i]["text"] for i in idxs if i < len(chunks)]
                q["source_location"] = _resolve_chunk_location(doc_id, c_texts)
                q["source_chunk_text"] = "\n\n".join(c_texts)

            _cache.put(cache_key, payload)
            await _cm.send(project_id, "QUIZ_READY", payload)

    # ---- FLASHCARD_GRADE -----------------------------------------
    elif event_type == "FLASHCARD_GRADE":
        _journal.append(
            JournalEntry(
                project_id=project_id,
                node_id=node_id,
                event_type=JournalEventType.FLASHCARD_GRADE,
                data=data,
            )
        )

    # ---- QUIZ_SUBMIT ---------------------------------------------
    elif event_type == "QUIZ_SUBMIT":
        _journal.append(
            JournalEntry(
                project_id=project_id,
                node_id=node_id,
                event_type=JournalEventType.QUIZ_SUBMIT,
                data=data,
            )
        )
        await _cm.send(
            project_id,
            "QUIZ_FEEDBACK",
            {"correct": data.get("correct"), "was_correct": data.get("was_correct")},
        )

    # ---- STUDY_BUDDY_INIT ----------------------------------------
    elif event_type == "STUDY_BUDDY_INIT":
        intention_level = data.get("intention", intention)
        feynman_mode = data.get("feynman_mode", False)
        full = ""
        with observe_operation(
            "rag.request", subsystem="rag", consumer="pair_buddy"
        ) as root_op:
            try:
                node_label = "project papers"
                chunks: list[dict[str, Any]] = []
                pair_doc_ids: list[str] = []
                context_tools_summary = ""
                if feynman_mode:
                    pair_doc_ids = await _session_all_doc_ids(project_id)
                    chunks = await _get_chunks(
                        project_id,
                        "project papers",
                        n=5,
                        document_ids=pair_doc_ids or None,
                        selection_anchors=_request_selection_anchors(data),
                        consumer="pair_buddy",
                    )
                    context_tools_summary = "project_context"
                project_memory, student_memory = await _recall_memory_channels(
                    project_id, node_label, ConsumerType.PAIR_BUDDY
                )
                student_profile = _with_agent_control(
                    project_id,
                    "pair_buddy",
                    f"PROJECT MEMORY:\n{project_memory or 'None'}\n\nSTUDENT MEMORY:\n{student_memory or 'None'}",
                )

                async for token in _get_study_buddy().generate_initial_question(
                    node_label, chunks, intention_level, student_profile,
                    is_merged=len(pair_doc_ids) > 1, merge_summary="",
                    context_tools_summary=context_tools_summary,
                ):
                    full += token
                    await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": token})
                root_op.add_count("returned_evidence_count", len(chunks))
            except Exception as exc:
                logger.exception("StudyBuddy init failed")
                root_op.record_exception(
                    exc,
                    escaped=False,
                    stage="context_composition",
                    category="pair_buddy_response_fallback",
                )
                root_op.mark_terminal("degraded")
                full = "Oops, I encountered an error. Could you try again?"
                await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": full})

            _journal.append(
                JournalEntry(
                    project_id=project_id,
                    node_id=node_id,
                    event_type=JournalEventType.FEYNMAN_TURN,
                    data={"event": "init", "response": full},
                )
            )
            await _cm.send(project_id, "STUDY_BUDDY_DONE", {})

    # ---- STUDY_BUDDY_TURN ----------------------------------------
    elif event_type == "STUDY_BUDDY_TURN":
        turn_started_at = time.monotonic()
        _log_pair_buddy_turn_phase(project_id, "received", turn_started_at)
        intention_level = data.get("intention", intention)
        history = data.get("history", [])
        student_text = data.get("student_text", "")
        feynman_mode = data.get("feynman_mode", False)
        flag_confirmed = data.get("flag_confirmed", None)
        project_context_enabled = bool(data.get("project_context_enabled", True))
        knowledge_mode = "net_support" if os.getenv("TAVILY_API_KEY") else "content_only"
        _pair_buddy_interaction = append_student_interaction(
            project_id=project_id, source="study_buddy_turn", agent_id="study_buddy", text=student_text,
        )
        await capture_explicit_profile_signal(
            project_id=project_id, interaction_id=_pair_buddy_interaction.interaction_id, user_text=student_text,
        )
        
        full = ""
        with observe_operation(
            "rag.request", subsystem="rag", consumer="pair_buddy"
        ) as root_op:
            try:
                node_label = "project papers"
                context_tool_plan = plan_research_context_tools(
                    agent_id="pair_buddy",
                    query=student_text,
                    selection_text="",
                    project_context_enabled=project_context_enabled,
                    allow_web=bool(os.getenv("TAVILY_API_KEY")),
                )
                if context_tool_plan.needs_memory_clarification:
                    full = context_tool_plan.clarification_prompt
                    _log_pair_buddy_turn_phase(project_id, "context_ready", turn_started_at)
                    _log_pair_buddy_turn_phase(project_id, "first_token", turn_started_at)
                    await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": full})
                    _log_pair_buddy_turn_phase(project_id, "completed", turn_started_at)
                    await _cm.send(project_id, "STUDY_BUDDY_DONE", {})
                    return
                pair_doc_ids: list[str] = []
                chunks: list[dict[str, Any]] = []
                force_project_context = project_context_enabled and (feynman_mode or context_tool_plan.use_project_context)
                if force_project_context:
                    pair_doc_ids = await _session_all_doc_ids(project_id)
                    chunks = await _get_chunks(
                        project_id, context_tool_plan.context_query or student_text or "project papers",
                        n=5,
                        document_ids=pair_doc_ids or None,
                        selection_anchors=_request_selection_anchors(data),
                        consumer="pair_buddy",
                    )
                memory_context = ""
                if context_tool_plan.use_memory_context:
                    memory_query = context_tool_plan.memory_query or student_text or "student profile"
                    project_memory, student_memory = await _recall_memory_channels(
                        project_id, memory_query, ConsumerType.PAIR_BUDDY
                    )
                    memory_context = (
                        f"PROJECT MEMORY (Chroma; project-only):\n{project_memory or 'None'}\n\n"
                        f"STUDENT MEMORY (Cognee; cross-project):\n{student_memory or 'None'}"
                    )
                student_profile = _with_agent_control(
                    project_id, "pair_buddy",
                    memory_context,
                )
                web_context = ""
                if context_tool_plan.use_web_search and not feynman_mode:
                    web_text, _entities = await _run_optional_web_context(student_text, "")
                    if web_text:
                        web_context = (
                            "WEB RESEARCH (optional context tool result; cite linked sources for web facts):\n"
                            f"{web_text}"
                        )
                context_tools_used = []
                if force_project_context:
                    context_tools_used.append("project_context")
                if context_tool_plan.use_memory_context:
                    context_tools_used.append("memory_context")
                if web_context:
                    context_tools_used.append("web_search")
                _log_pair_buddy_turn_phase(project_id, "context_ready", turn_started_at)
                first_token_logged = False
                async for token in _get_study_buddy().evaluate_and_ask_next(
                    node_label=node_label,
                    chunks=chunks,
                    familiarity=intention_level,
                    history=history,
                    student_answer=student_text,
                    student_profile=student_profile,
                    is_merged=len(pair_doc_ids) > 1,
                    merge_summary="",
                    feynman_mode=feynman_mode,
                    flag_confirmed=flag_confirmed,
                    project_id=project_id,
                    context_tools_summary="+".join(context_tools_used),
                    web_context=web_context,
                ):
                    if not first_token_logged:
                        _log_pair_buddy_turn_phase(project_id, "first_token", turn_started_at)
                        first_token_logged = True
                    full += token
                    await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": token})
                root_op.add_count("returned_evidence_count", len(chunks))
            except Exception as exc:
                logger.exception("StudyBuddy turn failed")
                root_op.record_exception(
                    exc,
                    escaped=False,
                    stage="context_composition",
                    category="pair_buddy_response_fallback",
                )
                root_op.mark_terminal("degraded")
                full = "Oops, I encountered an error. Could you try again?"
                _log_pair_buddy_turn_phase(project_id, "first_token", turn_started_at)
                await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": full})

            _journal.append(
                JournalEntry(
                    project_id=project_id,
                    node_id=node_id,
                    event_type=JournalEventType.FEYNMAN_TURN,
                    data={"student_text": student_text, "response": full},
                )
            )
            _log_pair_buddy_turn_phase(project_id, "completed", turn_started_at)
            await _cm.send(project_id, "STUDY_BUDDY_DONE", {})

    # ---- STUDY_BUDDY_AUDIO (voice → STT → ResearchMate turn) ---------------------
    elif event_type == "STUDY_BUDDY_AUDIO":
        if _is_demo_mode():
            await _cm.send(project_id, "STUDY_BUDDY_DONE", {"error": "Voice Pair Buddy is disabled in demo mode"})
            return
        audio_b64 = data.get("audio_base64", "")
        if not audio_b64:
            await _cm.send(project_id, "STUDY_BUDDY_DONE", {"error": "no audio"})
        elif not VoiceTranscriptionService.get().is_available:
            await _cm.send(project_id, "STUDY_BUDDY_DONE", {
                "error": "STT model not available"
            })
        else:
            audio_bytes = base64.b64decode(audio_b64)
            result = transcribe_voice(
                VoiceAudioInput(
                    content=audio_bytes,
                    mime_type=str(data.get("audio_mime_type") or "audio/wav"),
                ),
                "buddy",
                project_id=project_id,
                request_id=str(data.get("request_id") or new_request_id("voice")),
            )
            text = result.transcript if result.success else ""
            if not text:
                await _cm.send(project_id, "STUDY_BUDDY_DONE", {"error": "transcription returned empty"})
            else:
                await _cm.send(project_id, "STUDY_BUDDY_TRANSCRIBED", {"text": text})
                _audio_interaction = append_student_interaction(
                    project_id=project_id, source="study_buddy_audio", agent_id="study_buddy", text=text,
                )
                await capture_explicit_profile_signal(
                    project_id=project_id, interaction_id=_audio_interaction.interaction_id, user_text=text,
                )
                intention_level = data.get("intention", intention)
                history = data.get("history", [])
                feynman_mode = data.get("feynman_mode", False)
                flag_confirmed = data.get("flag_confirmed", None)
                
                full = ""
                try:
                    sb_node = _safe_get_node(project_id, node_id)
                    node_label = data.get("node_label") or sb_node.label
                    chunks = await _get_chunks(
                        project_id,
                        node_label,
                        n=5,
                        document_ids=sb_node.document_ids or None,
                        consumer="pair_buddy",
                    )
                    recent_summaries = _recent_session_summaries.get(project_id, [])
                    student_profile = await _build_agent_memory_context(
                        project_id=project_id,
                        agent_id="pair_buddy",
                        query=f"{node_label}\n{text}",
                        paper_chunks=chunks,
                        history=history,
                        recent_summaries=recent_summaries,
                        max_chars=9000,
                    )
                    async for token in _get_study_buddy().evaluate_and_ask_next(
                        node_label=node_label,
                        chunks=chunks,
                        intention=intention_level,
                        history=history,
                        student_answer=text,
                        student_profile=student_profile,
                        is_merged=sb_node.is_merged,
                        merge_summary=sb_node.merge_summary,
                        feynman_mode=feynman_mode,
                        flag_confirmed=flag_confirmed,
                        project_id=project_id
                    ):
                        full += token
                        await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": token})
                except Exception:
                    logger.exception("StudyBuddy audio failed")
                    full = "Oops, I encountered an error processing your voice."
                    await _cm.send(project_id, "STUDY_BUDDY_TOKEN", {"token": full})
                    
                _journal.append(JournalEntry(project_id=project_id, node_id=node_id,
                    event_type=JournalEventType.FEYNMAN_TURN,
                    data={"student_text": text, "input_type": "voice", "response": full}))
                await _cm.send(project_id, "STUDY_BUDDY_DONE", {})

    # ---- TRANSCRIBE_AUDIO (pure STT -> text only, no LM turn) -------------------
    # Dictation path shared by the chat mic and the Voice AI tab: transcribe and
    # hand back the text, nothing else. (STUDY_BUDDY_AUDIO above transcribes AND
    # runs a Socratic turn -> not what dictation wants.)
    elif event_type == "TRANSCRIBE_AUDIO_CHUNK":
        target = data.get("target", "")
        request_id = str(data.get("request_id") or new_request_id("stt"))
        try:
            assembled_data = _store_stt_audio_chunk(project_id, data)
        except Exception as exc:
            logger.warning("STT chunk upload failed: %s", exc)
            await _cm.send(project_id, "TRANSCRIBED", {
                "transcript": "",
                "text": "",
                "pause_text": "",
                "pauses": [],
                "warnings": [],
                "target": target,
                "request_id": request_id,
                "success": False,
                "error": "stt_chunk_upload_failed",
            })
            return
        if assembled_data is not None:
            await handle_event(project_id, "TRANSCRIBE_AUDIO", assembled_data)

    elif event_type == "TRANSCRIBE_AUDIO":
        request_id = str(data.get("request_id") or new_request_id("stt"))
        run_id = str(data.get("run_id") or default_run_id())
        target = data.get("target", "")
        try:
            typed_target = validate_voice_target(str(target))
        except ValueError:
            await _cm.send(project_id, "TRANSCRIBED", {
                "transcript": "",
                "text": "",
                "pause_text": "",
                "pauses": [],
                "warnings": [],
                "target": target,
                "request_id": request_id,
                "success": False,
                "error": "invalid_voice_target",
            })
            return
        if _is_demo_mode():
            await _cm.send(project_id, "TRANSCRIBED", {
                "transcript": "",
                "text": "",
                "target": target,
                "request_id": request_id,
                "success": False,
                "error": "demo_mode",
            })
            return
        try:
            audio_bytes = base64.b64decode(str(data.get("audio_base64") or ""), validate=True)
        except Exception:
            audio_bytes = b""
        voice_result = transcribe_voice(
            VoiceAudioInput(
                content=audio_bytes,
                mime_type=str(data.get("audio_mime_type") or "application/octet-stream"),
            ),
            typed_target,
            project_id=project_id,
            request_id=request_id,
        )
        payload = voice_result.model_dump()
        await _cm.send(project_id, "TRANSCRIBED", payload)
        return

    elif event_type == "STT_FRONTEND_TIMING_LOG":
        request_id = str(data.get("request_id") or new_request_id("stt"))
        run_id = str(data.get("run_id") or default_run_id())
        _log_frontend_timing(request_id, run_id, data.get("frontend_timing"))
    elif event_type == "VOICE_CLIENT_EVENT":
        action = str(data.get("action") or "").strip().lower()
        try:
            client_target = validate_voice_target(str(data.get("target") or ""))
        except ValueError:
            client_target = None
        if action in {"cancel", "retry"} and client_target is not None:
            record_voice_event(f"voice_{action}", {
                "request_id": data.get("request_id"),
                "project_id": project_id,
                "target": client_target,
                "client_cancelled": bool(data.get("client_cancelled", action == "cancel")),
                "backend_cancelled": bool(data.get("backend_cancelled", False)),
                "result_discarded": bool(data.get("result_discarded", action == "cancel")),
            })
    elif event_type == "STT_ERROR_LOG":
        payload = _base_voice_payload("stt_error", str(data.get("request_id") or new_request_id("stt")), str(data.get("run_id") or default_run_id()))
        payload.update({
            "benchmark_type": "browser_pipeline",
            "error": data.get("error"),
            "context": data.get("context"),
        })
        log_jsonl(ERROR_LOG_PATH, payload)

    # ---- WIKI_DEEPDIVE_REQUEST (on-demand YouTube videos, played in-app) ----
    elif event_type == "WIKI_DEEPDIVE_REQUEST":
        target = data.get("target", "wiki")
        if _is_demo_mode():
            await _cm.send(project_id, "WIKI_DEEPDIVE_VIDEOS", {
                "term": data.get("selection_text") or data.get("node_label", node_id),
                "videos": [],
                "target": target,
            })
            return
        term = data.get("selection_text") or data.get("node_label", node_id)
        try:
            videos = await _get_infinity().find_videos(term, intention)
        except Exception as e:
            print("deep dive search error:", e)
            videos = []
        await _cm.send(project_id, "WIKI_DEEPDIVE_VIDEOS", {"term": term, "videos": videos, "target": target})
        _journal.append(JournalEntry(
            project_id=project_id, node_id=node_id,
            event_type=JournalEventType.DEEP_DIVE, data={"term": term, "count": len(videos)},
        ))
        # Auto-summarize the top video (grounds Quiz/Flashcards on it).
        if videos:
            v = videos[0]
            await _summarize_and_ingest_video(project_id, v["video_id"], term, v["title"], intention, target)

    # ---- WIKI_DEEPDIVE_SUMMARIZE (student opened a different video) ----
    elif event_type == "WIKI_DEEPDIVE_SUMMARIZE":
        target = data.get("target", "wiki")
        if _is_demo_mode():
            await _cm.send(project_id, "WIKI_DEEPDIVE_SUMMARY", {
                "term": data.get("term", ""),
                "video_id": data.get("video_id", ""),
                "summary": "YouTube deep dives are disabled in demo mode.",
                "key_points": [],
                "target": target,
            })
            return
        await _summarize_and_ingest_video(
            project_id, data.get("video_id", ""), data.get("term", ""),
            data.get("title", ""), intention, target,
        )

    # ---- COMMIT_PROJECT (Push) --------------------------------
    elif event_type == "COMMIT_PROJECT":
        if _is_demo_mode():
            await _cm.send(project_id, "EVALUATION_DONE", {"new_nodes": [], "disabled": "demo"})
            return
        document_id = data.get("document_id", "")
        evaluator = EvaluatorAgent(journal_service=_journal)
        loop = asyncio.get_event_loop()
        # Trajectory-aware evaluation: feed the student's prior cross-session memory
        # so the idea-observer judges against their history, not this session alone.
        prior_context = await _build_agent_memory_context(
            project_id=project_id,
            agent_id="brain",
            query=data.get("topic", ""),
            max_chars=9000,
        )
        evaluator_chunks = await _get_chunks(
            project_id,
            data.get("topic", "") or "project claims methods findings limitations",
            n=20,
            document_ids=await _session_all_doc_ids(project_id) or None,
            consumer="project_graph",
        )
        evaluator_source_context = "\n\n".join(
            f"[{chunk.get('evidence_id')} | {chunk.get('filename') or chunk.get('document_id')} | "
            f"pages {chunk.get('page_start')}-{chunk.get('page_end')}]\n{chunk.get('text', '')}"
            for chunk in evaluator_chunks
        )
        assessments, session_summary = await loop.run_in_executor(
            None, evaluator.evaluate_session, project_id, prior_context, evaluator_source_context
        )
        
        # 1. Idea Observer flushes directly to frontend
        for a in assessments:
            payload = a.model_dump()
            await _cm.send(project_id, "NODE_ASSESSMENT", payload)
            _local_mem.append_trajectory(document_id, payload)
        
        journal = _journal.get_session(project_id)
        current_nodes = [n.model_dump() for n in _graph_mgr.list_nodes(project_id)]

        # 2. Graph Curator evaluates missing topics SYNCHRONOUSLY
        from app.agents.graph_curator import GraphCuratorAgent
        curator = GraphCuratorAgent(journal_service=_journal)
        spawned = await loop.run_in_executor(
            None, curator.curate, project_id, current_nodes, prior_context
        )

        # 3. Materialize each curator-spawned node into the graph and stream it live,
        # distinctly colored (origin="exploration") from planned curriculum nodes ->
        # a bare NodePatch can't create a properly-attached node (no depth/document_ids/
        # parent linkage) or notify the frontend, so we build real NodeData here.
        node_lookup = {n["id"]: n for n in current_nodes}
        new_nodes: List[NodeData] = []
        for n in spawned:
            if n.node_id in node_lookup:
                continue  # id collision -> skip rather than silently overwrite
            parent = node_lookup.get(n.parent_id)
            if not parent:
                continue  # defensive; graph_curator already filters to valid parents
            node = NodeData(
                id=n.node_id, label=n.label, description=n.description,
                depth=parent["depth"] + 1, parent_id=n.parent_id, status="ongoing",
                document_ids=parent.get("document_ids", []), origin="exploration",
            )
            _ground_graph_nodes(project_id, [node])
            _graph_mgr.add_node(project_id, node)
            node_lookup[node.id] = node.model_dump()
            new_nodes.append(node)
            await _cm.send(project_id, "GRAPH_NODE_ADDED", node.model_dump())
            await _cm.send(project_id, "GRAPH_EDGE_ADDED",
                           {"source": n.parent_id, "target": n.node_id, "relationship": "related"})

        # 4. Schedule slow post-commit work. The graph and snapshot are already
        # usable; Cognee, persona, citation, and intelligence refresh in the background.
        async def _run_commit_background() -> None:
            try:
                from app.services.project_activity import ProjectActivityService

                activity = ProjectActivityService()
                activity.record(project_id, "memory_adaptation", "Reviewing student interactions since last Commit", status="running")
                await run_commit_adaptation(project_id=project_id)
                activity.record(project_id, "memory_adaptation", "Student memory and persona updated from this Commit")
            except Exception as exc:
                print("Commit adaptation background failed:", exc)
            try:
                from app.services.citation_graph import CitationGraphService
                from app.services.project_activity import ProjectActivityService

                ProjectActivityService().record(project_id, "citation_graph_refresh", "Citation map refreshing", status="running")
                citation_graph = await CitationGraphService().refresh(project_id)
                ProjectActivityService().record(project_id, "citation_graph_refresh", "Citation map refreshed")
                await _cm.send(project_id, "CITATION_GRAPH_READY", citation_graph.model_dump())
            except Exception as exc:
                print("Citation graph refresh failed:", exc)
            try:
                from app.services.project_intelligence import ProjectIntelligenceService
                from app.services.project_service import ProjectService
                from app.services.project_activity import ProjectActivityService

                project = ProjectService().load(project_id)
                if project is not None:
                    ProjectIntelligenceService().update_from_commit(project, data.get("topic", ""), session_summary)
                    ProjectActivityService().record(project_id, "commit", "Project brief updated from Commit")
            except Exception as exc:
                print("Project brief update failed:", exc)

        asyncio.create_task(_run_commit_background())
        _recent_session_summaries.setdefault(project_id, []).append(session_summary)

        # 5. Commit Project Snapshot
        from app.services.project_commit import commit_project_snapshot
        all_nodes = _graph_mgr.list_nodes(project_id)
        file_ids: List[str] = []
        for n in all_nodes:
            for d in (n.document_ids or []):
                if d not in file_ids:
                    file_ids.append(d)
        await commit_project_snapshot(
            project_id, data.get("topic", ""), intention,
            [n.model_dump() for n in all_nodes],
            data.get("content_files", []), document_id, file_ids,
        )

        await _cm.send(project_id, "EVALUATION_DONE", {"new_nodes": [n.model_dump() for n in new_nodes], "background_jobs": True})


    # ---- UPDATE_NODE_STATUS (Frontend manually completes a node) ---------
    elif event_type == "UPDATE_NODE_STATUS":
        from app.schemas.graph import NodePatch
        patch = NodePatch(node_id=node_id, status=data.get("status", "ongoing"))
        _graph_mgr.apply_node_patch(project_id, patch)

    # ---- CACHE_CLEAR (dev) -----------------------------------------------
    elif event_type == "CACHE_CLEAR":
        count = _cache.clear()
        await _cm.send(project_id, "CACHE_CLEARED", {"count": count})

    # ---- CONTEXT_CARD_REQUEST (Infinite Wiki) ----------------------------
    elif event_type == "CONTEXT_CARD_REQUEST":
        selection_text = data.get("selection_text", "")
        intention = data.get("intention", "learn")
        anchor_id = f"wiki_{hash(selection_text) & 0xFFFFFF}"
        cache_key = _cache.make_key("WIKIPEDIA_CARD", intention, anchor_id, [], selection_text)
        cached = _cache.get(cache_key)
        if cached:
            await _cm.send(project_id, "WIKI_PAGE", cached)
            await _cm.send(project_id, "WIKI_DONE", {})
        else:
            try:
                page = await _wikipedia.resolve(selection_text)
                if page is None:
                    payload = {
                        "term": selection_text,
                        "title": selection_text,
                        "extract": "",
                        "url": "",
                        "image_url": "",
                        "thumbnail_url": "",
                        "sections": [],
                        "attribution": "No matching Wikipedia page found.",
                    }
                else:
                    payload = {"term": selection_text, **page.model_dump()}
                _cache.put(cache_key, payload)
                await _cm.send(project_id, "WIKI_PAGE", payload)
            except Exception:
                logger.exception("Wikipedia Infinite Wiki card failed")
                payload = {
                    "term": selection_text,
                    "title": selection_text,
                    "extract": "",
                    "url": "",
                    "image_url": "",
                    "thumbnail_url": "",
                    "sections": [],
                    "attribution": "Wikipedia lookup failed.",
                }
                await _cm.send(project_id, "WIKI_PAGE", payload)
            finally:
                await _cm.send(project_id, "WIKI_DONE", {})

    # ---- STYLE_FEEDBACK -------------------------------------------
    elif event_type == "STYLE_FEEDBACK":
        # Style feedback is a learner preference signal. Native Cognee retrieval
        # feedback is only valid for Cognee-issued recall ids, not app graph ids.
        feedback = data.get("feedback", "")
        if feedback:
            await _memory.record_style_feedback(project_id, feedback)

    # ---- CLOSE_PROJECT --------------------------------------------
    elif event_type == "CLOSE_PROJECT":
        if _is_demo_mode():
            await _cm.send(project_id, "SESSION_COMPLETE", {"markdown": "", "patches": [], "disabled": "demo"})
            return
        evaluator = EvaluatorAgent(journal_service=_journal)
        prior_context = await _build_agent_memory_context(
            project_id=project_id,
            agent_id="brain",
            query=data.get("topic", ""),
            max_chars=9000,
        )
        close_chunks = await _get_chunks(
            project_id,
            data.get("topic", "") or "project claims methods findings limitations",
            n=20,
            document_ids=await _session_all_doc_ids(project_id) or None,
            consumer="project_graph",
        )
        close_source_context = "\n\n".join(
            f"[{chunk.get('evidence_id')} | {chunk.get('source')}]\n{chunk.get('text', '')}"
            for chunk in close_chunks
        )
        assessments, session_summary = await asyncio.get_running_loop().run_in_executor(
            None,
            evaluator.evaluate_session,
            project_id,
            prior_context,
            close_source_context,
        )

        journal = _journal.get_session(project_id)
        all_nodes = _graph_mgr.list_nodes(project_id)
        markdown = build_summary_markdown(
            topic=data.get("topic", "Study Session"),
            intention=intention,
            nodes=all_nodes,
            journal=journal,
            session_summary=session_summary,
        )
        await _cm.send(
            project_id,
            "SESSION_COMPLETE",
            {"markdown": markdown, "patches": []},
        )
        # Fire-and-forget: review any student interactions still unreviewed
        # since the last Commit (the cursor-based batching means this is a
        # no-op if nothing changed since the last successful Commit).
        asyncio.create_task(run_commit_adaptation(project_id=project_id))

        # Same Session History commit Push does -> in case a session ends without a prior Push.
        from app.services.project_commit import commit_project_snapshot
        file_ids: List[str] = []
        for n in all_nodes:
            for d in (n.document_ids or []):
                if d not in file_ids:
                    file_ids.append(d)
        asyncio.create_task(commit_project_snapshot(
            project_id, data.get("topic", ""), intention,
            [n.model_dump() for n in all_nodes],
            data.get("content_files", []), data.get("document_id", ""), file_ids,
        ))

        # Bridge the background processing gap for the current session
        _recent_session_summaries.setdefault(project_id, []).append(session_summary)