admin / Synapse-Cortex

"""AI investigator: selects the most appropriate *enabled* Playbook for a
ticket by calling Claude with a tool whose `playbook_id` parameter is
structurally constrained to an enum of only the tenant's currently-enabled
playbook ids (plus a NONE sentinel). This prevents the model from selecting
a disabled or hallucinated playbook - a real correctness/security control,
not just a prompt hint.

The model is given ticket/asset context and per-playbook metadata only
(name, allowed target OS, approval level) - never the raw graph_json and
never credential data. It never writes commands itself; the caller extracts
the actual command list from the selected playbook's snapshot server-side
via app/ai/guardrails.parse_graph_to_commands.

get_anthropic_client() constructs the client lazily so a missing API key
surfaces as a clear error the first time an investigation is actually
attempted, not at import/app-boot time. It is exposed as a FastAPI
dependency (Depends(get_db)/Depends(get_session_user)) so tests can
override it via app.dependency_overrides with a fake client, mirroring the
get_executor() pattern in app/ai/executor.py.

It prefers the tenant's own Claude API key, saved via Admin -> AI Settings
and encrypted with the tenant's vault key, and falls back to the
ANTHROPIC_API_KEY environment variable when no tenant key is configured -
this keeps existing env-var-only deployments working unchanged.

It is used as a route's Depends(get_anthropic_client), so the missing-key
case raises HTTPException directly rather than InvestigatorError: FastAPI
only turns HTTPException into a clean JSON response when raised inside a
dependency - any other exception there becomes an unhandled 500 with no
detail, since the route body's own try/except never gets a chance to run.
"""
import os
import uuid
from dataclasses import dataclass
from typing import Optional

from fastapi import Depends, HTTPException, status
from sqlalchemy.orm import Session

from .. import models
from ..database import get_db
from ..deps import get_session_user
from .vault import VaultError, decrypt_secret

MODEL_ID = "claude-opus-4-8"

NONE_SENTINEL = "NONE"


class InvestigatorError(RuntimeError):
    pass


def get_anthropic_client(
    db: Session = Depends(get_db),
    user: Optional[models.User] = Depends(get_session_user),
):
    if not user:
        raise HTTPException(status.HTTP_401_UNAUTHORIZED)

    tenant = db.get(models.Tenant, user.tenant_id)
    api_key = None
    if tenant.anthropic_api_key_encrypted:
        try:
            api_key = decrypt_secret(tenant, tenant.anthropic_api_key_encrypted)
        except VaultError:
            api_key = None
    if not api_key:
        api_key = os.getenv("ANTHROPIC_API_KEY")

    if not api_key:
        raise HTTPException(
            status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=(
                "No Claude API key is configured. Save one from Admin → AI Settings, or set "
                "ANTHROPIC_API_KEY in the environment, to enable AI-driven playbook selection."
            ),
        )
    import anthropic

    try:
        return anthropic.Anthropic(api_key=api_key)
    except Exception as exc:
        raise HTTPException(status.HTTP_502_BAD_GATEWAY, detail=f"Could not initialize the Claude API client: {exc}")


def test_api_key(api_key: str) -> tuple[bool, str]:
    """Makes a single, minimal (max_tokens=1) call to verify a Claude API
    key is valid. Never touches a tenant's stored key or the
    ANTHROPIC_API_KEY environment variable - the caller passes whichever
    key it wants validated. Client construction is wrapped in the same
    try/except as the request itself - it can also raise (e.g. an
    incompatible installed SDK/httpx pairing), and this is a user-facing
    "is this key good" check, so nothing here should ever surface as an
    unhandled 500."""
    import anthropic

    try:
        client = anthropic.Anthropic(api_key=api_key)
        client.messages.create(
            model=MODEL_ID,
            max_tokens=1,
            messages=[{"role": "user", "content": "ping"}],
        )
    except anthropic.AuthenticationError as exc:
        return False, f"Authentication failed: {exc}"
    except Exception as exc:  # anthropic.APIError and subclasses, network/client-construction errors
        return False, f"Claude API request failed: {exc}"
    return True, "Claude API key is valid."


@dataclass
class InvestigationResult:
    selected_playbook_id: Optional[uuid.UUID]
    rationale: str
    # Only set when the selected playbook is a command runner (freeform=true):
    # the exact single shell command the model proposes to run. It is NOT
    # executed as-is - a human reviews (and may edit) it, and it passes the
    # full guardrail engine (including the non-overridable forbidden-command
    # baseline) before anything runs. See app/routers/remediation.py.
    suggested_command: Optional[str] = None


def _build_tool(playbook_ids: list[str]) -> dict:
    return {
        "name": "select_playbook",
        "description": (
            "Select the single most appropriate remediation playbook for this ticket "
            "from the enabled playbooks provided, or NONE if none of them apply. "
            "If (and only if) the selected playbook has \"freeform\": true, also provide "
            "suggested_command with the exact single shell command that should be run to "
            "resolve the ticket - a human will review and may edit it before it executes."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "playbook_id": {
                    "type": "string",
                    "enum": playbook_ids + [NONE_SENTINEL],
                    "description": "The id of the selected playbook, or the literal string NONE.",
                },
                "rationale": {
                    "type": "string",
                    "description": "One or two sentences explaining why this playbook (or NONE) was selected.",
                },
                "suggested_command": {
                    "type": "string",
                    "description": (
                        "Only when the selected playbook has freeform=true: the exact single shell "
                        "command to run. Omit for any normal (fixed-command) playbook."
                    ),
                },
            },
            "required": ["playbook_id", "rationale"],
        },
    }


def investigate(
    *,
    client,
    ticket_context: dict,
    asset_context: Optional[dict],
    playbooks: list[dict],
) -> InvestigationResult:
    """`playbooks` must already be filtered to the tenant's currently
    *enabled* playbooks - the tool's enum is built directly from whatever is
    passed in here, so an un-filtered list would let the model select a
    disabled playbook. Each entry: {id, name, allowed_target_os,
    required_approval_level}."""
    if not playbooks:
        return InvestigationResult(
            selected_playbook_id=None, rationale="No enabled playbooks are configured for this tenant."
        )

    playbook_ids = [str(p["id"]) for p in playbooks]
    tool = _build_tool(playbook_ids)

    playbook_summaries = [
        {
            "id": str(p["id"]),
            "name": p["name"],
            "allowed_target_os": p["allowed_target_os"],
            "required_approval_level": p["required_approval_level"],
            "freeform": bool(p.get("freeform")),
        }
        for p in playbooks
    ]

    has_client_context = bool(asset_context and asset_context.get("client_context"))
    client_context_guidance = (
        "The linked asset includes a \"client_context\" note written by an operator - this is "
        "authoritative context about the client, the asset's operational role, and any handling "
        "constraints (e.g. maintenance windows, do-not-reboot rules, SLA sensitivity). Read it "
        "carefully and let it govern your choice: prefer a safer or read-only playbook, or select "
        "NONE (deferring to a human), if the client_context indicates the available fixes could be "
        "disruptive or unsafe for this particular asset. Reflect how it shaped your decision in the "
        "rationale.\n\n"
        if has_client_context
        else ""
    )

    prompt = (
        "You are triaging an IT support ticket to select the correct automated remediation "
        "playbook. Consider the ticket details, the linked asset (if any), and each playbook's "
        "allowed target OS. Call select_playbook with your choice.\n\n"
        f"{client_context_guidance}"
        f"Ticket:\n{ticket_context}\n\n"
        f"Asset:\n{asset_context or 'No asset linked to this ticket.'}\n\n"
        f"Enabled playbooks:\n{playbook_summaries}"
    )

    try:
        response = client.messages.create(
            model=MODEL_ID,
            max_tokens=1024,
            tools=[tool],
            tool_choice={"type": "tool", "name": "select_playbook"},
            messages=[{"role": "user", "content": prompt}],
        )
    except InvestigatorError:
        raise
    except Exception as exc:  # anthropic.APIError and subclasses
        raise InvestigatorError(f"AI investigation request failed: {exc}") from exc

    tool_use = next((block for block in response.content if block.type == "tool_use"), None)
    if tool_use is None:
        raise InvestigatorError("AI investigation did not return a playbook selection.")

    selection = tool_use.input.get("playbook_id")
    rationale = tool_use.input.get("rationale") or ""
    suggested_command = tool_use.input.get("suggested_command") or None

    if selection == NONE_SENTINEL or selection not in playbook_ids:
        return InvestigationResult(selected_playbook_id=None, rationale=rationale or "No suitable playbook was found.")

    return InvestigationResult(
        selected_playbook_id=uuid.UUID(selection), rationale=rationale, suggested_command=suggested_command
    )