"""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.skill_registry import match_skills, skill_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 matched_skills: list[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"{skill_guidance_block(state.get('matched_skills'))}" 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"{skill_guidance_block(state.get('matched_skills'))}" 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"] if state.get("matched_skills"): report["matched_skills"] = [ {k: v for k, v in s.items() if k != "body"} for s in state["matched_skills"] ] 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 matched_skills: list[dict] = [] 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 rule_id = matched_rule.get("id") if matched_rule else None matched_skills = [ { "id": s.id, "name": s.name, "description": s.description, "body": s.body, "score": s.score, "priority": s.priority, } for s in match_skills(task.title, routing_text, matched_rule_id=rule_id) ] 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)] if not matched_skills: rule_id = matched_rule.get("id") if matched_rule else None matched_skills = [ { "id": s.id, "name": s.name, "description": s.description, "body": s.body, "score": s.score, "priority": s.priority, } for s in match_skills(task.title, routing_text, matched_rule_id=rule_id) ] 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, ) if matched_skills: names = ", ".join(s["name"] for s in matched_skills) await _emit( task_id, "skill", f"Matched skills: {names}", {"skills": [{k: v for k, v in s.items() if k != "body"} for s in matched_skills]}, ) 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, "matched_skills": matched_skills or None, } 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)