contract-clock-studio-v1/agent.py

"""ContractClock Studio full-stack A2A agent.

Deterministic product backend for the ContractClock one-page startup. The agent
extracts only explicit dated renewal obligations and notice windows, persists
contract timelines in user-scoped managed Postgres, and exposes the same typed
A2A tools to Agent API clients, MCP clients, and the packed React frontend.
"""
from __future__ import annotations

import asyncio
import base64
import binascii
import hashlib
import os
import re
from datetime import date, datetime, timedelta, timezone
from typing import Annotated, Any

from pydantic import BaseModel, Field

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


MAX_CONTRACT_TEXT_CHARS = 60_000
MAX_UPLOAD_BYTES = 512_000
ALLOWED_BROWSER_MEDIA_TYPES = (
    "text/plain",
    "text/markdown",
    "application/octet-stream",
)
DATE_RE = re.compile(r"\b(20\d{2}-\d{2}-\d{2})\b")
RENEWAL_DATE_RE = re.compile(
    r"\brenew(?:s|al|ed|ing)?\b[^.\n;]{0,160}?\b(?:on|as\s+of|effective(?:\s+on)?|"
    r"date(?:\s+is)?|until|through)\b\s*:?\s*(20\d{2}-\d{2}-\d{2})\b",
    re.IGNORECASE,
)
RENEWAL_RE = re.compile(r"\brenew(?:s|al|ed|ing)?\b", re.IGNORECASE)
NOTICE_RE = re.compile(
    r"(?:at\s+least\s+)?(?P<days>\d{1,3})\s+days?\s+(?:written\s+)?notice",
    re.IGNORECASE,
)
AMBIGUOUS_DATE_RE = re.compile(
    r"\b(sometime|spring|summer|fall|autumn|winter|soon|later|well\s+in\s+advance|next\s+\w+)\b",
    re.IGNORECASE,
)


class ContractClockStudioV1Config(BaseModel):
    pass


class BrowserDocument(BaseModel):
    filename: Annotated[str, Field(min_length=1, max_length=160)]
    media_type: Annotated[str, Field(min_length=1, max_length=120)]
    data_base64: Annotated[str, Field(min_length=1, max_length=700_000)]


class ContractClockStudioV1(A2AAgent[ContractClockStudioV1Config, PlatformUserAuth]):
    name = "contract-clock-studio-v1"
    description = (
        "ContractClock extracts explicit renewal and notice deadlines from "
        "contracts, saves user-scoped timelines, and reopens them in a packed "
        "one-page studio."
    )
    version = "0.1.5"

    config_model = ContractClockStudioV1Config
    auth_model = PlatformUserAuth

    state = State.DURABLE
    state_model = ContractClockStudioV1Config
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="contract-clock-studio-v1-data",
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic local extraction; no LLM credential required.",
    )
    tools_used = ("postgres", "mcp", "packed-frontend")

    @a2a.tool(
        description=(
            "Analyze pasted contract text, extract explicit renewal and notice "
            "deadlines, atomically persist the contract timeline and execution receipt."
        ),
        timeout_seconds=30,
        idempotent=False,
        cost_class="cheap",
    )
    async def analyze_contract(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: Annotated[str, Field(min_length=1, max_length=120)],
        text: Annotated[str, Field(min_length=1, max_length=MAX_CONTRACT_TEXT_CHARS)],
        title: Annotated[str, Field(min_length=1, max_length=200)] = "Untitled contract",
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        if tenant is None:
            return _auth_failure()
        analysis = _analyze_text(contract_id=contract_id, title=title, text=text)
        if not analysis["ok"]:
            return analysis
        try:
            await _save_analysis(
                tenant_key=tenant,
                contract_id=contract_id,
                title=title,
                text=text,
                analysis=analysis,
                skill_name="analyze_contract",
                task_id=getattr(ctx, "task_id", ""),
            )
        except DatabaseUnavailable as exc:
            return _db_failure(exc)
        return analysis

    @a2a.tool(
        description="Reload a previously analyzed contract deadline timeline from user-scoped storage.",
        timeout_seconds=20,
        idempotent=True,
        cost_class="cheap",
    )
    async def get_contract_deadlines(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: Annotated[str, Field(min_length=1, max_length=120)],
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        if tenant is None:
            return _auth_failure()
        try:
            record = await _load_contract(tenant_key=tenant, contract_id=contract_id)
        except DatabaseUnavailable as exc:
            return _db_failure(exc)
        if record is None:
            return {"ok": False, "code": "not_found", "contract_id": contract_id}
        return {
            "ok": True,
            "contract_id": contract_id,
            "title": record.get("title") or "Untitled contract",
            "deadlines": record.get("deadlines") or [],
            "receipt": record.get("receipt") or {},
            "updated_at": record.get("updated_at"),
        }

    @a2a.tool(
        description=(
            "Analyze a text contract uploaded by an external Agent API client. "
            "The FileUpload schema is retained for typed multipart clients."
        ),
        timeout_seconds=30,
        idempotent=False,
        cost_class="cheap",
    )
    async def analyze_contract_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: Annotated[str, Field(min_length=1, max_length=120)],
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=ALLOWED_BROWSER_MEDIA_TYPES,
                max_bytes=MAX_UPLOAD_BYTES,
                description="UTF-8 text contract document",
            ),
        ],
        title: Annotated[str, Field(min_length=1, max_length=200)] = "Uploaded contract",
    ) -> dict[str, Any]:
        try:
            payload = _read_uploaded_file(ctx, document)
        except UploadRejected as exc:
            return exc.result
        return await self.analyze_contract(ctx, contract_id=contract_id, text=payload, title=title)

    @a2a.tool(
        description=(
            "Browser-safe upload bridge accepting one bounded base64 JSON text "
            "document from the packed frontend."
        ),
        timeout_seconds=30,
        idempotent=False,
        cost_class="cheap",
    )
    async def analyze_contract_upload(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: Annotated[str, Field(min_length=1, max_length=120)],
        document: BrowserDocument,
        title: Annotated[str, Field(min_length=1, max_length=200)] = "Uploaded contract",
    ) -> dict[str, Any]:
        try:
            payload = _decode_browser_document(document)
        except UploadRejected as exc:
            return exc.result
        return await self.analyze_contract(ctx, contract_id=contract_id, text=payload, title=title)


class DatabaseUnavailable(RuntimeError):
    def __init__(self, code: str, message: str) -> None:
        self.code = code
        super().__init__(message)


class UploadRejected(ValueError):
    def __init__(self, code: str, message: str) -> None:
        self.result = {"ok": False, "code": code, "message": message}
        super().__init__(message)


def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str | None:
    auth = getattr(ctx, "auth", None)
    user_id = getattr(auth, "user_id", None)
    sub = (getattr(auth, "sub", None) or "").strip()
    if user_id is not None:
        return f"user:{user_id}"
    if sub:
        return f"user:{sub}"
    return None


def _auth_failure() -> dict[str, Any]:
    return {
        "ok": False,
        "code": "auth_required",
        "message": "A signed-in A2A Cloud user session is required to use ContractClock.",
    }


def _db_failure(exc: DatabaseUnavailable) -> dict[str, Any]:
    return {
        "ok": False,
        "code": exc.code,
        "message": "ContractClock storage is not ready. Retry after the managed database is available.",
    }


def _analyze_text(*, contract_id: str, title: str, text: str) -> dict[str, Any]:
    clean_text = (text or "").strip()
    if not clean_text:
        return {"ok": False, "code": "empty_contract", "contract_id": contract_id}
    if len(clean_text) > MAX_CONTRACT_TEXT_CHARS:
        return {"ok": False, "code": "text_too_large", "contract_id": contract_id}

    explicit_dates = []
    for match in DATE_RE.finditer(clean_text):
        try:
            explicit_dates.append(date.fromisoformat(match.group(1)))
        except ValueError:
            return {"ok": False, "code": "invalid_date", "contract_id": contract_id}

    if not explicit_dates:
        code = "ambiguous_date" if AMBIGUOUS_DATE_RE.search(clean_text) else "missing_explicit_date"
        return {"ok": False, "code": code, "contract_id": contract_id}

    deadlines: list[dict[str, str]] = []
    renewal_match = RENEWAL_DATE_RE.search(clean_text)
    if renewal_match:
        renewal_date = date.fromisoformat(renewal_match.group(1))
        deadlines.append({"kind": "renewal", "date": renewal_date.isoformat()})
        notice_match = NOTICE_RE.search(clean_text)
        if notice_match:
            notice_days = int(notice_match.group("days"))
            if 1 <= notice_days <= 366:
                deadlines.append(
                    {"kind": "notice", "date": (renewal_date - timedelta(days=notice_days)).isoformat()}
                )
    elif RENEWAL_RE.search(clean_text):
        return {
            "ok": False,
            "code": "ambiguous_renewal_date",
            "contract_id": contract_id,
        }
    else:
        # Explicit dates outside a renewal clause are kept as dated obligations;
        # relative invoice terms without a concrete anchor are intentionally not guessed.
        for item in explicit_dates:
            deadlines.append({"kind": "obligation", "date": item.isoformat()})

    if not deadlines:
        return {"ok": False, "code": "no_supported_deadline", "contract_id": contract_id}

    return {
        "ok": True,
        "contract_id": contract_id,
        "title": title,
        "deadlines": deadlines,
        "rationale": (
            "Only ISO-formatted explicit dates were extracted. Notice deadlines "
            "are calculated from explicit renewal dates and written-notice day counts."
        ),
    }


def _decode_browser_document(document: BrowserDocument) -> str:
    filename = document.filename.strip()
    if "/" in filename or "\\" in filename or filename in {".", ".."}:
        raise UploadRejected("invalid_filename", "Upload filename must be a simple file name.")
    media_type = document.media_type.split(";", 1)[0].strip().lower()
    if media_type not in ALLOWED_BROWSER_MEDIA_TYPES:
        raise UploadRejected("unsupported_media_type", "Upload a UTF-8 text contract file.")
    try:
        raw = base64.b64decode(document.data_base64, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise UploadRejected("invalid_base64", "Upload data must be valid base64.") from exc
    return _decode_uploaded_bytes(raw)


def _read_uploaded_file(ctx: RunContext[PlatformUserAuth], document: UploadedFile) -> str:
    media_type = document.media_type.split(";", 1)[0].strip().lower()
    if media_type not in ALLOWED_BROWSER_MEDIA_TYPES:
        raise UploadRejected("unsupported_media_type", "Upload a UTF-8 text contract file.")
    if document.size_bytes > MAX_UPLOAD_BYTES:
        raise UploadRejected("upload_too_large", "Contract uploads are limited to 512 KB.")
    try:
        raw = ctx.workspace.read_bytes(document.path)
    except Exception as exc:  # noqa: BLE001
        raise UploadRejected("upload_read_failed", "Could not read the uploaded contract file.") from exc
    return _decode_uploaded_bytes(raw)


def _decode_uploaded_bytes(raw: bytes) -> str:
    if not raw:
        raise UploadRejected("empty_upload", "Uploaded contract file is empty.")
    if len(raw) > MAX_UPLOAD_BYTES:
        raise UploadRejected("upload_too_large", "Contract uploads are limited to 512 KB.")
    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise UploadRejected("invalid_text_encoding", "Uploaded contract must be UTF-8 text.") from exc
    if not text.strip():
        raise UploadRejected("empty_upload", "Uploaded contract file is empty.")
    return text


async def _save_analysis(
    *,
    tenant_key: str,
    contract_id: str,
    title: str,
    text: str,
    analysis: dict[str, Any],
    skill_name: str,
    task_id: str,
) -> None:
    await asyncio.to_thread(
        _save_analysis_sync,
        tenant_key=tenant_key,
        contract_id=contract_id,
        title=title,
        text=text,
        analysis=analysis,
        skill_name=skill_name,
        task_id=task_id,
    )


async def _load_contract(*, tenant_key: str, contract_id: str) -> dict[str, Any] | None:
    return await asyncio.to_thread(_load_contract_sync, tenant_key=tenant_key, contract_id=contract_id)


def _connect_db() -> Any:
    dsn = os.environ.get("DATABASE_URL", "").strip()
    if not dsn:
        raise DatabaseUnavailable("database_unavailable", "DATABASE_URL is not configured")
    try:
        import psycopg
    except Exception as exc:  # noqa: BLE001
        raise DatabaseUnavailable("database_driver_missing", "psycopg is not installed") from exc
    try:
        return psycopg.connect(
            dsn,
            connect_timeout=5,
            options="-c statement_timeout=5000 -c idle_in_transaction_session_timeout=5000",
        )
    except Exception as exc:  # noqa: BLE001
        raise DatabaseUnavailable("database_unavailable", "could not connect to managed Postgres") from exc


def _save_analysis_sync(
    *,
    tenant_key: str,
    contract_id: str,
    title: str,
    text: str,
    analysis: dict[str, Any],
    skill_name: str,
    task_id: str,
) -> None:
    try:
        from psycopg.types.json import Jsonb
    except Exception as exc:  # noqa: BLE001
        raise DatabaseUnavailable("database_driver_missing", "psycopg JSON support is not installed") from exc

    now = datetime.now(timezone.utc).isoformat()
    input_hash = hashlib.sha256(
        f"{tenant_key}\0{contract_id}\0{text}".encode("utf-8")
    ).hexdigest()
    receipt = {
        "agent": ContractClockStudioV1.name,
        "version": ContractClockStudioV1.version,
        "skill": skill_name,
        "task_id": task_id,
        "input_hash": input_hash,
        "status": "ok",
        "created_at": now,
    }
    analysis_with_receipt = {**analysis, "receipt": receipt}

    try:
        with _connect_db() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO contract_clock_contracts
                        (tenant_key, contract_id, title, source_text_sha256, deadlines, analysis, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, NOW())
                    ON CONFLICT (tenant_key, contract_id) DO UPDATE SET
                        title = EXCLUDED.title,
                        source_text_sha256 = EXCLUDED.source_text_sha256,
                        deadlines = EXCLUDED.deadlines,
                        analysis = EXCLUDED.analysis,
                        updated_at = NOW()
                    """,
                    (
                        tenant_key,
                        contract_id,
                        title,
                        hashlib.sha256(text.encode("utf-8")).hexdigest(),
                        Jsonb(analysis["deadlines"]),
                        Jsonb(analysis_with_receipt),
                    ),
                )
                cur.execute(
                    """
                    INSERT INTO contract_clock_execution_receipts
                        (tenant_key, contract_id, skill_name, task_id, input_hash, status, result_summary)
                    VALUES (%s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        tenant_key,
                        contract_id,
                        skill_name,
                        task_id,
                        input_hash,
                        "ok",
                        Jsonb(
                            {
                                "deadline_count": len(analysis.get("deadlines") or []),
                                "contract_id": contract_id,
                            }
                        ),
                    ),
                )
    except DatabaseUnavailable:
        raise
    except Exception as exc:  # noqa: BLE001
        raise DatabaseUnavailable("database_write_failed", "could not persist contract analysis") from exc


def _load_contract_sync(*, tenant_key: str, contract_id: str) -> dict[str, Any] | None:
    try:
        with _connect_db() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT title, deadlines, analysis, updated_at
                    FROM contract_clock_contracts
                    WHERE tenant_key = %s AND contract_id = %s
                    """,
                    (tenant_key, contract_id),
                )
                row = cur.fetchone()
    except DatabaseUnavailable:
        raise
    except Exception as exc:  # noqa: BLE001
        raise DatabaseUnavailable("database_read_failed", "could not load contract timeline") from exc
    if row is None:
        return None
    title, deadlines, analysis, updated_at = row
    analysis = analysis or {}
    return {
        "title": title,
        "deadlines": deadlines or [],
        "receipt": analysis.get("receipt") or {},
        "updated_at": updated_at.isoformat() if hasattr(updated_at, "isoformat") else str(updated_at),
    }