contract-clock-studio-v1/agent.py

"""ContractClock Studio: deterministic contract deadline extraction.

This agent is intentionally no-LLM. It extracts only explicit ISO-dated
renewals and notice windows from contracts, persists user-scoped timelines in
managed Postgres, and never guesses ambiguous dates.
"""
from __future__ import annotations

import base64
import binascii
import hashlib
import json
import os
import re
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import date, datetime, timezone, timedelta
from typing import Annotated, Any, Iterator, Literal

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,
    WorkspaceAccess,
    WorkspaceMode,
)

MAX_TEXT_CHARS = 200_000
MAX_UPLOAD_BYTES = 256 * 1024
ALLOWED_MEDIA_TYPES = ("text/plain", "text/markdown", "application/octet-stream")
DATABASE_ENV = "DATABASE_URL"
DB_STATEMENT_TIMEOUT_MS = 5_000
DB_CONNECT_TIMEOUT_SECONDS = 3
VERSION = "0.1.1"

ISO_DATE_RE = re.compile(r"\b(20\d{2}|19\d{2})-(0[1-9]|1[0-2])-([0-2]\d|3[01])\b")
NOTICE_RE = re.compile(
    r"(?:at\s+least\s+)?(?P<days>\d{1,3})\s+days?\s+(?:prior\s+)?(?:written\s+)?notice",
    re.IGNORECASE,
)
AMBIGUOUS_DATE_RE = re.compile(
    r"\b(?:sometime|around|approximately|about|next\s+(?:spring|summer|fall|autumn|winter|month|year)|well\s+in\s+advance|tbd|to\s+be\s+determined)\b",
    re.IGNORECASE,
)
RENEWAL_RE = re.compile(r"\brenew(?:s|al|ed|ing)?\b", re.IGNORECASE)
DEADLINE_KIND_ORDER = {"renewal": 0, "notice": 1, "obligation": 2}


class ContractClockStudioV1Config(BaseModel):
    pass


class ContractClockState(BaseModel):
    """Durable-state marker; records live in managed Postgres tables."""

    durable_store: str = "managed_postgres"


class Deadline(BaseModel):
    kind: Literal["renewal", "notice", "obligation"]
    date: str
    summary: str
    source_text: str


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 AnalyzeResult(BaseModel):
    ok: bool
    contract_id: str
    title: str | None = None
    deadlines: list[Deadline] = Field(default_factory=list)
    code: str | None = None
    message: str | None = None
    warnings: list[str] = Field(default_factory=list)
    receipt_id: str | None = None
    persisted: bool = False


class ContractRecord(BaseModel):
    ok: bool
    contract_id: str
    title: str | None = None
    deadlines: list[Deadline] = Field(default_factory=list)
    code: str | None = None
    message: str | None = None
    receipt_id: str | None = None
    persisted_at: str | None = None


class ContractSummary(BaseModel):
    contract_id: str
    title: str
    deadline_count: int
    updated_at: str


class ListContractsResult(BaseModel):
    ok: bool
    contracts: list[ContractSummary] = Field(default_factory=list)
    code: str | None = None
    message: str | None = None


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


@dataclass(frozen=True)
class ExtractedTimeline:
    deadlines: list[Deadline]
    warnings: list[str]


class ContractClockStudioV1(A2AAgent[ContractClockStudioV1Config, PlatformUserAuth]):
    name = "contract-clock-studio-v1"
    description = (
        "Private ContractClock app: paste or upload a contract, deterministically "
        "extract explicit renewal and notice dates, persist the timeline in "
        "user-scoped managed Postgres, and reopen it later."
    )
    version = VERSION

    config_model = ContractClockStudioV1Config
    auth_model = PlatformUserAuth

    # Deterministic local logic only: no LLM credential is read or required.
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="No LLM calls. Deterministic extraction and managed Postgres persistence only.",
    )
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    state = State.DURABLE
    state_model = ContractClockState
    tools_used = ("postgres", "mcp")
    workspace_access = WorkspaceAccess.dynamic(
        max_files=4,
        allowed_modes=(WorkspaceMode.READ_ONLY,),
        require_reason=False,
        max_total_size_bytes=MAX_UPLOAD_BYTES,
    )
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="contract-clock-studio-v1-data",
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url=DATABASE_ENV),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )

    @a2a.tool(
        description="Analyze pasted contract text and persist explicit renewal and notice deadlines.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def analyze_contract(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: str,
        text: str,
        title: str,
    ) -> AnalyzeResult:
        try:
            tenant = tenant_key(ctx)
            clean_id = validate_contract_id(contract_id)
            clean_title = validate_title(title)
            clean_text = validate_text(text)
            timeline = extract_timeline(clean_text)
        except ValidationFailure as exc:
            return AnalyzeResult(
                ok=False,
                contract_id=contract_id,
                title=title,
                code=exc.code,
                message=exc.message,
            )

        payload = {
            "contract_id": clean_id,
            "title": clean_title,
            "text_sha256": sha256_text(clean_text),
            "deadlines": [deadline.model_dump(mode="json") for deadline in timeline.deadlines],
            "warnings": timeline.warnings,
        }
        receipt_id = production_receipt_id(tenant, "analyze_contract", payload)
        try:
            persist_contract_with_receipt(
                tenant,
                clean_id,
                clean_title,
                payload,
                receipt_id,
                "analyze_contract",
                "ok",
            )
        except Exception:
            return AnalyzeResult(
                ok=False,
                contract_id=clean_id,
                title=clean_title,
                deadlines=timeline.deadlines,
                code="persistence_unavailable",
                message="ContractClock could not save the timeline. Please retry after database setup is healthy.",
                warnings=timeline.warnings,
                receipt_id=receipt_id,
                persisted=False,
            )

        await ctx.emit_progress(f"Saved {len(timeline.deadlines)} deadlines for {clean_id}.")
        return AnalyzeResult(
            ok=True,
            contract_id=clean_id,
            title=clean_title,
            deadlines=timeline.deadlines,
            warnings=timeline.warnings,
            receipt_id=receipt_id,
            persisted=True,
        )

    @a2a.tool(
        description="Analyze a bounded browser-uploaded base64 text contract and persist its timeline.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def analyze_contract_browser_upload(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: str,
        title: str,
        document: BrowserDocument,
    ) -> AnalyzeResult:
        try:
            text = decode_browser_document(document)
        except ValidationFailure as exc:
            return AnalyzeResult(ok=False, contract_id=contract_id, title=title, code=exc.code, message=exc.message)
        return await self.analyze_contract(ctx, contract_id=contract_id, title=title, text=text)

    @a2a.tool(
        description="Analyze a typed external FileUpload text contract and persist its timeline.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def analyze_contract_file_upload(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: str,
        title: str,
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=ALLOWED_MEDIA_TYPES,
                max_bytes=MAX_UPLOAD_BYTES,
                description="Plain-text contract file, max 256 KiB.",
            ),
        ],
    ) -> AnalyzeResult:
        try:
            if document.size_bytes > MAX_UPLOAD_BYTES:
                raise ValidationFailure("file_too_large", "Upload is larger than 256 KiB.")
            if document.media_type not in ALLOWED_MEDIA_TYPES:
                raise ValidationFailure("unsupported_media_type", "Only plain text uploads are supported.")
            await ctx.workspace.request_access(
                files=[document.path],
                mode=WorkspaceMode.READ_ONLY,
                reason="Read the caller-uploaded contract text for deadline extraction.",
                purpose="contract upload analysis",
            )
            read_bytes = getattr(ctx.workspace, "read_bytes", None)
            if read_bytes is None:
                raise RuntimeError("workspace backend does not expose direct file reads")
            raw = read_bytes(document.path)
            if len(raw) > MAX_UPLOAD_BYTES:
                raise ValidationFailure("file_too_large", "Upload is larger than 256 KiB.")
            text = raw.decode("utf-8")
        except UnicodeDecodeError:
            return AnalyzeResult(ok=False, contract_id=contract_id, title=title, code="invalid_encoding", message="Upload must be UTF-8 text.")
        except ValidationFailure as exc:
            return AnalyzeResult(ok=False, contract_id=contract_id, title=title, code=exc.code, message=exc.message)
        except Exception:
            return AnalyzeResult(ok=False, contract_id=contract_id, title=title, code="file_read_failed", message="Could not read the uploaded file from the workspace grant.")
        return await self.analyze_contract(ctx, contract_id=contract_id, title=title, text=text)

    @a2a.tool(
        description="Reload a previously saved contract timeline for the signed-in user.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="deterministic",
    )
    async def get_contract_deadlines(
        self,
        ctx: RunContext[PlatformUserAuth],
        contract_id: str,
    ) -> ContractRecord:
        try:
            tenant = tenant_key(ctx)
            clean_id = validate_contract_id(contract_id)
            record = load_contract(tenant, clean_id)
        except ValidationFailure as exc:
            return ContractRecord(ok=False, contract_id=contract_id, code=exc.code, message=exc.message)
        except Exception:
            return ContractRecord(ok=False, contract_id=contract_id, code="persistence_unavailable", message="ContractClock could not load timelines from storage.")
        if record is None:
            return ContractRecord(ok=False, contract_id=clean_id, code="not_found", message="No saved timeline exists for this contract and user.")
        payload = record["payload"]
        deadlines = [Deadline.model_validate(item) for item in payload.get("deadlines", [])]
        return ContractRecord(
            ok=True,
            contract_id=clean_id,
            title=payload.get("title"),
            deadlines=deadlines,
            receipt_id=record.get("receipt_id"),
            persisted_at=record.get("updated_at"),
        )

    @a2a.tool(
        description="List saved contract timelines for the signed-in user.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="deterministic",
    )
    async def list_contracts(self, ctx: RunContext[PlatformUserAuth], limit: int = 20) -> ListContractsResult:
        try:
            tenant = tenant_key(ctx)
            safe_limit = max(1, min(int(limit), 50))
            rows = list_contract_rows(tenant, safe_limit)
        except ValidationFailure as exc:
            return ListContractsResult(ok=False, code=exc.code, message=exc.message)
        except Exception:
            return ListContractsResult(ok=False, code="persistence_unavailable", message="ContractClock could not list saved timelines.")
        return ListContractsResult(ok=True, contracts=[ContractSummary.model_validate(row) for row in rows])


def tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
    auth = ctx.auth
    principal = auth.user_id if auth.user_id is not None else (auth.sub or auth.email)
    if not principal:
        raise ValidationFailure("auth_required", "A signed-in platform user is required.")
    return f"user:{principal}"


def validate_contract_id(contract_id: str) -> str:
    value = (contract_id or "").strip()
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", value):
        raise ValidationFailure("invalid_contract_id", "contract_id must be 1-128 safe identifier characters.")
    return value


def validate_title(title: str) -> str:
    value = (title or "").strip()
    if not value or len(value) > 180:
        raise ValidationFailure("invalid_title", "title is required and must be at most 180 characters.")
    return value


def validate_text(text: str) -> str:
    value = text or ""
    if not value.strip():
        raise ValidationFailure("empty_contract", "Contract text is required.")
    if len(value) > MAX_TEXT_CHARS:
        raise ValidationFailure("text_too_large", "Contract text exceeds the 200,000 character limit.")
    return value


def decode_browser_document(document: BrowserDocument) -> str:
    name = document.filename.strip().replace("\\", "/").split("/")[-1]
    if not name or name in {".", ".."}:
        raise ValidationFailure("invalid_filename", "Upload filename is invalid.")
    if document.media_type not in ALLOWED_MEDIA_TYPES:
        raise ValidationFailure("unsupported_media_type", "Only plain text uploads are supported.")
    try:
        raw = base64.b64decode(document.data_base64, validate=True)
    except (binascii.Error, ValueError):
        raise ValidationFailure("invalid_base64", "Upload data_base64 must be valid base64.") from None
    if len(raw) > MAX_UPLOAD_BYTES:
        raise ValidationFailure("file_too_large", "Upload is larger than 256 KiB.")
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError:
        raise ValidationFailure("invalid_encoding", "Upload must be UTF-8 text.") from None


def extract_timeline(text: str) -> ExtractedTimeline:
    if AMBIGUOUS_DATE_RE.search(text) and not ISO_DATE_RE.search(text):
        raise ValidationFailure("ambiguous_date", "The contract uses ambiguous date language; no deadline was invented.")

    dates: list[tuple[date, str]] = []
    for match in ISO_DATE_RE.finditer(text):
        token = match.group(0)
        try:
            parsed = date.fromisoformat(token)
        except ValueError:
            raise ValidationFailure("invalid_date", f"Invalid calendar date: {token}.") from None
        dates.append((parsed, source_excerpt(text, match.start(), match.end())))

    if not dates:
        if RENEWAL_RE.search(text) or NOTICE_RE.search(text) or AMBIGUOUS_DATE_RE.search(text):
            raise ValidationFailure("ambiguous_date", "No explicit ISO renewal date was found; ContractClock will not guess.")
        raise ValidationFailure("no_dated_obligations", "No explicit ISO-dated obligations were found.")

    deadlines: list[Deadline] = []
    warnings: list[str] = []
    renewal_dates = []
    for parsed, excerpt in dates:
        date_token = parsed.isoformat()
        date_index = text.find(date_token)
        local_context = text[max(0, date_index - 160): date_index + 160] if date_index >= 0 else excerpt
        if RENEWAL_RE.search(excerpt) or RENEWAL_RE.search(local_context):
            renewal_dates.append(parsed)
            deadlines.append(Deadline(kind="renewal", date=date_token, summary="Automatic renewal date explicitly stated in contract.", source_text=excerpt))
        else:
            deadlines.append(Deadline(kind="obligation", date=date_token, summary="Explicit dated obligation stated in contract.", source_text=excerpt))

    notice_match = NOTICE_RE.search(text)
    if notice_match:
        days = int(notice_match.group("days"))
        if not renewal_dates:
            raise ValidationFailure("ambiguous_date", "Notice window found, but no explicit renewal date anchors the notice deadline.")
        anchor = min(renewal_dates)
        notice_date = anchor - timedelta(days=days)
        deadlines.append(
            Deadline(
                kind="notice",
                date=notice_date.isoformat(),
                summary=f"Notice deadline calculated as {days} days before renewal date {anchor.isoformat()}.",
                source_text=source_excerpt(text, notice_match.start(), notice_match.end()),
            )
        )
    elif RENEWAL_RE.search(text):
        warnings.append("No explicit notice-window duration was found for the renewal date.")

    deduped: dict[tuple[str, str], Deadline] = {}
    for deadline in sorted(deadlines, key=deadline_sort_key):
        deduped.setdefault((deadline.kind, deadline.date), deadline)
    return ExtractedTimeline(deadlines=list(deduped.values()), warnings=warnings)


def deadline_sort_key(deadline: Deadline) -> tuple[int, str, str]:
    return (DEADLINE_KIND_ORDER[deadline.kind], deadline.date, deadline.summary)


def source_excerpt(text: str, start: int, end: int, radius: int = 140) -> str:
    left = max(0, start - radius)
    right = min(len(text), end + radius)
    return re.sub(r"\s+", " ", text[left:right]).strip()


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def production_receipt_id(tenant: str, skill: str, payload: dict[str, Any]) -> str:
    material = json.dumps({"tenant": tenant, "skill": skill, "payload": payload}, sort_keys=True, separators=(",", ":"))
    return "ccr_" + hashlib.sha256(material.encode("utf-8")).hexdigest()[:24]


@contextmanager
def db_connection() -> Iterator[Any]:
    url = os.environ.get(DATABASE_ENV)
    if not url:
        raise RuntimeError("DATABASE_URL is not configured")
    import psycopg

    with psycopg.connect(
        url,
        connect_timeout=DB_CONNECT_TIMEOUT_SECONDS,
        options=f"-c statement_timeout={DB_STATEMENT_TIMEOUT_MS} -c idle_in_transaction_session_timeout={DB_STATEMENT_TIMEOUT_MS}",
    ) as conn:
        yield conn


def persist_contract_with_receipt(
    tenant: str,
    contract_id: str,
    title: str,
    payload: dict[str, Any],
    receipt_id: str,
    skill_name: str,
    status: str,
) -> None:
    receipt = {
        "receipt_id": receipt_id,
        "agent": ContractClockStudioV1.name,
        "agent_version": ContractClockStudioV1.version,
        "skill": skill_name,
        "contract_id": contract_id,
        "input_hash": sha256_text(json.dumps(payload, sort_keys=True)),
        "status": status,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    with db_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO contract_timelines (tenant_key, contract_id, title, payload, receipt_id, updated_at)
                VALUES (%s, %s, %s, %s::jsonb, %s, NOW())
                ON CONFLICT (tenant_key, contract_id)
                DO UPDATE SET title = EXCLUDED.title, payload = EXCLUDED.payload,
                              receipt_id = EXCLUDED.receipt_id, updated_at = NOW()
                """,
                (tenant, contract_id, title, json.dumps(payload), receipt_id),
            )
            cur.execute(
                """
                INSERT INTO execution_receipts (tenant_key, receipt_id, skill_name, contract_id, status, payload)
                VALUES (%s, %s, %s, %s, %s, %s::jsonb)
                ON CONFLICT (tenant_key, receipt_id) DO NOTHING
                """,
                (tenant, receipt_id, skill_name, contract_id, status, json.dumps(receipt)),
            )
        conn.commit()


def load_contract(tenant: str, contract_id: str) -> dict[str, Any] | None:
    with db_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT payload, receipt_id, updated_at
                FROM contract_timelines
                WHERE tenant_key = %s AND contract_id = %s
                """,
                (tenant, contract_id),
            )
            row = cur.fetchone()
    if row is None:
        return None
    payload, receipt_id, updated_at = row
    return {
        "payload": payload if isinstance(payload, dict) else json.loads(payload),
        "receipt_id": receipt_id,
        "updated_at": updated_at.isoformat() if hasattr(updated_at, "isoformat") else str(updated_at),
    }


def list_contract_rows(tenant: str, limit: int) -> list[dict[str, Any]]:
    with db_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT contract_id, title, jsonb_array_length(payload->'deadlines') AS deadline_count, updated_at
                FROM contract_timelines
                WHERE tenant_key = %s
                ORDER BY updated_at DESC
                LIMIT %s
                """,
                (tenant, limit),
            )
            rows = cur.fetchall()
    return [
        {
            "contract_id": contract_id,
            "title": title,
            "deadline_count": int(deadline_count or 0),
            "updated_at": updated_at.isoformat() if hasattr(updated_at, "isoformat") else str(updated_at),
        }
        for contract_id, title, deadline_count, updated_at in rows
    ]