invoice-guard-studio-v1/agent.py

"""InvoiceGuard Studio v1.

A deterministic full-stack A2A product for reviewing pasted or uploaded
invoices. The agent does not call an LLM: it validates invoice records,
detects duplicate invoice numbers and stated-total mismatches, persists
user-scoped review cases/decisions/receipts in managed Postgres, and exposes
small typed tools that are also served through MCP.
"""
from __future__ import annotations

import base64
import asyncio
import csv
import hashlib
import io
import json
import os
import re
from contextlib import contextmanager
from datetime import UTC, datetime
from decimal import Decimal, InvalidOperation, ROUND_HALF_UP
from typing import Annotated, Any

from pydantic import BaseModel, ConfigDict, Field, ValidationError

import a2a_pack as a2a
from a2a_pack import (
    A2AAgent,
    AgentDatabase,
    AgentDatabaseEnv,
    AgentDatabaseMigrations,
    AgentPlatformResources,
    FileUpload,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    UploadedFile,
    WorkspaceAccess,
    WorkspaceMode,
)

MAX_INVOICES = 50
MAX_UPLOADS = 5
MAX_UPLOAD_BYTES = 256 * 1024
MAX_TEXT_FIELD_LENGTH = 180
AGENT_VERSION = "0.1.3"
ACCEPTED_MEDIA_TYPES = {"application/json", "text/json", "text/csv", "text/plain"}
DB_CONNECT_OPTIONS = "-c statement_timeout=5000 -c lock_timeout=3000 -c idle_in_transaction_session_timeout=10000"


class InvoiceGuardStudioV1Config(BaseModel):
    """No operator configuration is required for this deterministic product."""


class BrowserDocument(BaseModel):
    filename: str = Field(min_length=1, max_length=180)
    media_type: str = Field(min_length=1, max_length=120)
    data_base64: str = Field(min_length=1, max_length=((MAX_UPLOAD_BYTES * 4) // 3) + 16)


class InvoiceInput(BaseModel):
    model_config = ConfigDict(extra="allow")

    invoice_number: str = Field(min_length=1, max_length=120)
    vendor: str = Field(min_length=1, max_length=MAX_TEXT_FIELD_LENGTH)
    subtotal: Any
    tax: Any
    total: Any


class DecisionInput(BaseModel):
    decision_id: str = Field(min_length=1, max_length=120)
    invoice_number: str | None = Field(default=None, max_length=120)
    decision: str = Field(min_length=1, max_length=40)
    note: str = Field(default="", max_length=1000)


class InvoiceGuardStudioV1(A2AAgent[InvoiceGuardStudioV1Config, PlatformUserAuth]):
    name = "invoice-guard-studio-v1"
    description = (
        "One-page InvoiceGuard app for deterministic duplicate invoice number "
        "and stated-total mismatch review with user-scoped Postgres persistence."
    )
    version = AGENT_VERSION

    config_model = InvoiceGuardStudioV1Config
    auth_model = PlatformUserAuth

    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic invoice validation; no LLM credentials are required.",
    )
    workspace_access = WorkspaceAccess.dynamic(
        max_files=8,
        allowed_modes=(WorkspaceMode.READ_ONLY,),
        require_reason=False,
        max_total_size_bytes=MAX_UPLOAD_BYTES * MAX_UPLOADS,
    )
    tools_used = ("managed-postgres", "mcp")
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="invoice-guard-studio-v1-data",
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )

    @a2a.tool(
        description="Review pasted invoice objects, detect duplicate invoice numbers and stated-total mismatches, persist a user-scoped case, and record a production receipt.",
        timeout_seconds=30,
        idempotent=False,
        cost_class="deterministic",
    )
    async def review_invoices(
        self,
        ctx: RunContext[PlatformUserAuth],
        case_id: Annotated[
            str,
            Field(min_length=1, max_length=120, pattern=r"^[A-Za-z0-9_.:-]+$"),
        ],
        invoices: Annotated[list[InvoiceInput], Field(min_length=1, max_length=MAX_INVOICES)],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        result = _review_payload(case_id, invoices)
        if not result["ok"]:
            await ctx.emit_error(result["message"], code=result["code"])
            return result
        persisted = await asyncio.to_thread(
            _persist_case_with_receipt,
            tenant,
            result,
            "json",
        )
        await ctx.emit_progress(
            f"reviewed {result['invoice_count']} invoice(s): {result['duplicate_count']} duplicate group(s), {result['total_mismatch_count']} total mismatch(es)"
        )
        return {**result, "receipt": persisted["receipt"], "saved_at": persisted["saved_at"]}

    @a2a.tool(
        description="Review invoices uploaded by browser JSON/base64 bridge, then persist the same review case and receipt as pasted invoices.",
        timeout_seconds=30,
        idempotent=False,
        cost_class="deterministic",
    )
    async def review_invoice_uploads(
        self,
        ctx: RunContext[PlatformUserAuth],
        case_id: Annotated[
            str,
            Field(min_length=1, max_length=120, pattern=r"^[A-Za-z0-9_.:-]+$"),
        ],
        documents: Annotated[list[BrowserDocument], Field(min_length=1, max_length=MAX_UPLOADS)],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        parsed = _parse_browser_documents(documents)
        if not parsed["ok"]:
            await ctx.emit_error(parsed["message"], code=parsed["code"])
            return parsed
        result = _review_payload(case_id, parsed["invoices"])
        if not result["ok"]:
            await ctx.emit_error(result["message"], code=result["code"])
            return result
        persisted = await asyncio.to_thread(
            _persist_case_with_receipt,
            tenant,
            result,
            "browser_upload",
        )
        return {**result, "receipt": persisted["receipt"], "saved_at": persisted["saved_at"]}

    @a2a.tool(
        description="Review a typed Agent API file upload containing JSON, CSV, or plain-text invoice rows.",
        timeout_seconds=30,
        idempotent=False,
        cost_class="deterministic",
    )
    async def review_invoice_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        case_id: Annotated[
            str,
            Field(min_length=1, max_length=120, pattern=r"^[A-Za-z0-9_.:-]+$"),
        ],
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=sorted(ACCEPTED_MEDIA_TYPES),
                max_bytes=MAX_UPLOAD_BYTES,
                description="JSON array/object, CSV, or simple key-value text invoice data.",
            ),
        ],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        if document.size_bytes > MAX_UPLOAD_BYTES:
            return _validation_error("upload_too_large", f"Upload exceeds {MAX_UPLOAD_BYTES} bytes.")
        if document.media_type not in ACCEPTED_MEDIA_TYPES:
            return _validation_error("unsupported_media_type", "Upload must be JSON, CSV, or plain text.")
        data = ctx.workspace.read_bytes(document.path)
        parsed = _parse_upload_bytes(document.filename, document.media_type, data)
        if not parsed["ok"]:
            await ctx.emit_error(parsed["message"], code=parsed["code"])
            return parsed
        result = _review_payload(case_id, parsed["invoices"])
        if not result["ok"]:
            await ctx.emit_error(result["message"], code=result["code"])
            return result
        persisted = await asyncio.to_thread(
            _persist_case_with_receipt,
            tenant,
            result,
            "file_upload",
        )
        return {**result, "receipt": persisted["receipt"], "saved_at": persisted["saved_at"]}

    @a2a.tool(
        description="Reopen a previously saved user-scoped invoice review case, including decisions and receipt metadata.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="deterministic",
    )
    async def get_invoice_case(
        self,
        ctx: RunContext[PlatformUserAuth],
        case_id: Annotated[
            str,
            Field(min_length=1, max_length=120, pattern=r"^[A-Za-z0-9_.:-]+$"),
        ],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        try:
            clean_case_id = _clean_case_id(case_id)
        except ValueError as exc:
            return _validation_error("invalid_case_id", str(exc), case_id=str(case_id or ""))
        record = await asyncio.to_thread(_load_case, tenant, clean_case_id)
        if record is None:
            return _validation_error("case_not_found", "No invoice review case exists for this signed-in user and case_id.", case_id=case_id)
        return record

    @a2a.tool(
        description="Record or update a user decision for a saved invoice review case.",
        timeout_seconds=15,
        idempotent=False,
        cost_class="deterministic",
    )
    async def record_decision(
        self,
        ctx: RunContext[PlatformUserAuth],
        case_id: Annotated[
            str,
            Field(min_length=1, max_length=120, pattern=r"^[A-Za-z0-9_.:-]+$"),
        ],
        decision: DecisionInput,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        try:
            clean_case_id = _clean_case_id(case_id)
        except ValueError as exc:
            return _validation_error("invalid_case_id", str(exc), case_id=str(case_id or ""))
        updated = await asyncio.to_thread(
            _persist_decision,
            tenant,
            clean_case_id,
            decision.model_dump(),
        )
        if updated is None:
            return _validation_error("case_not_found", "Save or review the invoice case before recording decisions.", case_id=case_id)
        return updated


def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
    stable_id = ctx.auth.user_id if ctx.auth.user_id is not None else ctx.auth.sub
    if not stable_id:
        raise PermissionError("stable platform identity required")
    return f"user:{stable_id}"


def _clean_case_id(case_id: str) -> str:
    cleaned = str(case_id or "").strip()
    if not cleaned or len(cleaned) > 120 or not re.fullmatch(r"[A-Za-z0-9_.:-]+", cleaned):
        raise ValueError("case_id must be 1-120 chars using letters, numbers, dot, colon, underscore, or dash")
    return cleaned


def _money(value: Any) -> Decimal | None:
    if value is None or value == "":
        return None
    try:
        return Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)
    except (InvalidOperation, ValueError):
        return None


def _money_json(value: Decimal) -> float:
    return float(value.quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))


def _review_payload(case_id: str, raw_invoices: list[InvoiceInput | dict[str, Any]]) -> dict[str, Any]:
    try:
        clean_case_id = _clean_case_id(case_id)
    except ValueError as exc:
        return _validation_error("invalid_case_id", str(exc), case_id=str(case_id or ""))
    if not isinstance(raw_invoices, list) or not raw_invoices:
        return _validation_error("no_invoices", "Provide at least one invoice.", case_id=clean_case_id)
    if len(raw_invoices) > MAX_INVOICES:
        return _validation_error("too_many_invoices", f"At most {MAX_INVOICES} invoices are allowed per review.", case_id=clean_case_id)

    invoices: list[dict[str, Any]] = []
    for index, raw_item in enumerate(raw_invoices):
        if isinstance(raw_item, InvoiceInput):
            item = raw_item.model_dump()
        elif isinstance(raw_item, dict):
            try:
                item = InvoiceInput.model_validate(raw_item).model_dump()
            except ValidationError:
                item = raw_item
        else:
            return _validation_error("invalid_invoice", f"Invoice at index {index} must be an object.", case_id=clean_case_id)
        missing = [name for name in ("invoice_number", "vendor", "subtotal", "tax", "total") if item.get(name) in (None, "")]
        invoice_number = str(item.get("invoice_number", "")).strip()
        vendor = str(item.get("vendor", "")).strip()
        if not invoice_number or not vendor:
            missing.extend(name for name, value in (("invoice_number", invoice_number), ("vendor", vendor)) if not value and name not in missing)
        if missing:
            return _validation_error(
                "incomplete_invoice",
                f"Invoice at index {index} is missing required field(s): {', '.join(missing)}.",
                case_id=clean_case_id,
            )
        if len(invoice_number) > 120 or len(vendor) > MAX_TEXT_FIELD_LENGTH:
            return _validation_error("invalid_invoice", f"Invoice at index {index} exceeds text field length limits.", case_id=clean_case_id)
        subtotal = _money(item.get("subtotal"))
        tax = _money(item.get("tax"))
        total = _money(item.get("total"))
        if subtotal is None or tax is None or total is None:
            return _validation_error("invalid_amount", f"Invoice at index {index} has a non-numeric amount.", case_id=clean_case_id)
        invoices.append(
            {
                "index": index,
                "invoice_number": invoice_number,
                "vendor": vendor,
                "subtotal": _money_json(subtotal),
                "tax": _money_json(tax),
                "total": _money_json(total),
                "expected_total": _money_json(subtotal + tax),
            }
        )

    groups: dict[str, list[int]] = {}
    mismatches: list[dict[str, Any]] = []
    for inv in invoices:
        key = inv["invoice_number"].casefold()
        groups.setdefault(key, []).append(inv["index"])
        if Decimal(str(inv["total"])).quantize(Decimal("0.01")) != Decimal(str(inv["expected_total"])).quantize(Decimal("0.01")):
            mismatches.append(
                {
                    "index": inv["index"],
                    "invoice_number": inv["invoice_number"],
                    "vendor": inv["vendor"],
                    "subtotal": inv["subtotal"],
                    "tax": inv["tax"],
                    "stated_total": inv["total"],
                    "expected_total": inv["expected_total"],
                    "difference": _money_json(Decimal(str(inv["total"])) - Decimal(str(inv["expected_total"]))),
                }
            )
    duplicates = [
        {"invoice_number": invoices[indexes[0]]["invoice_number"], "indexes": indexes, "count": len(indexes)}
        for indexes in groups.values()
        if len(indexes) > 1
    ]
    return {
        "ok": True,
        "case_id": clean_case_id,
        "invoice_count": len(invoices),
        "duplicate_count": len(duplicates),
        "total_mismatch_count": len(mismatches),
        "duplicates": duplicates,
        "total_mismatches": mismatches,
        "invoices": invoices,
        "decisions": [],
    }


def _validation_error(code: str, message: str, *, case_id: str | None = None) -> dict[str, Any]:
    payload: dict[str, Any] = {"ok": False, "code": code, "message": message}
    if case_id is not None:
        payload["case_id"] = case_id
    return payload


def _safe_filename(name: str) -> bool:
    clean = str(name or "").strip()
    return bool(clean and len(clean) <= 180 and "/" not in clean and "\\" not in clean and "\x00" not in clean)


def _parse_browser_documents(documents: list[BrowserDocument | dict[str, Any]]) -> dict[str, Any]:
    if not documents:
        return _validation_error("no_uploads", "Attach at least one invoice document.")
    if len(documents) > MAX_UPLOADS:
        return _validation_error("too_many_uploads", f"At most {MAX_UPLOADS} uploads are allowed.")
    invoices: list[dict[str, Any]] = []
    for raw_document in documents:
        document = raw_document if isinstance(raw_document, BrowserDocument) else BrowserDocument.model_validate(raw_document)
        if not _safe_filename(document.filename):
            return _validation_error("invalid_filename", "Upload filename is missing or unsafe.")
        if document.media_type not in ACCEPTED_MEDIA_TYPES:
            return _validation_error("unsupported_media_type", "Upload must be JSON, CSV, or plain text.")
        try:
            data = base64.b64decode(document.data_base64, validate=True)
        except Exception:
            return _validation_error("invalid_base64", "Upload data_base64 is not valid base64.")
        if len(data) > MAX_UPLOAD_BYTES:
            return _validation_error("upload_too_large", f"Each upload must be at most {MAX_UPLOAD_BYTES} bytes.")
        parsed = _parse_upload_bytes(document.filename, document.media_type, data)
        if not parsed["ok"]:
            return parsed
        invoices.extend(parsed["invoices"])
    return {"ok": True, "invoices": invoices}


def _parse_upload_bytes(filename: str, media_type: str, data: bytes) -> dict[str, Any]:
    try:
        text = data.decode("utf-8-sig")
    except UnicodeDecodeError:
        return _validation_error("invalid_text_encoding", "Upload must be UTF-8 text.")
    if media_type in {"application/json", "text/json"} or filename.lower().endswith(".json"):
        try:
            payload = json.loads(text)
        except json.JSONDecodeError:
            return _validation_error("invalid_json", "Upload JSON could not be parsed.")
        if isinstance(payload, dict) and isinstance(payload.get("invoices"), list):
            return {"ok": True, "invoices": payload["invoices"]}
        if isinstance(payload, list):
            return {"ok": True, "invoices": payload}
        return _validation_error("invalid_json_shape", "JSON upload must be an invoice array or {invoices: [...]} object.")
    if media_type == "text/csv" or filename.lower().endswith(".csv"):
        rows = list(csv.DictReader(io.StringIO(text)))
        return {"ok": True, "invoices": rows}
    # Plain text bridge: accept one JSON object per line for simple browser demos.
    invoices: list[dict[str, Any]] = []
    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            item = json.loads(line)
        except json.JSONDecodeError:
            return _validation_error("invalid_text_invoice", "Plain text uploads must contain one JSON invoice object per line.")
        if not isinstance(item, dict):
            return _validation_error("invalid_text_invoice", "Plain text upload lines must be JSON objects.")
        invoices.append(item)
    return {"ok": True, "invoices": invoices}


@contextmanager
def _db_connection():
    database_url = os.environ.get("DATABASE_URL")
    if not database_url:
        raise RuntimeError("DATABASE_URL is not configured for the managed Postgres resource")
    import psycopg

    with psycopg.connect(database_url, options=DB_CONNECT_OPTIONS) as conn:
        yield conn


def _now_iso() -> str:
    return datetime.now(UTC).isoformat()


def _receipt_for(tenant: str, case: dict[str, Any], source: str, saved_at: str) -> dict[str, Any]:
    public_case = {
        "case_id": case["case_id"],
        "duplicate_count": case["duplicate_count"],
        "total_mismatch_count": case["total_mismatch_count"],
        "invoice_count": case["invoice_count"],
        "source": source,
    }
    digest = hashlib.sha256(json.dumps(public_case, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest()
    return {
        "receipt_id": f"invoice_guard:{case['case_id']}:{digest[:16]}",
        "kind": "production_execution_receipt",
        "agent": "invoice-guard-studio-v1",
        "agent_version": AGENT_VERSION,
        "skill": "review_invoices",
        "tenant_hash": hashlib.sha256(tenant.encode("utf-8")).hexdigest()[:16],
        "input_hash": digest,
        "created_at": saved_at,
    }


def _persist_case_with_receipt(tenant: str, case: dict[str, Any], source: str) -> dict[str, Any]:
    saved_at = _now_iso()
    receipt = _receipt_for(tenant, case, source, saved_at)
    case_payload = {**case, "receipt": receipt, "saved_at": saved_at, "updated_at": saved_at}
    with _db_connection() as conn:
        with conn.transaction():
            conn.execute(
                """
                INSERT INTO invoice_cases (
                    tenant_key, case_id, invoice_count, duplicate_count, total_mismatch_count,
                    invoices, duplicates, total_mismatches, decisions, receipt, created_at, updated_at
                ) VALUES (%s, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb, %s::jsonb, NOW(), NOW())
                ON CONFLICT (tenant_key, case_id) DO UPDATE SET
                    invoice_count = EXCLUDED.invoice_count,
                    duplicate_count = EXCLUDED.duplicate_count,
                    total_mismatch_count = EXCLUDED.total_mismatch_count,
                    invoices = EXCLUDED.invoices,
                    duplicates = EXCLUDED.duplicates,
                    total_mismatches = EXCLUDED.total_mismatches,
                    receipt = EXCLUDED.receipt,
                    updated_at = NOW()
                """,
                (
                    tenant,
                    case["case_id"],
                    case["invoice_count"],
                    case["duplicate_count"],
                    case["total_mismatch_count"],
                    json.dumps(case["invoices"]),
                    json.dumps(case["duplicates"]),
                    json.dumps(case["total_mismatches"]),
                    json.dumps(case.get("decisions", [])),
                    json.dumps(receipt),
                ),
            )
            conn.execute(
                """
                INSERT INTO invoice_execution_receipts (tenant_key, receipt_id, case_id, skill, input_hash, receipt, created_at)
                VALUES (%s, %s, %s, %s, %s, %s::jsonb, NOW())
                ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET receipt = EXCLUDED.receipt, created_at = NOW()
                """,
                (tenant, receipt["receipt_id"], case["case_id"], "review_invoices", receipt["input_hash"], json.dumps(receipt)),
            )
    return {"receipt": receipt, "saved_at": saved_at, "case": case_payload}


def _row_to_case(row: Any) -> dict[str, Any]:
    payload = {
        "ok": True,
        "case_id": row["case_id"],
        "invoice_count": row["invoice_count"],
        "duplicate_count": row["duplicate_count"],
        "total_mismatch_count": row["total_mismatch_count"],
        "invoices": row["invoices"],
        "duplicates": row["duplicates"],
        "total_mismatches": row["total_mismatches"],
        "decisions": row["decisions"],
        "receipt": row["receipt"],
        "saved_at": row["created_at"].isoformat() if hasattr(row["created_at"], "isoformat") else str(row["created_at"]),
        "updated_at": row["updated_at"].isoformat() if hasattr(row["updated_at"], "isoformat") else str(row["updated_at"]),
    }
    return payload


def _load_case(tenant: str, case_id: str) -> dict[str, Any] | None:
    clean_case_id = _clean_case_id(case_id)
    with _db_connection() as conn:
        conn.row_factory = _dict_row_factory()
        row = conn.execute(
            """
            SELECT case_id, invoice_count, duplicate_count, total_mismatch_count,
                   invoices, duplicates, total_mismatches, decisions, receipt, created_at, updated_at
            FROM invoice_cases
            WHERE tenant_key = %s AND case_id = %s
            """,
            (tenant, clean_case_id),
        ).fetchone()
    return _row_to_case(row) if row else None


def _persist_decision(tenant: str, case_id: str, decision: dict[str, Any]) -> dict[str, Any] | None:
    clean_case_id = _clean_case_id(case_id)
    decided_at = _now_iso()
    decision_payload = {**decision, "decided_at": decided_at}
    with _db_connection() as conn:
        conn.row_factory = _dict_row_factory()
        with conn.transaction():
            row = conn.execute(
                """
                SELECT decisions FROM invoice_cases
                WHERE tenant_key = %s AND case_id = %s
                FOR UPDATE
                """,
                (tenant, clean_case_id),
            ).fetchone()
            if row is None:
                return None
            decisions = [item for item in (row["decisions"] or []) if item.get("decision_id") != decision_payload["decision_id"]]
            decisions.append(decision_payload)
            conn.execute(
                """
                UPDATE invoice_cases
                SET decisions = %s::jsonb, updated_at = NOW()
                WHERE tenant_key = %s AND case_id = %s
                """,
                (json.dumps(decisions), tenant, clean_case_id),
            )
        committed = conn.execute(
            """
            SELECT case_id, invoice_count, duplicate_count, total_mismatch_count,
                   invoices, duplicates, total_mismatches, decisions, receipt, created_at, updated_at
            FROM invoice_cases
            WHERE tenant_key = %s AND case_id = %s
            """,
            (tenant, clean_case_id),
        ).fetchone()
        return _row_to_case(committed) if committed else None


def _dict_row_factory():
    from psycopg.rows import dict_row

    return dict_row