Agentic-OS/backend/app/agent/graph.py

"""LangGraph troubleshooting agent — vessel-aware, multi-device diagnostics."""
from __future__ import annotations

import asyncio
import json
import logging
from datetime import datetime, timezone
from typing import Any, TypedDict

from langchain_core.messages import HumanMessage, SystemMessage
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.agent.device_routing import select_devices_for_issue
from app.agent.json_utils import parse_json as _parse_json
from app.agent.playbooks import playbook_for
from app.agent.pfsense_profiles import run_pfsense_diagnostics
from app.agent.proxmox_profiles import enrich_proxmox_device, run_proxmox_diagnostics
from app.agent.asterisk_profiles import asterisk_container_name
from app.agent.tool_runner import invoke_connect, invoke_tool
from app.core.crypto import decrypt_secret
from app.config import settings
from app.database import SessionLocal
from app.inventory_catalog import catalog_by_key
from app.mcp_manager import get_manager
from app.models.enums import ApprovalStatus, DeviceType, TaskStatus
from app.models.inventory import Device, Vessel
from app.models.task import ApprovalRequest, Task
from app.services import balance, integrations
from app.services.diagnostic_format import (
    diagnostics_to_markdown,
    prepare_report_diagnostics,
)
from app.services.events import publish_event
from app.services.device_defaults import resolve_slot_defaults
from app.services.investigation_context import gather_investigation_context
from app.services.approval_filter import split_proposed_fixes
from app.services.troubleshooting_rules import rule_guidance_block
from app.services.llm import triage_model, triage_model_info
from app.services.llm_usage import LlmUsageTracker, merge_run_history
from app.services.model_config import get_routing_config
from app.services.model_router import plan_reasoning_tiers, reason_with_cascade
from app.services.task_runtime import (
    clear_task_cancel,
    clear_task_min_tier,
    ensure_not_cancelled,
    get_task_min_tier,
    is_task_cancel_requested,
    TaskCancelledError,
)

logger = logging.getLogger(__name__)

MAX_DIAGNOSTICS_PER_DEVICE = 6

# In-memory LLM usage collectors keyed by task_id (one agent run at a time per task)
_run_trackers: dict[int, LlmUsageTracker] = {}


class AgentState(TypedDict, total=False):
    task_id: int
    title: str
    issue: str
    details: dict | None
    vessel_id: int | None
    vessel_name: str | None
    vessel_site: str | None
    vessel_public_ip: str | None
    devices: list[dict]
    triage_result: dict
    diagnostics: list[dict]
    analysis: dict
    report: dict
    prior_context: dict | None
    investigation_context: dict | None
    matched_rule: dict | None
    follow_up: str | None
    run_number: int
    prior_runs: list[dict]


def _tracker(task_id: int) -> LlmUsageTracker:
    if task_id not in _run_trackers:
        _run_trackers[task_id] = LlmUsageTracker()
    return _run_trackers[task_id]


async def _emit(task_id: int, kind: str, message: str, payload: dict | None = None) -> None:
    await publish_event(task_id, kind, message, payload)


async def _progress(task_id: int, phase: str, message: str, **extra) -> None:
    await _emit(task_id, "progress", message, {"phase": phase, **extra})


def _device_dict(d: Device) -> dict:
    key = d.catalog_key
    entry = catalog_by_key().get(key) if key else None
    secret = decrypt_secret(d.secret_enc)
    username = d.username
    port = d.port

    if entry:
        defaults = resolve_slot_defaults(key, entry)
        if not secret:
            secret = defaults.secret
        if not username:
            username = defaults.username or entry.default_username
        if port is None and defaults.port is not None:
            port = defaults.port

    base = {
        "id": d.id,
        "name": d.name,
        "catalog_key": key,
        "device_type": d.device_type.value,
        "address": d.address,
        "port": port,
        "mcp_server": d.mcp_server,
        "username": username,
        "secret": secret,
        "asterisk_container": asterisk_container_name(key) if key else None,
    }
    if key == "proxmox":
        return enrich_proxmox_device(base)
    return base


async def _load_vessel_devices(db: AsyncSession, vessel_id: int) -> list[dict]:
    res = await db.execute(
        select(Vessel).options(selectinload(Vessel.devices)).where(Vessel.id == vessel_id)
    )
    vessel = res.scalar_one_or_none()
    if not vessel:
        return []
    devices = [_device_dict(d) for d in vessel.devices]
    docker = next((d for d in devices if d.get("catalog_key") == "docker_vm"), None)
    if docker and docker.get("secret"):
        for d in devices:
            if d.get("catalog_key") in ("asterisk_geneseasx", "asterisk") and not d.get("secret"):
                d["secret"] = docker["secret"]
            if d.get("catalog_key") in ("asterisk_geneseasx", "asterisk") and not d.get("port"):
                d["port"] = docker.get("port")
    return devices


def _prior_context_block(prior: dict | None, follow_up: str | None) -> str:
    if not prior and not follow_up:
        return ""
    lines = ["\n\n--- Previous investigation context ---"]
    if prior:
        if prior.get("summary"):
            lines.append(f"Prior summary: {prior['summary']}")
        if prior.get("root_cause"):
            lines.append(f"Prior root cause: {prior['root_cause']}")
        if prior.get("resolution"):
            lines.append(f"Prior resolution attempt: {prior['resolution']}")
        if prior.get("devices_checked"):
            lines.append(f"Already checked: {', '.join(prior['devices_checked'])}")
        if prior.get("resolved") is False:
            lines.append("Status: issue NOT resolved — continue investigating.")
    if follow_up:
        lines.append(f"Follow-up request: {follow_up}")
    return "\n".join(lines)


def _investigation_context_block(ctx: dict | None) -> str:
    """Format memory + Obsidian hits for LLM prompts."""
    if not ctx:
        return ""
    memory = ctx.get("memory_hits") or []
    obsidian = ctx.get("obsidian_hits") or []
    if not memory and not obsidian:
        return ""

    lines = ["\n\n--- Related past investigations (memory + Obsidian) ---"]
    lines.append(f"Search query: {ctx.get('query', '')}")

    if memory:
        lines.append("\nProject memory:")
        for hit in memory[:6]:
            title = hit.get("title") or "untitled"
            mtype = hit.get("memory_type") or "note"
            content = (hit.get("content") or "").replace("\n", " ")[:350]
            proj = hit.get("project_id") or "?"
            vessel = hit.get("vessel")
            vessel_tag = f" vessel={vessel}" if vessel else ""
            lines.append(f"- [{proj}/{mtype}{vessel_tag}] {title}: {content}")

    if obsidian:
        lines.append("\nObsidian notes:")
        for hit in obsidian[:5]:
            path = hit.get("path") or "?"
            title = hit.get("title") or path
            excerpt = (hit.get("excerpt") or "").replace("\n", " ")[:300]
            vessel = hit.get("vessel")
            vessel_tag = f" (vessel: {vessel})" if vessel else ""
            lines.append(f"- {path}{vessel_tag} — {title}: {excerpt}")

    lines.append(
        "Use these prior findings to avoid repeating ruled-out causes and to reuse known fixes."
    )
    return "\n".join(lines)


async def context_node(state: AgentState) -> AgentState:
    """Load related memory and Obsidian notes before triage."""
    tid = state["task_id"]
    await ensure_not_cancelled(tid)
    await _progress(tid, "context", "Searching project memory and Obsidian…")
    await _emit(tid, "context", "Searching project memory and Obsidian for related past issues...")
    try:
        ctx = await gather_investigation_context(
            state["title"],
            state["issue"],
            vessel_name=state.get("vessel_name"),
            vessel_id=state.get("vessel_id"),
        )
    except Exception as exc:  # noqa: BLE001
        logger.warning("investigation context failed: %s", exc)
        ctx = {"memory_hits": [], "obsidian_hits": [], "error": str(exc)}

    mem_n = ctx.get("memory_count", len(ctx.get("memory_hits") or []))
    obs_n = ctx.get("obsidian_count", len(ctx.get("obsidian_hits") or []))
    msg = f"Found {mem_n} memory record(s) and {obs_n} Obsidian note(s) related to this issue."
    await _emit(
        tid,
        "context",
        msg,
        {
            "query": ctx.get("query"),
            "keywords": ctx.get("keywords"),
            "memory_count": mem_n,
            "obsidian_count": obs_n,
            "memory_titles": [h.get("title") for h in (ctx.get("memory_hits") or [])[:5]],
            "obsidian_paths": [h.get("path") for h in (ctx.get("obsidian_hits") or [])[:5]],
        },
    )
    return {"investigation_context": ctx}


async def _run_device_diagnostics(
    task_id: int,
    device: dict,
    manager,
    *,
    title: str = "",
    issue: str = "",
) -> list[dict]:
    """Connect to one onboard device and run read-only playbook diagnostics."""
    from app.models.enums import DeviceType as DT

    results: list[dict] = []
    dtype = DT(device["device_type"])
    label = device.get("catalog_key") or device.get("name") or dtype.value

    if device.get("catalog_key") == "proxmox" or dtype == DT.proxmox:
        return await run_proxmox_diagnostics(
            task_id,
            device,
            manager,
            _emit,
            max_diagnostics=MAX_DIAGNOSTICS_PER_DEVICE,
            task_title=title,
            task_issue=issue,
        )

    if device.get("catalog_key") == "pfsense" or dtype == DT.pfsense:
        return await run_pfsense_diagnostics(
            task_id,
            device,
            manager,
            _emit,
            title=title,
            issue=issue,
            max_diagnostics=8,
        )

    pb = playbook_for(dtype, device.get("mcp_server"), device.get("catalog_key"))

    if not pb:
        await _emit(task_id, "diagnose", f"No playbook for {label}.")
        return results

    server = pb["mcp_server"]
    server_state = manager.servers.get(server or "")
    if not server or not server_state or server_state.status != "loaded":
        await _emit(
            task_id,
            "diagnose",
            f"MCP '{server}' not available for {label}.",
            {"device": label, "server": server},
        )
        return results

    await _emit(task_id, "connect", f"Connecting to {label} via {server}...", {"device": label})
    connect = pb.get("connect")
    if connect:
        tool, builder = connect
        await invoke_connect(task_id, label, server, tool, builder(device), manager, _emit)

    for tool, builder in pb["diagnostics"][:MAX_DIAGNOSTICS_PER_DEVICE]:
        args = builder(device)
        res, ok = await invoke_tool(
            task_id,
            label,
            server,
            tool,
            args,
            manager,
            _emit,
            task_title=title,
            task_issue=issue,
        )
        results.append(
            {
                "device": label,
                "tool": tool,
                "ok": ok,
                "output": res.get("text", ""),
            }
        )
    return results


async def triage_node(state: AgentState) -> AgentState:
    tid = state["task_id"]
    await ensure_not_cancelled(tid)
    devices = state.get("devices") or []
    run_number = state.get("run_number", 1)
    await _progress(tid, "triage", f"Local model triaging (run {run_number})…")
    await _emit(tid, "triage", f"Local model triaging (run {run_number})...")
    device_summary = [
        {"label": d.get("catalog_key") or d.get("device_type"), "name": d.get("name")} for d in devices
    ]
    prompt = (
        "You are the triage brain of a vessel troubleshooting OS. Classify the issue. "
        "Prefer local handling when possible; only recommend cloud tiers for complex cases. "
        'Respond ONLY as JSON: {"severity":"low|medium|high|critical",'
        '"category":"short label","summary":"1-2 sentences",'
        '"needs_cloud_reasoning":true|false,'
        '"recommended_tier":"local|economy|premium",'
        '"plan":["step1","step2"]}'
    )
    ctx = (
        f"Title: {state['title']}\nIssue: {state['issue']}\n"
        f"Onboard devices to check: {json.dumps(device_summary)}"
        f"{_investigation_context_block(state.get('investigation_context'))}"
        f"{rule_guidance_block(state.get('matched_rule'))}"
        f"{_prior_context_block(state.get('prior_context'), state.get('follow_up'))}"
    )
    triage_info = await triage_model_info()
    try:
        resp = await _tracker(tid).invoke(
            "triage",
            await triage_model(),
            triage_info,
            [SystemMessage(content=prompt), HumanMessage(content=ctx)],
        )
        triage = _parse_json(resp.content) or {"summary": str(resp.content)[:500]}
    except Exception as exc:  # noqa: BLE001
        logger.warning("triage failed: %s", exc)
        triage = {"severity": "medium", "needs_cloud_reasoning": True, "error": str(exc)}
    await _emit(tid, "triage", triage.get("summary", "triage complete"), triage)
    return {"triage_result": triage}


async def diagnose_node(state: AgentState) -> AgentState:
    tid = state["task_id"]
    await ensure_not_cancelled(tid)
    devices = state.get("devices") or []
    diagnostics: list[dict] = []

    if not devices:
        await _emit(tid, "diagnose", "No vessel devices matched this issue — nothing to diagnose.")
        return {"diagnostics": diagnostics}

    await _progress(
        tid,
        "diagnose",
        f"Running diagnostics on {len(devices)} device(s)…",
        device_count=len(devices),
    )
    manager = get_manager()
    for device in devices:
        await ensure_not_cancelled(tid)
        diagnostics.extend(
            await _run_device_diagnostics(
                tid,
                device,
                manager,
                title=state.get("title") or "",
                issue=state.get("issue") or "",
            )
        )
    return {"diagnostics": diagnostics}


async def reason_node(state: AgentState) -> AgentState:
    tid = state["task_id"]
    await ensure_not_cancelled(tid)
    triage = state.get("triage_result", {})
    diagnostics = state.get("diagnostics", [])
    cfg = await get_routing_config()
    min_tier = await get_task_min_tier(tid)

    await _progress(tid, "reason", "Analyzing diagnostics with LLM…")
    tiers, routing_reason = plan_reasoning_tiers(
        triage,
        diagnostics,
        prior_context=state.get("prior_context"),
        follow_up=state.get("follow_up"),
        run_number=state.get("run_number", 1),
        cfg=cfg,
        min_tier=min_tier,
    )
    await _emit(
        tid,
        "reason",
        f"Reasoning cascade: {' → '.join(tiers)} ({routing_reason})",
        {"tiers": tiers, "routing_reason": routing_reason},
    )

    diag_text = "\n\n".join(
        f"### [{d.get('device', '?')}] {d['tool']} ({'ok' if d['ok'] else 'fail'})\n{d['output'][:1500]}"
        for d in diagnostics
    ) or "No live diagnostics were collected."

    prompt = (
        "You are an expert vessel infrastructure analyst. Given the issue and diagnostic output "
        "from onboard devices, determine root cause and resolution. "
        "Only analyze devices that were actually checked — do not speculate about unrelated layers. "
        "If prior investigation context or related memory/Obsidian notes are provided, build on them — "
        "do not repeat already-ruled-out causes. Prefer fixes that worked before when still applicable. "
        "If the user asked for a list, inventory, status dump, or config details, include the actual data "
        "from diagnostics in summary and resolution — do not only describe how to find it in the UI. "
        "If you are uncertain, set needs_escalation true and confidence below 0.65 so a stronger model can retry. "
        "IMPORTANT for proposed_fixes: list ONLY optional future write actions that would modify production "
        "config (firewall rules, NAT, Asterisk dialplan, etc.). Set is_config_change true ONLY for those. "
        "Do NOT list connect steps, diagnostics already run, or manual advice — use steps_taken for those. "
        "For read-only health/status checks with no change requested, use an empty proposed_fixes array. "
        'Respond ONLY as JSON: {"summary":"...","root_cause":"... or null",'
        '"steps_taken":["..."],"proposed_fixes":[{"description":"...","is_config_change":true|false,'
        '"tool":"optional-mcp-write-tool","args":{}}],"resolved":true|false,"resolution":"...",'
        '"confidence":0.0-1.0,"needs_escalation":true|false}'
    )
    ctx = (
        f"Title: {state['title']}\nIssue: {state['issue']}\nTriage: {json.dumps(triage)}\n\n{diag_text}"
        f"{_investigation_context_block(state.get('investigation_context'))}"
        f"{rule_guidance_block(state.get('matched_rule'))}"
        f"{_prior_context_block(state.get('prior_context'), state.get('follow_up'))}"
    )
    messages = [SystemMessage(content=prompt), HumanMessage(content=ctx)]
    try:
        analysis, routing_meta = await reason_with_cascade(
            _tracker(tid), tiers, messages, _parse_json, cfg=cfg, diagnostics=diagnostics, task_id=tid
        )
        analysis["model_routing"] = {**routing_meta, "routing_reason": routing_reason}
    except Exception as exc:  # noqa: BLE001
        logger.warning("reason failed: %s", exc)
        analysis = {"summary": f"reasoning error: {exc}", "resolved": False, "proposed_fixes": []}
    await _emit(
        tid,
        "reason",
        analysis.get("summary", "analysis complete"),
        {
            "model": analysis.get("_model"),
            "tier": analysis.get("_tier_used"),
            "routing": analysis.get("model_routing"),
        },
    )
    return {"analysis": analysis}


def _unique_strings(values: list[Any]) -> list[str]:
    out: list[str] = []
    seen: set[str] = set()
    for value in values:
        text = str(value).strip()
        if not text or text in seen:
            continue
        seen.add(text)
        out.append(text)
    return out


def _scope_checked(devices: list[dict], diagnostics: list[dict]) -> list[dict]:
    out: list[dict] = []
    for device in devices:
        key = device.get("catalog_key") or device.get("name") or "device"
        tools = [
            d.get("tool")
            for d in diagnostics
            if (d.get("device") == key or d.get("device") == device.get("name")) and d.get("tool")
        ]
        out.append(
            {
                "device": key,
                "name": device.get("name") or key,
                "checks": _unique_strings(tools),
                "ok": all(bool(d.get("ok")) for d in diagnostics if d.get("device") == key),
            }
        )
    return out


def _report_findings(analysis: dict, diagnostics: list[dict], devices_checked: list[str]) -> list[dict]:
    findings: list[dict] = []
    if analysis.get("root_cause"):
        findings.append(
            {
                "title": "Root cause",
                "summary": analysis.get("root_cause"),
                "description": analysis.get("resolution"),
            }
        )
    failed = [d for d in diagnostics if d.get("ok") is False]
    for d in failed[:5]:
        findings.append(
            {
                "title": f"{d.get('device', 'device')} check failed",
                "summary": d.get("tool"),
                "description": (d.get("output") or "")[:300],
            }
        )
    if not findings and devices_checked:
        findings.append(
            {
                "title": "No critical issue found",
                "summary": f"Completed checks for {', '.join(devices_checked)}.",
                "description": analysis.get("summary") or analysis.get("resolution"),
            }
        )
    return findings


async def report_node(state: AgentState) -> AgentState:
    tid = state["task_id"]
    await ensure_not_cancelled(tid)
    analysis = state.get("analysis", {})
    triage = state.get("triage_result", {})
    diagnostics = state.get("diagnostics", [])
    run_number = state.get("run_number", 1)
    follow_up = state.get("follow_up")

    await _progress(tid, "report", "Writing report and saving notes…")
    planned_steps = _unique_strings(list(triage.get("plan") or []))
    actions_taken = _unique_strings(list(analysis.get("steps_taken") or []))
    resolved = bool(analysis.get("resolved"))
    proposed_all = analysis.get("proposed_fixes") or []
    proposed, informational = split_proposed_fixes(proposed_all)

    llm_summary = _tracker(tid).summary()

    report_diagnostics = prepare_report_diagnostics(diagnostics)
    devices_checked = [d.get("catalog_key") or d.get("name") for d in state.get("devices") or []]
    devices_checked = _unique_strings(devices_checked)
    scope_checked = _scope_checked(state.get("devices") or [], diagnostics)
    tools_run = [
        {"device": d.get("device"), "tool": d.get("tool"), "ok": d.get("ok")}
        for d in diagnostics
    ]
    findings = _report_findings(analysis, diagnostics, devices_checked)
    executive_summary = analysis.get("summary") or (
        f"Checked {', '.join(devices_checked)}." if devices_checked else "No device checks were run."
    )

    report: dict[str, Any] = {
        "summary": analysis.get("summary"),
        "executive_summary": executive_summary,
        "severity": triage.get("severity"),
        "root_cause": analysis.get("root_cause"),
        "planned_steps": planned_steps,
        "actions_taken": actions_taken,
        "tools_run": tools_run,
        # Backward compatible, no tool-name spam.
        "steps": actions_taken or planned_steps,
        "devices_checked": devices_checked,
        "scope_checked": scope_checked,
        "findings": findings,
        "diagnostics": report_diagnostics,
        "resolved": resolved,
        "resolution": analysis.get("resolution"),
        "proposed_fixes": proposed_all,
        "proposed_fixes_actionable": proposed,
        "recommendations": informational,
        "follow_up": follow_up,
        "model_routing": analysis.get("model_routing"),
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "vessel_id": state.get("vessel_id"),
        "vessel_name": state.get("vessel_name"),
        "vessel_site": state.get("vessel_site"),
        "vessel_public_ip": state.get("vessel_public_ip"),
    }
    inv_ctx = state.get("investigation_context")
    if inv_ctx:
        report["prior_knowledge"] = {
            "query": inv_ctx.get("query"),
            "memory_count": inv_ctx.get("memory_count"),
            "obsidian_count": inv_ctx.get("obsidian_count"),
            "memory_titles": [h.get("title") for h in (inv_ctx.get("memory_hits") or [])[:8]],
            "obsidian_paths": [h.get("path") for h in (inv_ctx.get("obsidian_hits") or [])[:8]],
        }
    if state.get("matched_rule"):
        report["matched_rule"] = state["matched_rule"]

    report = merge_run_history(state.get("prior_runs") or [], run_number, report, llm_summary)

    async with SessionLocal() as db:
        created: list[ApprovalRequest] = []
        for fix in proposed:
            row = ApprovalRequest(
                task_id=tid,
                tool_name=fix.get("tool") or "manual-config-change",
                tool_args=fix.get("args") or {},
                risk=fix.get("description"),
                status=ApprovalStatus.pending,
            )
            db.add(row)
            created.append(row)
        await db.flush()
        for row in created:
            await _emit(
                tid,
                "approval",
                f"Approval required: {row.tool_name} — {(row.risk or '')[:120]}",
                {
                    "approval_id": row.id,
                    "tool_name": row.tool_name,
                    "risk": row.risk,
                    "status": "pending",
                },
            )
        if created:
            res = await db.execute(select(Task).where(Task.id == tid))
            task_row = res.scalar_one_or_none()
            if task_row:
                task_row.status = TaskStatus.waiting_approval
        await db.commit()

    cost_msg = (
        f"Run cost ${report['run_cost_usd']:.4f} · Total ${report['total_cost_usd']:.4f}"
        if report["total_cost_usd"] > 0
        else f"Run cost $0.00 (local) · Total ${report['total_cost_usd']:.4f}"
    )
    await _emit(tid, "report", f"Writing memory and Obsidian note… {cost_msg}", {"llm_usage": report["llm_usage"]})

    title = state["title"]
    run_label = f" (run {run_number})" if run_number > 1 else ""
    vessel_name = state.get("vessel_name")
    vessel_label = f"[{vessel_name}] " if vessel_name else ""
    content = (
        f"Vessel: {vessel_name or 'unknown'}\n"
        f"Issue: {state['issue']}\n\n"
        f"Run: {run_number}\n"
        f"Devices checked: {', '.join(report['devices_checked']) or 'none'}\n\n"
        f"Root cause: {report['root_cause']}\n\nResolution: {report['resolution']}\n\n"
        f"LLM cost this run: ${report['run_cost_usd']:.4f} (total ${report['total_cost_usd']:.4f})\n\n"
        f"Findings:\n" + "\n".join(f"- {f.get('title')}: {f.get('summary')}" for f in findings)
    )
    if report_diagnostics:
        content += "\n\nDiagnostic tools run:\n" + "\n".join(
            f"- [{d.get('device')}] {d.get('tool')}" for d in report_diagnostics
        )
    memory_id = await integrations.create_task_memory(
        title=f"{vessel_label}[Task] {title}{run_label}",
        content=content,
        memory_type="fix" if resolved else "note",
        tags=["agentic-os", triage.get("category", "troubleshooting")],
        importance=4 if triage.get("severity") in ("high", "critical") else 3,
        vessel_name=vessel_name,
        task_id=tid,
    )
    obsidian_path, obsidian_push_error = await integrations.write_obsidian_note(
        title=f"{vessel_label}{title}{run_label}".strip(),
        issue=state["issue"] if not follow_up else f"{state['issue']}\n\nFollow-up: {follow_up}",
        steps=actions_taken or planned_steps,
        cause=report["root_cause"],
        resolution=report["resolution"],
        resolved=resolved,
        vessel_name=vessel_name,
        task_id=tid,
        diagnostics_md=diagnostics_to_markdown(report_diagnostics),
    )

    report["memory_id"] = memory_id
    report["obsidian_path"] = obsidian_path
    report["obsidian_push_error"] = obsidian_push_error
    report["memory_project_id"] = settings.memory_task_project_id
    complete_msg = "Report complete."
    if memory_id:
        complete_msg += f" Memory saved ({settings.memory_task_project_id})."
    if obsidian_path and not obsidian_push_error:
        complete_msg += f" Obsidian note pushed ({obsidian_path})."
    elif obsidian_path:
        complete_msg += f" Obsidian note saved locally ({obsidian_path}); git push failed."
    await _emit(
        tid,
        "report",
        complete_msg,
        {
            "memory_id": memory_id,
            "memory_project_id": settings.memory_task_project_id,
            "obsidian_path": obsidian_path,
            "obsidian_push_error": obsidian_push_error,
            "llm_usage": report["llm_usage"],
            "run_cost_usd": report["run_cost_usd"],
            "total_cost_usd": report["total_cost_usd"],
        },
    )
    return {"report": report}


_compiled_graph = None


def get_graph():
    """Return the process-wide compiled LangGraph, building it once."""
    global _compiled_graph
    if _compiled_graph is None:
        _compiled_graph = build_graph()
    return _compiled_graph


def build_graph():
    from langgraph.graph import END, START, StateGraph

    g = StateGraph(AgentState)
    g.add_node("context_step", context_node)
    g.add_node("triage_step", triage_node)
    g.add_node("diagnose_step", diagnose_node)
    g.add_node("reason_step", reason_node)
    g.add_node("report_step", report_node)
    g.add_edge(START, "context_step")
    g.add_edge("context_step", "triage_step")
    g.add_edge("triage_step", "diagnose_step")
    g.add_edge("diagnose_step", "reason_step")
    g.add_edge("reason_step", "report_step")
    g.add_edge("report_step", END)
    return g.compile()


async def _finalize_cancelled(task_id: int) -> None:
    already_cancelled = False
    async with SessionLocal() as db:
        res = await db.execute(select(Task).where(Task.id == task_id))
        task = res.scalar_one_or_none()
        if task:
            already_cancelled = task.status == TaskStatus.cancelled
            if not already_cancelled:
                task.status = TaskStatus.cancelled
                await db.commit()
    if not already_cancelled:
        await _emit(task_id, "cancelled", "Task stopped.", persist=True)
    await clear_task_cancel(task_id)


async def run_task_agent(task_id: int) -> None:
    async with SessionLocal() as db:
        res = await db.execute(select(Task).where(Task.id == task_id))
        task = res.scalar_one_or_none()
        if not task:
            return

        details = dict(task.details or {})
        continue_ctx = details.pop("_continue", None)
        continue_mode = continue_ctx is not None

        prior_runs: list[dict] = details.get("runs") or []
        if not prior_runs and task.report:
            prior_runs = task.report.get("runs") or []

        if continue_mode:
            run_number = len(prior_runs) + 1
            follow_up = continue_ctx.get("follow_up", "")
            prior_context = continue_ctx.get("prior_report") or task.report
        else:
            run_number = 1
            follow_up = None
            prior_context = None
            prior_runs = []

        issue_text = task.issue
        routing_text = f"{issue_text}\n{follow_up}" if follow_up else issue_text

        devices: list[dict] = []
        matched_rule: dict | None = None
        vessel_name: str | None = None
        vessel_site: str | None = None
        vessel_public_ip: str | None = None
        if task.vessel_id:
            all_devices = await _load_vessel_devices(db, task.vessel_id)
            devices, rule = select_devices_for_issue(
                all_devices, routing_text, title=task.title
            )
            matched_rule = rule.as_dict() if rule else None
            vres = await db.execute(select(Vessel).where(Vessel.id == task.vessel_id))
            vessel = vres.scalar_one_or_none()
            if vessel:
                vessel_name = vessel.name
                vessel_site = vessel.site
                vessel_public_ip = vessel.public_ip
        elif task.device_id:
            dres = await db.execute(select(Device).where(Device.id == task.device_id))
            d = dres.scalar_one_or_none()
            if d:
                devices = [_device_dict(d)]

        task.status = TaskStatus.running
        task.details = details
        await db.commit()

    if await is_task_cancel_requested(task_id):
        await _finalize_cancelled(task_id)
        _run_trackers.pop(task_id, None)
        return

    labels = [d.get("catalog_key") or d.get("name") for d in devices]
    run_label = f" (run {run_number})" if run_number > 1 else ""
    await _emit(task_id, "log", f"Agent run started{run_label}.")
    if follow_up:
        await _emit(task_id, "log", f"Follow-up: {follow_up}")
    if vessel_name:
        await _emit(
            task_id,
            "log",
            f"Vessel: {vessel_name}"
            + (f" ({vessel_public_ip})" if vessel_public_ip else ""),
            {"vessel_name": vessel_name, "vessel_public_ip": vessel_public_ip},
        )
    await _emit(
        task_id,
        "log",
        f"Will diagnose {len(devices)} device(s): {', '.join(labels) or 'none'}",
        {"devices": labels, "run_number": run_number},
    )
    if matched_rule:
        await _emit(
            task_id,
            "rule",
            f"Matched rule: {matched_rule.get('name')} → {', '.join(matched_rule.get('devices') or [])}",
            matched_rule,
        )

    remaining = await balance.latest_remaining_usd()
    from app.config import settings

    if remaining is not None and remaining < settings.min_balance_usd:
        msg = f"Aborting: remaining balance ${remaining:.2f} below minimum ${settings.min_balance_usd:.2f}."
        await _emit(task_id, "error", msg)
        async with SessionLocal() as db:
            res = await db.execute(select(Task).where(Task.id == task_id))
            task = res.scalar_one()
            task.status = TaskStatus.failed
            task.error = msg
            await db.commit()
        _run_trackers.pop(task_id, None)
        return

    initial: AgentState = {
        "task_id": task_id,
        "title": task.title,
        "issue": issue_text,
        "details": task.details,
        "vessel_id": task.vessel_id,
        "vessel_name": vessel_name,
        "vessel_site": vessel_site,
        "vessel_public_ip": vessel_public_ip,
        "devices": devices,
        "prior_context": prior_context,
        "follow_up": follow_up,
        "run_number": run_number,
        "prior_runs": prior_runs if continue_mode else [],
        "matched_rule": matched_rule,
    }

    try:
        # Create the per-task usage tracker inside the try so the finally below
        # always cleans it up, even if setup above this point raised.
        _run_trackers[task_id] = LlmUsageTracker()
        graph = get_graph()
        final_state = await graph.ainvoke(initial)
        report = final_state.get("report", {})
        async with SessionLocal() as db:
            res = await db.execute(select(Task).where(Task.id == task_id))
            task = res.scalar_one()
            task.report = report
            task.memory_id = report.get("memory_id")
            task.obsidian_path = report.get("obsidian_path")
            task.details = {
                **(task.details or {}),
                "runs": report.get("runs", []),
                "total_cost_usd": report.get("total_cost_usd", 0.0),
            }
            pend = await db.execute(
                select(ApprovalRequest).where(
                    ApprovalRequest.task_id == task_id,
                    ApprovalRequest.status == ApprovalStatus.pending,
                )
            )
            has_pending = pend.scalars().first() is not None
            if report.get("resolved") and not has_pending:
                task.status = TaskStatus.succeeded
            elif has_pending:
                task.status = TaskStatus.waiting_approval
            elif continue_mode or not report.get("resolved"):
                task.status = TaskStatus.succeeded  # completed run but unresolved — user can continue
            else:
                task.status = TaskStatus.succeeded
            await db.commit()
            final_status = task.status.value
        await _emit(
            task_id,
            "log",
            f"Agent run finished ({final_status}). Cost: ${report.get('run_cost_usd', 0):.4f} this run, "
            f"${report.get('total_cost_usd', 0):.4f} total.",
            {"llm_usage": report.get("llm_usage"), "run_number": run_number},
        )
    except TaskCancelledError:
        logger.info("Task %s cancelled by user", task_id)
        await _finalize_cancelled(task_id)
    except asyncio.CancelledError:
        logger.info("Task %s asyncio cancelled", task_id)
        await _finalize_cancelled(task_id)
    except Exception as exc:  # noqa: BLE001
        logger.exception("Agent run failed for task %s", task_id)
        await _emit(task_id, "error", f"Agent run failed: {exc}")
        async with SessionLocal() as db:
            res = await db.execute(select(Task).where(Task.id == task_id))
            task = res.scalar_one()
            task.status = TaskStatus.failed
            task.error = str(exc)
            await db.commit()
    finally:
        _run_trackers.pop(task_id, None)
        try:
            await clear_task_min_tier(task_id)
        except Exception:  # noqa: BLE001
            logger.debug("clear_task_min_tier failed for task %s", task_id)