backend/agents.py

"""
The swarm. Four specialist roles built on Gemma-4-on-Cerebras:

  Analyst    (fan-out, per finding)  -> category, CVSS, exploitability, FP-likelihood
  Remediator (conditional, per High/Critical) -> fix + minimal patch hint
  Commander  (once, over the batch)  -> exec rollup + prioritized actions
  Vision RCA (multimodal, on demand) -> root-cause + patch from a screenshot

The speed thesis lives in `run_batch`: analysts fan out concurrently (bounded by a
semaphore) so the wall-clock is ~one slow call, not the sum of all calls. The
sequential baseline runs the identical work one-at-a-time — same model, honest delta.
"""
from __future__ import annotations

import asyncio
import time
from typing import Any, AsyncIterator, Optional

import httpx

# ---- strict JSON schemas (constrain Gemma 4 to clean, parseable output) ----

ANALYST_SCHEMA = {
    "type": "object", "additionalProperties": False,
    "properties": {
        "category": {"type": "string"},
        "severity_cvss": {"type": "number"},
        "severity_label": {"type": "string", "enum": ["critical", "high", "medium", "low", "info"]},
        "exploitability": {"type": "string", "enum": ["trivial", "plausible", "theoretical", "none"]},
        "false_positive_likelihood": {"type": "number"},
        "rationale": {"type": "string"},
        "affected_asset": {"type": "string"},
    },
    "required": ["category", "severity_cvss", "severity_label", "exploitability",
                 "false_positive_likelihood", "rationale", "affected_asset"],
}

REMEDIATOR_SCHEMA = {
    "type": "object", "additionalProperties": False,
    "properties": {
        "fix_summary": {"type": "string"},
        "patch_hint": {"type": "string"},
        "effort": {"type": "string", "enum": ["S", "M", "L"]},
    },
    "required": ["fix_summary", "patch_hint", "effort"],
}

COMMANDER_SCHEMA = {
    "type": "object", "additionalProperties": False,
    "properties": {
        "headline": {"type": "string"},
        "risk_paragraph": {"type": "string"},
        "top_actions": {"type": "array", "items": {"type": "string"}},
        "counts": {
            "type": "object", "additionalProperties": False,
            "properties": {k: {"type": "integer"} for k in
                           ["critical", "high", "medium", "low", "info"]},
            "required": ["critical", "high", "medium", "low", "info"],
        },
    },
    "required": ["headline", "risk_paragraph", "top_actions", "counts"],
}

VISION_SCHEMA = {
    "type": "object", "additionalProperties": False,
    "properties": {
        "root_cause": {"type": "string"},
        "evidence_from_image": {"type": "string"},
        "concrete_patch": {"type": "string"},
        "confidence": {"type": "number"},
    },
    "required": ["root_cause", "evidence_from_image", "concrete_patch", "confidence"],
}

ANALYST_SYS = (
    "You are a senior application-security analyst triaging one scanner finding. "
    "Judge real risk, not the scanner's raw label. Estimate a CVSS-style 0-10 score, "
    "whether it's actually exploitable, and how likely it's a false positive. "
    "Be terse and specific. Output only the requested JSON."
)
REMEDIATOR_SYS = (
    "You are a remediation engineer. Given a triaged finding, give the smallest correct "
    "fix and a one-line patch hint a developer can act on immediately. Output only JSON."
)
COMMANDER_SYS = (
    "You are the incident commander. Given the triaged batch, write a crisp executive "
    "rollup: one headline, one risk paragraph, and the top prioritized actions. "
    "Count findings by severity. Output only JSON."
)
VISION_SYS = (
    "You are a security analyst doing root-cause analysis from a screenshot of code, a "
    "dashboard, or an alert. Identify the actual root cause, cite what you see in the image, "
    "and give a concrete patch. Output only JSON."
)


def _analyst_user(f: dict) -> str:
    return (f"scanner={f.get('scanner','?')} rule={f.get('rule','?')} "
            f"raw_severity={f.get('raw_severity','?')}\n"
            f"path={f.get('path','')}\n"
            f"message: {f.get('message','')}\n"
            f"snippet: {f.get('snippet','')[:400]}")


class Swarm:
    def __init__(self, client, *, max_concurrency: int = 10, remediate_min_cvss: float = 7.0):
        self.client = client
        self.sem = asyncio.Semaphore(max_concurrency)
        self.remediate_min_cvss = remediate_min_cvss

    # ---- single agents ----
    async def analyst(self, http: httpx.AsyncClient, f: dict) -> dict:
        r = await self.client.chat(http, system=ANALYST_SYS, user=_analyst_user(f),
                                   json_schema=ANALYST_SCHEMA, max_tokens=350)
        return {"finding": f, "analysis": r.json or {}, "timing": r.timing}

    async def remediator(self, http: httpx.AsyncClient, f: dict, analysis: dict) -> Optional[dict]:
        a = (f"finding: {f.get('message','')}\ncategory: {analysis.get('category')}\n"
             f"cvss: {analysis.get('severity_cvss')}\nsnippet: {f.get('snippet','')[:300]}")
        r = await self.client.chat(http, system=REMEDIATOR_SYS, user=a,
                                   json_schema=REMEDIATOR_SCHEMA, max_tokens=250)
        return r.json

    async def vision_rca(self, http: httpx.AsyncClient, *, context: str,
                         image_b64: str, media_type: str) -> dict:
        r = await self.client.chat(http, system=VISION_SYS, user=context,
                                   image_b64=image_b64, media_type=media_type,
                                   json_schema=VISION_SCHEMA, max_tokens=400)
        return {"rca": r.json or {}, "timing": r.timing}

    async def commander(self, http: httpx.AsyncClient, triaged: list[dict]) -> dict:
        lines = []
        for t in triaged[:60]:
            a = t["analysis"]
            lines.append(f"- [{a.get('severity_label','?')}/{a.get('severity_cvss','?')}] "
                         f"{a.get('category','?')}: {t['finding'].get('message','')[:90]}")
        r = await self.client.chat(http, system=COMMANDER_SYS,
                                   user="Triaged findings:\n" + "\n".join(lines),
                                   json_schema=COMMANDER_SCHEMA, max_tokens=500)
        return {"rollup": r.json or {}, "timing": r.timing}

    # ---- one finding, full pipeline (analyst -> conditional remediator) ----
    async def _one(self, http: httpx.AsyncClient, f: dict) -> dict:
        res = await self.analyst(http, f)
        a = res["analysis"]
        if float(a.get("severity_cvss") or 0) >= self.remediate_min_cvss:
            res["remediation"] = await self.remediator(http, f, a)
        return res

    # ---- batch: stream results as they complete; bounded concurrency ----
    async def run_batch(self, findings: list[dict], *, parallel: bool) -> AsyncIterator[dict]:
        t0 = time.perf_counter()
        triaged: list[dict] = []
        calls = 0
        ttfts: list[float] = []
        comp_tokens = 0

        async with httpx.AsyncClient() as http:
            async def guarded(f):
                async with self.sem:
                    return await self._one(http, f)

            if parallel:
                tasks = [asyncio.create_task(guarded(f)) for f in findings]
                for fut in asyncio.as_completed(tasks):
                    res = await fut
                    triaged.append(res)
                    calls += 1 + (1 if "remediation" in res else 0)
                    if res["timing"].ttft_ms:
                        ttfts.append(res["timing"].ttft_ms)
                    comp_tokens += res["timing"].completion_tokens
                    yield {"type": "finding", "data": _wire(res),
                           "elapsed_ms": (time.perf_counter() - t0) * 1000}
            else:
                for f in findings:
                    res = await self._one(http, f)
                    triaged.append(res)
                    calls += 1 + (1 if "remediation" in res else 0)
                    if res["timing"].ttft_ms:
                        ttfts.append(res["timing"].ttft_ms)
                    comp_tokens += res["timing"].completion_tokens
                    yield {"type": "finding", "data": _wire(res),
                           "elapsed_ms": (time.perf_counter() - t0) * 1000}

            roll = await self.commander(http, triaged)
            calls += 1
            yield {"type": "rollup", "data": roll["rollup"],
                   "elapsed_ms": (time.perf_counter() - t0) * 1000}

        wall = (time.perf_counter() - t0) * 1000
        yield {"type": "summary", "data": {
            "wall_ms": round(wall, 1),
            "calls": calls,
            "findings": len(findings),
            "avg_ttft_ms": round(sum(ttfts) / len(ttfts), 1) if ttfts else None,
            "completion_tokens": comp_tokens,
            "throughput_tps": round(comp_tokens / (wall / 1000), 1) if wall else 0,
            "mode": "parallel" if parallel else "sequential",
        }}


def _wire(res: dict) -> dict:
    """Strip non-serializable timing object down to what the UI needs."""
    a = res["analysis"]
    return {
        "id": res["finding"]["id"],
        "rule": res["finding"].get("rule", ""),
        "scanner": res["finding"].get("scanner", ""),
        "message": res["finding"].get("message", ""),
        "path": res["finding"].get("path", ""),
        "category": a.get("category", ""),
        "cvss": a.get("severity_cvss", 0),
        "label": a.get("severity_label", "info"),
        "exploitability": a.get("exploitability", ""),
        "fp": a.get("false_positive_likelihood", 0),
        "rationale": a.get("rationale", ""),
        "remediation": res.get("remediation"),
        "ttft_ms": round(res["timing"].ttft_ms, 1) if res["timing"].ttft_ms else None,
    }