openworker/coworker/personas/registry.py

"""Persona registry — the installed personas + their lifecycle state.

Unifies two sources behind one `id → Agent` resolver: the core surfaces (Code / Chat /
Cowork) wrap their existing agent builders (exact prompts preserved), and markdown manifests
(Ops today; third-party dirs in Phase 2) load through ``PersonaManifest``. Lifecycle —
installed → enabled → surfaced, plus a default — is persisted to a small JSON file.

A session is born from exactly one persona (recorded as ``SessionRecord.agent``); resolving an
id always returns its Agent even if the persona was later disabled, so live sessions keep
working. Disable/surface only affect what the *new-session* picker offers.
"""

from __future__ import annotations

import json
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

from ..agents.base import Agent
from ..agents.chat import chat_agent
from ..agents.code import CODE_CAPABILITIES, code_agent
from ..agents.cowork import COWORK_CAPABILITIES, cowork_agent
from .manifest import PersonaManifest, load_manifest_file

DEFAULT_PERSONA_ID = "cowork"


@dataclass
class PersonaState:
    enabled: bool = True
    surfaced: bool = True


@dataclass
class PersonaEntry:
    id: str
    name: str
    icon: str = ""
    tagline: str = ""
    needs_workspace: bool = True
    builtin: bool = True
    family: str = "knowledge"
    # The persona's workspace requirement (git|project|deliverable|none) — surfaced to the GUI so it
    # can detect project-scoped personas (git/project) uniformly. Manifest-backed personas carry it
    # verbatim; builtins set it at registration to match their family/needs_workspace.
    workspace: str = "deliverable"
    tools: list[str] = field(default_factory=list)
    default_surfaced: bool = (
        True  # whether it shows in the picker before any user choice
    )
    # Whether it ships enabled before any user choice. Builtins default on (UX-029: the
    # composer picker is their front door) — except retired ones (Chat: Coworker covers
    # quick Q&A). Installed third-party personas always start disabled pending consent.
    default_enabled: bool = True
    _builder: Optional[Callable[[], Agent]] = None
    manifest: Optional[PersonaManifest] = None

    def agent(self) -> Agent:
        if self._builder is not None:
            return self._builder()
        assert self.manifest is not None
        return self.manifest.to_agent()


class PersonaRegistry:
    def __init__(
        self,
        *,
        builtin_dir: Optional[str | Path] = None,
        extra_dirs: Optional[list[str | Path]] = None,
        state_path: Optional[str | Path] = None,
        installed_dir: Optional[str | Path] = None,
    ) -> None:
        self.state_path = Path(state_path) if state_path else None
        # Managed area where installed personas are *snapshotted* (copied) at install time, so a
        # persona's definition is stable and self-contained — independent of the user's source dir.
        if installed_dir is not None:
            self.installed_dir: Optional[Path] = Path(installed_dir)
        elif self.state_path is not None:
            self.installed_dir = self.state_path.parent / "personas-installed"
        else:
            self.installed_dir = None
        self._entries: dict[str, PersonaEntry] = {}
        self._enabled: dict[str, bool] = {}
        self._surfaced: dict[str, bool] = {}
        # Sharing v1 (OPE-7): install provenance per installed persona —
        # {version, source, installed_at} — drives the "replaces vN" note on re-install.
        self._installed_meta: dict[str, dict] = {}
        self._default = DEFAULT_PERSONA_ID
        self._load_builtin(builtin_dir)
        for d in extra_dirs or []:
            self._load_dir(d, builtin=False)
        self._load_state()
        self._load_installed()  # re-load snapshots from prior installs

    # -- loading ----------------------------------------------------------------
    def _register_builder(
        self,
        id,
        name,
        icon,
        tagline,
        builder,
        needs_workspace,
        family,
        tools,
        workspace="deliverable",
        default_surfaced=True,
        default_enabled=True,
    ) -> None:
        self._entries[id] = PersonaEntry(
            id=id,
            name=name,
            icon=icon,
            tagline=tagline,
            needs_workspace=needs_workspace,
            builtin=True,
            family=family,
            workspace=workspace,
            tools=list(tools),
            default_surfaced=default_surfaced,
            default_enabled=default_enabled,
            _builder=builder,
        )

    def _load_builtin(self, builtin_dir: Optional[str | Path]) -> None:
        # Core surfaces keep their exact prompts via the existing builders. Cowork (the default)
        # leads; Chat is RETIRED (owner call 2026-08-11: Coworker covers quick Q&A) — it ships
        # disabled and unsurfaced, recoverable from Settings ▸ Coworkers.
        self._register_builder(
            "cowork",
            "OpenWorker",
            "cowork",
            "Produce a deliverable — research, analysis, scripts",
            cowork_agent,
            True,
            "knowledge",
            COWORK_CAPABILITIES,
            workspace="deliverable",
        )
        self._register_builder(
            "code",
            "Code",
            "code",
            "Work in a codebase — files, git, shell",
            code_agent,
            True,
            "code",
            CODE_CAPABILITIES,
            workspace="git",
        )
        self._register_builder(
            "chat",
            "Chat",
            "chat",
            "Quick questions — no workspace",
            chat_agent,
            False,
            "knowledge",
            [],
            workspace="none",
            default_surfaced=False,
            default_enabled=False,
        )
        # Markdown-backed built-ins (Ops, …) — dogfood the manifest path.
        d = Path(builtin_dir) if builtin_dir else Path(__file__).parent / "builtin"
        self._load_dir(d, builtin=True)

    def _load_dir(self, directory: str | Path, *, builtin: bool) -> None:
        d = Path(directory)
        if not d.is_dir():
            return
        for md in sorted(d.glob("*.md")):
            self._register_manifest(
                load_manifest_file(md, builtin=builtin), builtin=builtin
            )
        # Bundle subdirs (OPE-58): <dir>/<id>/manifest.md with an optional sibling
        # skills/ folder — the same self-contained shape an install snapshot uses, so a
        # persona's skills live with it instead of leaking into a shared flat dir.
        for sub in sorted(p for p in d.iterdir() if p.is_dir()):
            md = sub / "manifest.md"
            if md.is_file():
                self._register_manifest(
                    load_manifest_file(md, builtin=builtin), builtin=builtin
                )

    def _register_manifest(self, m, *, builtin: bool) -> None:
        self._entries[m.id] = PersonaEntry(
            id=m.id,
            name=m.name,
            icon=m.icon,
            tagline=m.tagline,
            needs_workspace=m.needs_workspace,
            builtin=builtin,
            family=m.family,
            workspace=m.workspace,
            tools=list(m.tools),
            manifest=m,
            # Team workers never surface in the picker: they are purpose-built to be
            # STAFFED by a lead, not started solo (their prompts talk to a lead, not
            # a human). They stay enabled so the staffing gate can resolve them.
            default_surfaced=m.team != "worker",
        )

    def _load_installed(self) -> None:
        if not (self.installed_dir and self.installed_dir.is_dir()):
            return
        for sub in sorted(self.installed_dir.iterdir()):
            if sub.is_dir():
                self._load_dir(sub, builtin=False)

    def _load_state(self) -> None:
        if self.state_path and self.state_path.is_file():
            data = json.loads(self.state_path.read_text(encoding="utf-8"))
            self._enabled = dict(data.get("enabled", {}))
            self._surfaced = dict(data.get("surfaced", {}))
            self._installed_meta = dict(data.get("installed_meta", {}))
            self._default = data.get("default", DEFAULT_PERSONA_ID)

    def save(self) -> None:
        if not self.state_path:
            return
        self.state_path.parent.mkdir(parents=True, exist_ok=True)
        self.state_path.write_text(
            json.dumps(
                {
                    "enabled": self._enabled,
                    "surfaced": self._surfaced,
                    "installed_meta": self._installed_meta,
                    "default": self._default,
                },
                indent=2,
            ),
            encoding="utf-8",
        )

    # -- queries ----------------------------------------------------------------
    def ids(self) -> list[str]:
        return list(self._entries)

    def get(self, persona_id: str) -> Optional[PersonaEntry]:
        return self._entries.get(persona_id)

    def is_enabled(self, persona_id: str) -> bool:
        # Explicit state (either way) always wins. Absent a user choice, the entry's
        # default applies: builtins ship enabled — the composer picker is their front door
        # (UX-029, supersedes the 2026-07-09 Coworker-only default that fit the old hidden
        # ▾ menu) — except retired ones (Chat). Installed third-party personas stay
        # disabled until the user consents from the risk screen.
        if persona_id in self._enabled:
            return bool(self._enabled[persona_id])
        entry = self._entries.get(persona_id)
        if entry is not None and entry.builtin:
            return entry.default_enabled
        return persona_id == self._default or persona_id == DEFAULT_PERSONA_ID

    def is_surfaced(self, persona_id: str) -> bool:
        # User choice wins; otherwise the persona's default (Chat defaults hidden).
        if persona_id in self._surfaced:
            return self._surfaced[persona_id]
        entry = self._entries.get(persona_id)
        return entry.default_surfaced if entry else True

    def default_id(self) -> str:
        # The configured default if it's enabled, else cowork if present, else any enabled one.
        if self._default in self._entries and self.is_enabled(self._default):
            return self._default
        if DEFAULT_PERSONA_ID in self._entries and self.is_enabled(DEFAULT_PERSONA_ID):
            return DEFAULT_PERSONA_ID
        for pid in self._entries:
            if self.is_enabled(pid):
                return pid
        return DEFAULT_PERSONA_ID

    def agent(self, persona_id: Optional[str]) -> Agent:
        """Resolve a persona id to its Agent. Unknown ids fall back to the default persona;
        a known-but-disabled id still resolves (live sessions keep working)."""
        entry = self._entries.get(persona_id or "")
        if entry is None:
            entry = self._entries.get(self.default_id())
        if entry is None:
            raise KeyError(f"no persona to resolve for {persona_id!r}")
        return entry.agent()

    def sidebar(self) -> list[dict]:
        """Session surfaces for the new-session picker: enabled AND surfaced, in order."""
        out = []
        for e in self._entries.values():
            if self.is_enabled(e.id) and self.is_surfaced(e.id):
                out.append(
                    {
                        "name": e.id,
                        "title": e.name,
                        "needs_workspace": e.needs_workspace,
                        "icon": e.icon,
                        "tagline": e.tagline,
                        "default": e.id == self.default_id(),
                    }
                )
        return out

    def list_all(self) -> list[dict]:
        """Every installed persona + its lifecycle state — for the Personas settings panel."""
        return [
            {
                "id": e.id,
                "name": e.name,
                "icon": e.icon,
                "tagline": e.tagline,
                "needs_workspace": e.needs_workspace,
                "builtin": e.builtin,
                "family": e.family,
                "workspace": e.workspace,
                "tools": e.tools,
                "enabled": self.is_enabled(e.id),
                "surfaced": self.is_surfaced(e.id),
                "default": e.id == self.default_id(),
                "version": e.manifest.version if e.manifest else "",
                "installed_at": self._installed_meta.get(e.id, {}).get("installed_at", ""),
            }
            for e in self._entries.values()
        ]

    # -- mutations --------------------------------------------------------------
    def set_enabled(self, persona_id: str, enabled: bool) -> None:
        if persona_id not in self._entries:
            raise KeyError(persona_id)
        self._enabled[persona_id] = bool(enabled)
        if enabled:
            # Enabling implies surfacing (installs land unsurfaced, and "enabled but
            # invisible in the picker" is never what a user just asked for). They can
            # still untick "In picker" afterwards to hide it.
            self._surfaced[persona_id] = True
        self.save()

    def set_surfaced(self, persona_id: str, surfaced: bool) -> None:
        if persona_id not in self._entries:
            raise KeyError(persona_id)
        self._surfaced[persona_id] = bool(surfaced)
        self.save()

    def set_default(self, persona_id: str) -> None:
        if persona_id not in self._entries:
            raise KeyError(persona_id)
        self._default = persona_id
        self._enabled[persona_id] = True  # a default must be enabled
        self.save()

    def uninstall(self, persona_id: str) -> None:
        """Remove an installed persona: registry entry, lifecycle state, and its snapshot
        dir. Built-ins can't be uninstalled (disable them instead). Live sessions born
        from it resolve to the default persona afterwards (same as any unknown id)."""
        entry = self._entries.get(persona_id)
        if entry is None:
            raise KeyError(persona_id)
        if entry.builtin:
            raise ValueError(f"{persona_id} is built-in and cannot be deleted")
        del self._entries[persona_id]
        self._enabled.pop(persona_id, None)
        self._surfaced.pop(persona_id, None)
        if self._default == persona_id:
            self._default = DEFAULT_PERSONA_ID
        if self.installed_dir is not None:
            snap = self.installed_dir / persona_id
            if snap.is_dir():
                shutil.rmtree(snap)
        self.save()

    # -- install (third-party personas) -----------------------------------------
    def install_from_dir(self, directory: str | Path) -> list[dict]:
        """Install persona(s) from a local directory by **snapshotting** their manifests into our
        managed area (so the definition is stable, independent of the source dir). Returns a
        consent summary per persona; each lands **disabled + unsurfaced** pending the user's
        consent — the caller enables them only after the user approves the declared capabilities.

        NOTE: re-installing an updated persona overwrites the snapshot; live sessions on it simply
        resume with the new prompt/tools. We accept that for now (see PERSONAS.md)."""
        from .loading import consent_summary

        d = Path(directory)
        if not d.is_dir():
            raise FileNotFoundError(f"not a directory: {d}")
        mds = sorted(d.glob("*.md"))
        if not mds:
            raise FileNotFoundError(f"no persona manifests (*.md) in {d}")

        summaries: list[dict] = []
        for md in mds:
            m = load_manifest_file(md, builtin=False)  # validate before snapshotting
            replaces = self._replaces_of(m)
            snapshot = self._snapshot(md, m.id)
            installed = load_manifest_file(snapshot, builtin=False) if snapshot else m
            self._register_manifest(installed, builtin=False)
            # Consent rules (sharing v1): a fresh install always lands disabled pending
            # consent. An UPDATE keeps the user's enabled state — unless its capability
            # set GREW, which is a new decision, never a silent upgrade.
            if replaces is None or replaces.get("capabilities_grew"):
                self._enabled[m.id] = False
                self._surfaced[m.id] = False
            self._installed_meta[m.id] = {
                "version": installed.version,
                "source": str(md),
                "installed_at": self._now_stamp(),
            }
            summary = consent_summary(installed)
            summary["replaces"] = replaces
            summaries.append(summary)
        self.save()
        return summaries

    @staticmethod
    def _now_stamp() -> str:
        from datetime import date

        return date.today().isoformat()

    def _replaces_of(self, incoming) -> Optional[dict]:
        """When re-installing an already-installed persona id: what the new copy
        replaces ({version, installed_at, capabilities_grew}), else None."""
        from .loading import capability_set

        existing = self._entries.get(incoming.id)
        if existing is None or existing.builtin or existing.manifest is None:
            return None
        meta = self._installed_meta.get(incoming.id, {})
        grew = bool(capability_set(incoming) - capability_set(existing.manifest))
        return {
            "version": meta.get("version") or existing.manifest.version or "",
            "installed_at": meta.get("installed_at", ""),
            "capabilities_grew": grew,
        }

    def export_persona(self, persona_id: str, dest_dir: str | Path) -> dict:
        """Sharing v1 export: zip the persona's bundle (manifest + skills/) into
        ``dest_dir``. The zip's contents ARE the import format — extract or point the
        installer at it and the round trip is lossless."""
        import zipfile

        entry = self._entries.get(persona_id)
        if entry is None or entry.manifest is None or not entry.manifest.source:
            return {"ok": False, "error": "this coworker has no shareable bundle"}
        src_md = Path(entry.manifest.source)
        if not src_md.is_file():
            return {"ok": False, "error": "the coworker's bundle files are missing"}
        dest = Path(dest_dir).expanduser()
        if not dest.is_dir():
            return {"ok": False, "error": "destination folder does not exist"}
        version = entry.manifest.version
        zip_name = f"{persona_id}-coworker{('-v' + version) if version else ''}.zip"
        zip_path = dest / zip_name
        skills_dir = src_md.parent / "skills"
        try:
            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
                zf.write(src_md, "manifest.md")
                if skills_dir.is_dir():
                    for p in sorted(skills_dir.rglob("*")):
                        if p.is_file():
                            zf.write(p, str(Path("skills") / p.relative_to(skills_dir)))
        except OSError as e:
            return {"ok": False, "error": f"could not write the archive: {e}"}
        return {"ok": True, "path": str(zip_path)}

    def install_from_zip(self, data: bytes, filename: str = "") -> list[dict]:
        """Install persona(s) from a shared bundle zip (the export format). The archive
        is extracted to a temp dir with a zip-slip guard, then installed like a local
        directory — landing disabled pending consent like every install."""
        import io
        import tempfile
        import zipfile

        with tempfile.TemporaryDirectory(prefix="ocw-persona-zip-") as tmp:
            root = Path(tmp)
            try:
                with zipfile.ZipFile(io.BytesIO(data)) as zf:
                    for info in zf.infolist():
                        name = info.filename
                        target = (root / name).resolve()
                        if not str(target).startswith(str(root.resolve())):
                            raise FileNotFoundError(f"unsafe path in archive: {name}")
                    zf.extractall(root)
            except zipfile.BadZipFile as e:
                raise FileNotFoundError(f"not a valid bundle archive: {e}") from e
            # Accept both layouts: files at the root, or a single wrapping folder
            # (how macOS zips a directory).
            candidates = [root, *[p for p in root.iterdir() if p.is_dir()]]
            for d in candidates:
                if list(d.glob("*.md")) or (d / "manifest.md").is_file():
                    return self.install_from_dir(d)
            raise FileNotFoundError(
                f"no persona manifest found in {filename or 'the archive'}"
            )

    def _snapshot(self, md: Path, persona_id: str) -> Optional[Path]:
        """Copy a manifest into the managed install area; return the snapshot path (or None if no
        managed area is configured, e.g. an ephemeral in-memory registry)."""
        if self.installed_dir is None:
            return None
        dest_dir = self.installed_dir / persona_id
        dest_dir.mkdir(parents=True, exist_ok=True)
        dest = dest_dir / "manifest.md"
        shutil.copy2(md, dest)
        # Bundle shape (OPE-58 / sharing v1): a `skills/` dir next to the manifest travels
        # with the snapshot, so a persona's skills stay stable independent of the source.
        src_skills = md.parent / "skills"
        if src_skills.is_dir():
            shutil.copytree(src_skills, dest_dir / "skills", dirs_exist_ok=True)
        return dest

    def install_from_git(
        self, url: str, *, cache_base: Optional[str | Path] = None, clone=None
    ) -> list[dict]:
        """Clone a persona repo and install its personas (disabled pending consent)."""
        from .loading import clone_persona_repo, git_clone

        base = (
            Path(cache_base)
            if cache_base
            else (
                (self.state_path.parent if self.state_path else Path.cwd())
                / "persona-cache"
            )
        )
        dest = clone_persona_repo(url, base, clone=clone or git_clone)
        return self.install_from_dir(dest)


# -- module singleton (used by agents.get_agent / list_agents) ------------------
_singleton: Optional[PersonaRegistry] = None


def get_registry() -> PersonaRegistry:
    global _singleton
    if _singleton is None:
        from ..secrets import state_dir

        _singleton = PersonaRegistry(state_path=state_dir() / "personas.json")
    return _singleton


def set_registry(registry: PersonaRegistry) -> None:
    """Install a registry as the process singleton (the manager does this with its data dir)."""
    global _singleton
    _singleton = registry