Gitea-Tools/mcp_session_state.py

"""Durable MCP session validation state shared across daemon process pools (#559).

The IDE often routes sequential MCP tool calls to different daemon processes.
Session-scoped proofs (workflow load, review decision lock) must therefore
survive process boundaries while remaining fail-closed against spoofing.

Security notes (extends #211):
- Never store under host-global ``/tmp`` (world-writable, spoofable).
- Default root is ``~/.cache/gitea-tools/session-state`` (mode ``0o700``).
- Files are written atomically with mode ``0o600``.
- Records are keyed by remote + org + repo + profile identity, not by PID.
- TTL prevents indefinitely stale reuse across unrelated sessions.
"""

from __future__ import annotations

import fcntl
import json
import os
import re
import tempfile
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any

STATE_DIR_ENV = "GITEA_MCP_SESSION_STATE_DIR"
DEFAULT_STATE_DIR = os.path.expanduser("~/.cache/gitea-tools/session-state")
TTL_HOURS_ENV = "GITEA_MCP_SESSION_STATE_TTL_HOURS"
DEFAULT_TTL_HOURS = 4.0

KIND_WORKFLOW_LOAD = "review_workflow_load"
KIND_DECISION_LOCK = "review_decision_lock"
KIND_REVIEW_DRAFT = "review_draft"
# Durable marker set when a worker session attempts a direct stable-branch push
# or a root-checkout local commit (#671). Keyed per profile identity like the
# other session proofs; a contaminated session fails closed on gated mutations
# until a reconciler audits and clears it.
KIND_STABLE_BRANCH_CONTAMINATION = "stable_branch_contamination"



_SAFE_SEGMENT_RE = re.compile(r"[^A-Za-z0-9._+-]+")
SESSION_PROFILE_LOCK_ENV = "GITEA_SESSION_PROFILE_LOCK"


def default_state_dir() -> str:
    raw = (os.environ.get(STATE_DIR_ENV) or DEFAULT_STATE_DIR).strip()
    return raw or DEFAULT_STATE_DIR


def ttl_hours() -> float:
    raw = (os.environ.get(TTL_HOURS_ENV) or "").strip()
    if not raw:
        return DEFAULT_TTL_HOURS
    try:
        value = float(raw)
    except ValueError:
        return DEFAULT_TTL_HOURS
    return value if value > 0 else DEFAULT_TTL_HOURS


def _sanitize_segment(value: str) -> str:
    text = (value or "").strip()
    if not text:
        return "_"
    return _SAFE_SEGMENT_RE.sub("_", text)


def current_profile_identity(
    profile_name: str | None = None,
    session_profile_lock: str | None = None,
    profile_identity: str | None = None,
) -> str:
    """Resolve the session profile identity used as the durable key."""
    env_lock = (os.environ.get(SESSION_PROFILE_LOCK_ENV) or "").strip()
    explicit = (profile_identity or session_profile_lock or "").strip()
    lock = (explicit or env_lock or "").strip()
    name = (profile_name or "").strip()
    return lock or name or "unknown-profile"


def state_key(
    *,
    kind: str,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    profile_identity: str | None = None,
) -> str:
    """Build durable filename key.

    Session proofs are one-active-per-profile (workflow load / decision lock),
    so the primary key is kind + profile identity. Remote/org/repo are stored
    inside the payload and validated on load (#559), which lets a later daemon
    process recover state without already knowing the remote argument.
    """
    # Keep remote/org/repo parameters for API stability / future kinds; they are
    # intentionally not part of the filename for session-scoped proofs.
    _ = (remote, org, repo)
    return "-".join(
        _sanitize_segment(part)
        for part in (
            kind,
            profile_identity or "unknown-profile",
        )
    )


def state_file_path(
    *,
    kind: str,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    profile_identity: str | None = None,
    state_dir: str | None = None,
) -> str:
    root = (state_dir or default_state_dir()).strip()
    name = state_key(
        kind=kind,
        remote=remote,
        org=org,
        repo=repo,
        profile_identity=profile_identity,
    )
    return os.path.join(root, f"{name}.json")


def _ensure_state_dir(state_dir: str | None = None) -> str:
    root = (state_dir or default_state_dir()).strip()
    os.makedirs(root, mode=0o700, exist_ok=True)
    return root


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


def _parse_iso(value: str | None) -> datetime | None:
    text = (value or "").strip()
    if not text:
        return None
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        parsed = datetime.fromisoformat(text)
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


@contextmanager
def _exclusive_file_lock(lock_path: str):
    os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True)
    fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX)
        yield fd
    finally:
        try:
            fcntl.flock(fd, fcntl.LOCK_UN)
        finally:
            os.close(fd)


def _read_json(path: str) -> dict[str, Any] | None:
    if not path or not os.path.exists(path):
        return None
    try:
        with open(path, encoding="utf-8") as handle:
            data = json.load(handle)
    except (OSError, json.JSONDecodeError):
        return None
    return data if isinstance(data, dict) else None


def _write_json(path: str, data: dict[str, Any]) -> None:
    parent = os.path.dirname(path) or "."
    os.makedirs(parent, mode=0o700, exist_ok=True)
    payload = json.dumps(data, indent=2, sort_keys=True) + "\n"
    fd, temp_path = tempfile.mkstemp(prefix=".session-", suffix=".json", dir=parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temp_path, 0o600)
        os.replace(temp_path, path)
        try:
            os.chmod(path, 0o600)
        except OSError:
            pass
    finally:
        if os.path.exists(temp_path):
            try:
                os.remove(temp_path)
            except OSError:
                pass


def identity_match_reasons(
    record: dict[str, Any] | None,
    *,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    profile_identity: str | None = None,
) -> list[str]:
    """Return fail-closed reasons when durable record identity does not match."""
    if record is None:
        return []
    reasons: list[str] = []
    expected_profile = current_profile_identity(profile_identity=profile_identity)
    stored_profile = (
        (record.get("session_profile_lock") or record.get("profile_identity") or "")
        .strip()
    )
    if stored_profile and expected_profile and stored_profile != expected_profile:
        if expected_profile != "unknown-profile":
            reasons.append(
                "session state profile identity mismatch "
                f"(stored={stored_profile!r}, active={expected_profile!r}; fail closed)"
            )

    for field, expected in (
        ("remote", remote),
        ("org", org),
        ("repo", repo),
    ):
        want = (expected or "").strip()
        have = (str(record.get(field) or "")).strip()
        if want and have and want != have:
            reasons.append(
                f"session state {field} mismatch "
                f"(stored={have!r}, expected={want!r}; fail closed)"
            )

    recorded_at = _parse_iso(record.get("recorded_at") or record.get("updated_at"))
    if recorded_at is None:
        reasons.append("session state missing recorded_at timestamp (fail closed)")
    else:
        age = _now_utc() - recorded_at
        if age > timedelta(hours=ttl_hours()):
            reasons.append(
                f"session state expired after {ttl_hours():g}h (fail closed)"
            )
        if age < timedelta(0):
            reasons.append("session state recorded_at is in the future (fail closed)")
    return reasons


def load_state(
    *,
    kind: str,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    profile_identity: str | None = None,
    state_dir: str | None = None,
) -> dict[str, Any] | None:
    """Load durable state payload when identity checks pass."""
    profile = current_profile_identity(profile_identity=profile_identity)
    path = state_file_path(
        kind=kind,
        remote=remote,
        org=org,
        repo=repo,
        profile_identity=profile,
        state_dir=state_dir,
    )
    lock_path = f"{path}.lock"
    with _exclusive_file_lock(lock_path):
        envelope = _read_json(path)
    if not envelope:
        return None
    payload = envelope.get("payload")
    if not isinstance(payload, dict):
        return None
    # Identity fields live on both envelope and payload for convenience.
    merged = dict(payload)
    for key in (
        "kind",
        "remote",
        "org",
        "repo",
        "profile_identity",
        "session_profile_lock",
        "recorded_at",
        "updated_at",
        "writer_pid",
    ):
        if key in envelope and key not in merged:
            merged[key] = envelope[key]
    reasons = identity_match_reasons(
        merged,
        remote=remote,
        org=org,
        repo=repo,
        profile_identity=profile,
    )
    if reasons:
        return None
    return merged


def save_state(
    *,
    kind: str,
    payload: dict[str, Any] | None,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    profile_identity: str | None = None,
    state_dir: str | None = None,
) -> dict[str, Any] | None:
    """Persist or clear durable state for the given session identity key."""
    profile = current_profile_identity(
        profile_name=payload.get("session_profile") if payload else None,
        session_profile_lock=(
            (payload or {}).get("session_profile_lock") or profile_identity
        ),
    )
    # Prefer explicit args over payload fields for key location.
    key_remote = remote if remote is not None else (payload or {}).get("remote")
    key_org = org if org is not None else (payload or {}).get("org")
    key_repo = repo if repo is not None else (payload or {}).get("repo")

    root = _ensure_state_dir(state_dir)
    path = state_file_path(
        kind=kind,
        remote=key_remote,
        org=key_org,
        repo=key_repo,
        profile_identity=profile,
        state_dir=root,
    )
    lock_path = f"{path}.lock"
    with _exclusive_file_lock(lock_path):
        if payload is None:
            for candidate in (path, lock_path):
                if os.path.exists(candidate):
                    try:
                        os.remove(candidate)
                    except OSError:
                        pass
            return None

        now = _now_utc().isoformat().replace("+00:00", "Z")
        body = dict(payload)
        body.setdefault("session_pid", os.getpid())
        body["writer_pid"] = os.getpid()
        body["profile_identity"] = profile
        if not (body.get("session_profile_lock") or "").strip():
            body["session_profile_lock"] = profile
        body["recorded_at"] = body.get("recorded_at") or now
        body["updated_at"] = now
        if key_remote is not None:
            body["remote"] = key_remote
        if key_org is not None:
            body["org"] = key_org
        if key_repo is not None:
            body["repo"] = key_repo

        envelope = {
            "kind": kind,
            "remote": key_remote,
            "org": key_org,
            "repo": key_repo,
            "profile_identity": profile,
            "session_profile_lock": body.get("session_profile_lock"),
            "recorded_at": body["recorded_at"],
            "updated_at": body["updated_at"],
            "writer_pid": body["writer_pid"],
            "payload": body,
        }
        _write_json(path, envelope)
        return dict(body)


def clear_state(
    *,
    kind: str,
    remote: str | None = None,
    org: str | None = None,
    repo: str | None = None,
    profile_identity: str | None = None,
    state_dir: str | None = None,
) -> None:
    save_state(
        kind=kind,
        payload=None,
        remote=remote,
        org=org,
        repo=repo,
        profile_identity=profile_identity,
        state_dir=state_dir,
    )