Jyotisha/scripts/vedastro_foreground.py

"""Foreground VedAstro wait, cancel, and snapshot-cache coordination.

Join, budget, and worker count bound the same wall clock, so they are
declared together and must be changed together.

join: how long the consultation thread waits for the official gateway
      (default 1.5s, cap 3s).
budget: how long a worker may keep running after it has started.
workers: process-wide pool. Extra work queues; join still caps the wait.

A timed-out join must cancel its future so queued work never starts.

A gateway run that ends without official data (blocked, rate limited, local
fallback, its own timeout) is remembered for that chart and UTC day in the
existing negative snapshot store, so later turns that day answer at once
instead of starting the gateway again (issue #130). A success still wins:
the positive snapshot is looked up first.
budget <= k * join (k=2) keeps leftover occupancy of an already-running
worker from stretching later requests past the join cap. The previous
default budget of 8s with a 1.5s join left cancelled-too-late workers
holding the two-thread pool for the rest of the 8s.

Report profile (body ``vedastro_evidence_budget="report"``, gap 2 of
docs/operations/outage-behaviour-2026-11-01.md, BUG-1508). A report used to
wait synchronously on the main-entry overview and the rectification gate,
about 128s when VedAstro hangs, past the frontend's 90s. A report now runs
the same deferred session as consultation, with a longer wait so a healthy
VedAstro still lands in the report:

deadline: measured from session start, so it overlaps local compute
          (default 25s, cap 45s).
budget:   the per-request timeout already configured for the adapter, never
          more than k * deadline.

When the deadline passes while the gateway is still running, the report
continues with VedAstro marked ``official_blocked`` /
``report_vedastro_deadline_exceeded`` and that is remembered for the chart
and UTC day, so the next report answers at once. The worker keeps running;
if it later succeeds the positive snapshot wins, if it fails its own failure
replaces the deadline entry.
"""
from __future__ import annotations

import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
from typing import Any

try:
    from scripts.vedastro_consultation_snapshot import (
        DEFERRED_IN_CONSULTATION,
        lookup_negative_snapshot,
        store_negative_snapshot,
    )
    from scripts.vedastro_runtime_context import temporary_timeout_seconds
    from scripts.vedastro_snapshot_cache import (
        annotate_stale_gateway,
        is_cacheable_gateway,
        lookup_snapshot,
        official_snapshot_reference_date,
        requires_today_snapshot,
        store_snapshot,
    )
except ModuleNotFoundError:  # pragma: no cover - script execution path
    from vedastro_consultation_snapshot import (
        DEFERRED_IN_CONSULTATION,
        lookup_negative_snapshot,
        store_negative_snapshot,
    )
    from vedastro_runtime_context import temporary_timeout_seconds
    from vedastro_snapshot_cache import (
        annotate_stale_gateway,
        is_cacheable_gateway,
        lookup_snapshot,
        official_snapshot_reference_date,
        requires_today_snapshot,
        store_snapshot,
    )

# --- join / budget / workers (one declaration block; change together) ---
FOREGROUND_VEDASTRO_JOIN_SECONDS_DEFAULT = 1.5
FOREGROUND_VEDASTRO_JOIN_SECONDS_MAX = 3.0
FOREGROUND_VEDASTRO_BUDGET_JOIN_RATIO = 2.0
FOREGROUND_VEDASTRO_WORKERS_DEFAULT = 2
# ----------------------------------------------------------------------

# --- report profile: deadline / budget (change together, see docstring) ---
REPORT_VEDASTRO_EVIDENCE_BUDGET = "report"
REPORT_VEDASTRO_DEADLINE_SECONDS_DEFAULT = 25.0
REPORT_VEDASTRO_DEADLINE_SECONDS_MAX = 45.0
REPORT_VEDASTRO_DEADLINE_EXCEEDED = "report_vedastro_deadline_exceeded"
# ----------------------------------------------------------------------

_FOREGROUND_VEDASTRO_WORKERS = max(
    int(os.environ.get("JYOTISH_FOREGROUND_VEDASTRO_WORKERS", str(FOREGROUND_VEDASTRO_WORKERS_DEFAULT))),
    1,
)
_FOREGROUND_VEDASTRO_EXECUTOR = ThreadPoolExecutor(
    max_workers=_FOREGROUND_VEDASTRO_WORKERS,
    thread_name_prefix="jyotish-vedastro-fg",
)


def _foreground_vedastro_join_seconds() -> float:
    raw = str(os.environ.get("JYOTISH_FOREGROUND_VEDASTRO_JOIN_SECONDS", str(FOREGROUND_VEDASTRO_JOIN_SECONDS_DEFAULT))).strip()
    try:
        return min(max(float(raw), 0.0), FOREGROUND_VEDASTRO_JOIN_SECONDS_MAX)
    except ValueError:
        return FOREGROUND_VEDASTRO_JOIN_SECONDS_DEFAULT


def _foreground_vedastro_budget_seconds() -> float:
    join = _foreground_vedastro_join_seconds()
    cap = FOREGROUND_VEDASTRO_BUDGET_JOIN_RATIO * join
    raw = str(os.environ.get("JYOTISH_FOREGROUND_VEDASTRO_BUDGET_SECONDS", "")).strip()
    if not raw:
        return cap
    try:
        budget = float(raw)
    except ValueError:
        return cap
    if budget < 0:
        return cap
    return min(budget, cap)


def is_report_vedastro_budget(body: dict | None) -> bool:
    payload = body if isinstance(body, dict) else {}
    return str(payload.get("vedastro_evidence_budget") or "").strip().lower() == REPORT_VEDASTRO_EVIDENCE_BUDGET


def defers_optional_vedastro(body: dict | None) -> bool:
    """Consult asks with the flag; a report budget implies it (BUG-1508)."""
    payload = body if isinstance(body, dict) else {}
    return bool(payload.get("defer_optional_external_evidence")) or is_report_vedastro_budget(payload)


def _report_vedastro_deadline_seconds() -> float:
    raw = str(os.environ.get("JYOTISH_REPORT_VEDASTRO_DEADLINE_SECONDS", str(REPORT_VEDASTRO_DEADLINE_SECONDS_DEFAULT))).strip()
    try:
        return min(max(float(raw), 0.0), REPORT_VEDASTRO_DEADLINE_SECONDS_MAX)
    except ValueError:
        return REPORT_VEDASTRO_DEADLINE_SECONDS_DEFAULT


def _report_vedastro_budget_seconds() -> float:
    """The adapter's own per-request timeout; a report never raises it."""
    try:
        from scripts.vedastro_service_adapter import _timeout_seconds as configured_timeout
    except ModuleNotFoundError:  # pragma: no cover - script execution path
        from vedastro_service_adapter import _timeout_seconds as configured_timeout
    cap = FOREGROUND_VEDASTRO_BUDGET_JOIN_RATIO * _report_vedastro_deadline_seconds()
    return min(float(configured_timeout()), cap)


def _blocked_foreground_vedastro(*, reason: str, error_type: str | None = None) -> dict:
    packet = {
        "scope": "vedastro_gateway_run",
        "status": "official_blocked",
        "official_closure_state": "official_blocked",
        "official_closure_reason": reason,
    }
    if error_type:
        packet["error_type"] = error_type
    return packet


def _run_foreground_vedastro_gateway(
    handler,
    body: dict,
    cancel_event: threading.Event | None = None,
    *,
    budget_seconds: float | None = None,
) -> dict:
    if cancel_event is not None and cancel_event.is_set():
        return _blocked_foreground_vedastro(reason="foreground_optional_evidence_timeout")
    budget = _foreground_vedastro_budget_seconds() if budget_seconds is None else budget_seconds
    try:
        with temporary_timeout_seconds(budget):
            result = handler._compute_vedastro_gateway_run(body)
    except Exception as exc:
        return _blocked_foreground_vedastro(
            reason="gateway_invocation_error",
            error_type=type(exc).__name__,
        )
    return result if isinstance(result, dict) else _blocked_foreground_vedastro(reason="gateway_invocation_error")


def _remember_gateway_outcome(body: dict, result: dict) -> None:
    """Success into the snapshot cache; a failure that really ran, for the day."""
    try:
        if is_cacheable_gateway(result):
            store_snapshot(body, result)
        elif result.get("official_closure_reason") != "foreground_optional_evidence_timeout":
            store_negative_snapshot(body, result)
    except (OSError, ValueError):
        pass


def _failed_today(body: dict, today: str) -> dict | None:
    """Today's remembered failure. A consult-gate deferral never ran, so it does not count."""
    try:
        negative = lookup_negative_snapshot(body, today=today)
    except (OSError, ValueError):
        return None
    if negative is None or negative.get("official_closure_reason") == DEFERRED_IN_CONSULTATION:
        return None
    return negative


def _run_foreground_vedastro_gateway_and_store(
    handler,
    body: dict,
    cancel_event: threading.Event | None = None,
    budget_seconds: float | None = None,
) -> dict:
    result = _run_foreground_vedastro_gateway(handler, body, cancel_event, budget_seconds=budget_seconds)
    # A run cancelled before it started never asked VedAstro: nothing to remember.
    if not (cancel_event is not None and cancel_event.is_set() and result.get("official_closure_reason") == "foreground_optional_evidence_timeout"):
        _remember_gateway_outcome(body, result)
    return result


def _join_foreground_vedastro(future, *, timeout: float, cancel_event: threading.Event | None = None) -> dict:
    if future is None:
        return _blocked_foreground_vedastro(reason="foreground_optional_evidence_timeout")
    try:
        result = future.result(timeout=timeout)
    except FuturesTimeoutError:
        if cancel_event is not None:
            cancel_event.set()
        future.cancel()
        return _blocked_foreground_vedastro(reason="foreground_optional_evidence_timeout")
    except Exception as exc:
        return _blocked_foreground_vedastro(
            reason="gateway_invocation_error",
            error_type=type(exc).__name__,
        )
    return result if isinstance(result, dict) else _blocked_foreground_vedastro(reason="gateway_invocation_error")


def should_include_western_evidence_packet(body: dict | None, *, surface: str = "") -> bool:
    payload = body if isinstance(body, dict) else {}
    if payload.get("include_western_evidence_packet") or payload.get("return_western_evidence_packet"):
        return True
    if payload.get("return_high_rigor_shape"):
        return True
    if isinstance(payload.get("western_oracle_payload"), dict) or isinstance(payload.get("western_astrology_oracle"), dict):
        return True
    if isinstance(payload.get("western_evidence_packet"), dict):
        return True
    return str(surface or "") in {"skill_mcp", "professional_reading_web", "mcp"}


class ForegroundVedastroSession:
    def __init__(self, handler, body: dict, *, defer: bool) -> None:
        self.handler = handler
        self.body = dict(body)
        self.defer = bool(defer)
        self.cached_gateway: dict | None = None
        self.future = None
        self.cancel_event = threading.Event()
        self.meta: dict[str, Any] = {}
        self.report = is_report_vedastro_budget(self.body)
        self.started_at = time.monotonic()

    def _budget_seconds(self) -> float | None:
        return _report_vedastro_budget_seconds() if self.report else None

    def _join_seconds(self) -> float:
        if not self.report:
            return _foreground_vedastro_join_seconds()
        return max(self.started_at + _report_vedastro_deadline_seconds() - time.monotonic(), 0.0)

    def _report_deadline_exceeded(self, result: dict) -> dict:
        """A report gave up on a gateway that is still running: remember it for today."""
        future = self.future
        if (
            not self.report
            or result.get("official_closure_reason") != "foreground_optional_evidence_timeout"
            or future is None
            or future.cancelled()
            or future.done()
        ):
            return result
        packet = _blocked_foreground_vedastro(reason=REPORT_VEDASTRO_DEADLINE_EXCEEDED)
        packet["deadline_seconds"] = _report_vedastro_deadline_seconds()
        try:
            store_negative_snapshot(self.body, packet)
        except (OSError, ValueError):
            pass
        return packet

    def start(self) -> "ForegroundVedastroSession":
        self.started_at = time.monotonic()
        today = official_snapshot_reference_date(self.body)
        hit = lookup_snapshot(self.body, today=today)
        if hit is not None:
            record = hit["record"]
            freshness = hit["freshness"]
            gateway = record["gateway"]
            if freshness == "stale":
                gateway = annotate_stale_gateway(
                    gateway,
                    reference_date=str(record.get("reference_date") or ""),
                    served_on_utc_date=today,
                )
                refresh_body = dict(self.body)
                refresh_body["reference_date"] = today
                refresh_body["today"] = today
                refresh_body["current_date"] = today
                self.future = _FOREGROUND_VEDASTRO_EXECUTOR.submit(
                    _run_foreground_vedastro_gateway_and_store,
                    self.handler,
                    refresh_body,
                    None,
                    self._budget_seconds(),
                )
            self.cached_gateway = gateway
            self.meta = {
                "freshness": freshness,
                "reference_date": str(record.get("reference_date") or today),
                "served_on_utc_date": today,
                "refresh_submitted": freshness == "stale",
            }
            return self
        failed = _failed_today(self.body, today)
        if failed is not None:
            self.cached_gateway = failed
            self.meta = {
                "freshness": "negative",
                "reference_date": today,
                "served_on_utc_date": today,
                "refresh_submitted": False,
            }
            return self
        self.meta = {
            "freshness": "miss",
            "reference_date": today,
            "served_on_utc_date": today,
            "refresh_submitted": False,
            "require_today": requires_today_snapshot(self.body),
        }
        if self.defer:
            self.future = _FOREGROUND_VEDASTRO_EXECUTOR.submit(
                _run_foreground_vedastro_gateway_and_store,
                self.handler,
                dict(self.body),
                self.cancel_event,
                self._budget_seconds(),
            )
        return self

    def finish(self) -> dict:
        if self.cached_gateway is not None:
            return self.cached_gateway
        if self.defer:
            result = _join_foreground_vedastro(
                self.future,
                timeout=self._join_seconds(),
                cancel_event=self.cancel_event,
            )
            return self._report_deadline_exceeded(result)
        try:
            result = self.handler._compute_vedastro_gateway_run(self.body)
        except Exception as exc:
            return _blocked_foreground_vedastro(
                reason="gateway_invocation_error",
                error_type=type(exc).__name__,
            )
        if not isinstance(result, dict):
            return _blocked_foreground_vedastro(reason="gateway_invocation_error")
        _remember_gateway_outcome(self.body, result)
        return result


def start_foreground_vedastro(handler, body: dict, *, defer_optional_external_evidence: bool) -> ForegroundVedastroSession:
    return ForegroundVedastroSession(
        handler,
        body,
        defer=defer_optional_external_evidence,
    ).start()


def finish_foreground_vedastro(session: ForegroundVedastroSession | None) -> dict:
    if session is None:
        return _blocked_foreground_vedastro(reason="foreground_optional_evidence_timeout")
    return session.finish()