python-codex 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

workspace_server/workspaces.py

@@ -1,628 +0,0 @@
-import asyncio
-import json
-import os
-import tempfile
-import typing
-from dataclasses import dataclass
-from pathlib import Path
-from uuid import uuid4
-
-from pycodex.utils import uuid7_string
-
-SessionFactory = typing.Callable[[], object]
-
-
-def default_board_path() -> Path:
-    return Path(tempfile.gettempdir()) / "pcws-{0}.html".format(uuid4().hex[:8])
-
-
-class WorkspaceStateStore:
-    def __init__(self, board_path: "typing.Union[Path, None]") -> None:
-        self.path = (
-            None if board_path is None else board_path.with_suffix(".pycodex-ws.json")
-        )
-
-    def load_tabs(self) -> "typing.List[typing.Dict[str, object]]":
-        if self.path is None or not self.path.is_file():
-            return []
-
-        try:
-            payload = json.loads(
-                self.path.read_text(encoding="utf-8", errors="replace") or "{}"
-            )
-        except (OSError, ValueError):
-            return []
-
-        tabs = payload.get("tabs") if isinstance(payload, dict) else None
-        if not isinstance(tabs, list):
-            return []
-
-        result = []
-        for tab in tabs:
-            if not isinstance(tab, dict):
-                continue
-            title = str(tab.get("title") or "").strip()
-            rollout_path = str(tab.get("rollout_path") or "").strip()
-            restored = {"title": title, "rollout_path": rollout_path}
-            if tab.get("fork") is True:
-                restored["fork"] = True
-            result.append(restored)
-        return result
-
-    def save_tabs(self, tabs: "typing.Iterable[typing.Dict[str, object]]") -> None:
-        if self.path is None:
-            return
-
-        state_tabs = [
-            {
-                "title": str(tab.get("title") or ""),
-                "rollout_path": str(tab.get("rollout_path") or ""),
-                **({"fork": True} if tab.get("fork") is True else {}),
-            }
-            for tab in tabs
-        ]
-        payload = (
-            json.dumps(
-                {"version": 1, "tabs": state_tabs},
-                ensure_ascii=False,
-                indent=2,
-            )
-            + "\n"
-        )
-        self.path.parent.mkdir(parents=True, exist_ok=True)
-        self.path.write_text(payload, encoding="utf-8")
-
-
-@dataclass(frozen=True)
-class WorkspaceDefinition:
-    workspace_id: "str"
-    board_path: "typing.Union[Path, None]"
-    work_dir: "Path"
-
-
-def load_workspace_definitions(
-    config_path: "typing.Union[str, Path]",
-) -> "typing.List[WorkspaceDefinition]":
-    path = Path(config_path).expanduser().resolve()
-    if not path.exists():
-        return []
-    try:
-        payload = json.loads(path.read_text(encoding="utf-8", errors="replace"))
-    except OSError as exc:
-        raise ValueError("unable to read workspace config: {0}".format(exc)) from exc
-    except ValueError as exc:
-        raise ValueError("workspace config must be valid JSON") from exc
-
-    if isinstance(payload, dict):
-        workspaces = payload.get("workspaces")
-    else:
-        workspaces = payload
-    if not isinstance(workspaces, list):
-        raise ValueError("workspace config must contain a workspace list")
-
-    result = []
-    seen_ids = set()
-    seen_boards = set()
-    for index, item in enumerate(workspaces, start=1):
-        if not isinstance(item, dict):
-            raise ValueError("workspace entry {0} must be an object".format(index))
-        workspace_id = normalize_workspace_id(str(item.get("id") or "").strip())
-        if not workspace_id:
-            workspace_id = _next_workspace_id(seen_ids)
-        if workspace_id in seen_ids:
-            raise ValueError("duplicate workspace name: {0}".format(workspace_id))
-        seen_ids.add(workspace_id)
-
-        work_dir_value = item.get("work_dir")
-        if work_dir_value is None:
-            work_dir_value = item.get("cwd")
-        if not str(work_dir_value or "").strip():
-            raise ValueError(
-                "workspace `{0}` is missing `work_dir`".format(workspace_id)
-            )
-        work_dir = _resolve_workspace_path(str(work_dir_value), path.parent)
-        if not work_dir.is_dir():
-            raise ValueError(
-                "workspace `{0}` work_dir does not exist: {1}".format(
-                    workspace_id,
-                    work_dir,
-                )
-            )
-
-        board_value = item.get("board")
-        board_path = None
-        if board_value is not False:
-            if not isinstance(board_value, str) or not board_value.strip():
-                raise ValueError(
-                    "workspace `{0}` requires a board path or false".format(
-                        workspace_id
-                    )
-                )
-            board_path = _resolve_workspace_path(board_value, path.parent)
-            if board_path in seen_boards:
-                raise ValueError("duplicate workspace board: {0}".format(board_path))
-            seen_boards.add(board_path)
-            if not board_path.parent.is_dir():
-                raise ValueError(
-                    "workspace `{0}` board parent directory does not exist: {1}".format(
-                        workspace_id,
-                        board_path.parent,
-                    )
-                )
-
-        result.append(
-            WorkspaceDefinition(
-                workspace_id=workspace_id,
-                board_path=board_path,
-                work_dir=work_dir,
-            )
-        )
-    return result
-
-
-def save_workspace_definitions(
-    config_path: "typing.Union[str, Path]",
-    definitions: "typing.Iterable[WorkspaceDefinition]",
-) -> None:
-    path = Path(config_path).expanduser().resolve()
-    base_dir = path.parent
-    payload = {
-        "workspaces": [
-            _workspace_definition_to_json(definition, base_dir)
-            for definition in definitions
-        ]
-    }
-    path.parent.mkdir(parents=True, exist_ok=True)
-    path.write_text(
-        json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
-        encoding="utf-8",
-    )
-
-
-def _workspace_definition_to_json(
-    definition: "WorkspaceDefinition",
-    base_dir: "Path",
-) -> "typing.Dict[str, object]":
-    result = {
-        "id": definition.workspace_id,
-        "work_dir": _format_path_for_workspace_config(definition.work_dir, base_dir),
-    }
-    if definition.board_path is not None:
-        result["board"] = _format_path_for_workspace_config(
-            definition.board_path,
-            base_dir,
-        )
-    else:
-        result["board"] = False
-    return result
-
-
-def _format_path_for_workspace_config(path: "Path", base_dir: "Path") -> "str":
-    resolved = path.resolve()
-    try:
-        relative = os.path.relpath(str(resolved), str(base_dir.resolve()))
-    except ValueError:
-        return str(resolved)
-    if relative == ".":
-        return "."
-    if relative.startswith("..") or os.path.isabs(relative):
-        return str(resolved)
-    return relative
-
-
-def _resolve_workspace_path(value: str, base_dir: "Path") -> "Path":
-    path = Path(str(value)).expanduser()
-    if not path.is_absolute():
-        path = base_dir / path
-    return path.resolve()
-
-
-def _next_workspace_id(existing: "typing.Container[str]") -> str:
-    index = 1
-    while True:
-        candidate = "workspace-{0}".format(index)
-        if candidate not in existing:
-            return candidate
-        index += 1
-
-
-def normalize_workspace_id(workspace_id: str) -> str:
-    text = str(workspace_id or "").strip()
-    if not text:
-        return ""
-    allowed = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-")
-    if any(char not in allowed for char in text):
-        raise ValueError(
-            "workspace id may only contain letters, numbers, underscore, and dash: {0}".format(
-                text
-            )
-        )
-    return text
-
-
-class WorkspaceSessionManager:
-    def __init__(
-        self,
-        session_factory: "SessionFactory",
-        board_path: "typing.Union[Path, None]" = None,
-        persist_callback: "typing.Union[typing.Callable[[], None], None]" = None,
-    ) -> None:
-        self._session_factory = session_factory
-        self._sessions: "typing.Dict[str, object]" = {}
-        self._session_order: "typing.List[str]" = []
-        self._state_watchers: "typing.Dict[str, asyncio.Task]" = {}
-        self._tab_states: "typing.Dict[str, typing.Dict[str, object]]" = {}
-        self._lock = None
-        self._state_store = WorkspaceStateStore(board_path)
-        self._persist_callback = persist_callback
-
-    def _get_lock(self):
-        if self._lock is None:
-            self._lock = asyncio.Lock()
-        return self._lock
-
-    def set_persist_callback(
-        self,
-        callback: "typing.Union[typing.Callable[[], None], None]",
-    ) -> None:
-        self._persist_callback = callback
-
-    async def start(self) -> None:
-        state_tabs = self._state_store.load_tabs()
-        if not state_tabs:
-            await self.create_session()
-            return
-        for tab in state_tabs:
-            await self._create_session(
-                title=str(tab.get("title") or ""),
-                rollout_path=str(tab.get("rollout_path") or ""),
-                fork=tab.get("fork") is True,
-            )
-        for session_id in self._session_order:
-            self._subscribe_session_state(session_id)
-        self.persist_workspace_state()
-
-    async def close(self) -> None:
-        sessions = list(self._sessions.values())
-        watchers = list(self._state_watchers.values())
-        try:
-            for session in sessions:
-                await session.close()
-            if watchers:
-                await asyncio.gather(*watchers)
-        finally:
-            for watcher in watchers:
-                watcher.cancel()
-            if watchers:
-                await asyncio.gather(*watchers, return_exceptions=True)
-            self._sessions.clear()
-            self._session_order = []
-            self._state_watchers.clear()
-            self._tab_states.clear()
-
-    async def create_session(
-        self,
-        title: str = "",
-        rollout_path: str = "",
-        fork: bool = False,
-    ) -> str:
-        session_id = await self._create_session(title, rollout_path, fork)
-        self._subscribe_session_state(session_id)
-        self.persist_workspace_state()
-        return session_id
-
-    async def _create_session(
-        self,
-        title: str = "",
-        rollout_path: str = "",
-        fork: bool = False,
-    ) -> str:
-        async with self._get_lock():
-            session_id = uuid7_string()
-            session = self._session_factory()
-            await session.start()
-
-            try:
-                if rollout_path or title:
-                    await session.restore_from_rollout(
-                        rollout_path, title=title, fork=fork
-                    )
-            except BaseException:
-                await session.close()
-                raise
-
-            self._sessions[session_id] = session
-            self._session_order.append(session_id)
-            self._tab_states[session_id] = {
-                "title": str(session_summary(session).get("title") or ""),
-                "rollout_path": rollout_path,
-                **({"fork": True} if fork else {}),
-            }
-            return session_id
-
-    def _subscribe_session_state(self, session_id: str) -> None:
-        session = self._sessions[session_id]
-        subscriber = session.subscribe()
-        self._state_watchers[session_id] = asyncio.create_task(
-            self._watch_session_state(session_id, session, subscriber)
-        )
-
-    async def close_session(self, session_id: str) -> None:
-        async with self._get_lock():
-            if len(self._session_order) <= 1:
-                raise ValueError("cannot close the last session")
-            session = self._sessions.pop(session_id, None)
-            if session is None:
-                raise KeyError(session_id)
-            watcher = self._state_watchers.pop(session_id, None)
-            self._tab_states.pop(session_id, None)
-            self._session_order = [
-                item for item in self._session_order if item != session_id
-            ]
-        if watcher is not None:
-            watcher.cancel()
-            await asyncio.gather(watcher, return_exceptions=True)
-        await session.close()
-        self.persist_workspace_state()
-
-    async def _watch_session_state(self, session_id: str, session, subscriber) -> None:
-        try:
-            while True:
-                event = await subscriber.get()
-                if event is None:
-                    return
-                if not isinstance(event, dict) or "snapshot" not in event:
-                    continue
-                reason = (
-                    event["payload"]["reason"]
-                    if event.get("kind") == "session_state"
-                    else None
-                )
-                previous = self._tab_states[session_id]
-                state = dict(previous)
-                snapshot = event["snapshot"]
-                state["title"] = str(snapshot.get("title") or "")
-                if reason == "identity" and state["rollout_path"]:
-                    # Recreate an unwritten fork from its recorded ancestor.
-                    state["fork"] = True
-                path = str(snapshot["recorded_rollout_path"] or "")
-                if path or not snapshot["rollout_path"]:
-                    state["rollout_path"] = path
-                    state.pop("fork", None)
-                self._tab_states[session_id] = state
-                if state != previous:
-                    self.persist_workspace_state()
-        finally:
-            session.unsubscribe(subscriber)
-
-    def persist_workspace_state(self) -> None:
-        tabs = []
-        for session_id in self._session_order:
-            session = self._sessions.get(session_id)
-            if session is None:
-                continue
-            state = self._tab_states[session_id]
-            tabs.append(state)
-        self._state_store.save_tabs(tabs)
-        if self._persist_callback is not None:
-            self._persist_callback()
-
-    def get(self, session_id: "typing.Union[str, None]" = None) -> object:
-        resolved_id = self.resolve_session_id(session_id)
-        try:
-            return self._sessions[resolved_id]
-        except KeyError:
-            raise KeyError(resolved_id)
-
-    def resolve_session_id(self, session_id: "typing.Union[str, None]" = None) -> str:
-        if session_id:
-            return str(session_id)
-        if not self._session_order:
-            raise KeyError("no sessions")
-        return self._session_order[0]
-
-    def list_sessions(self) -> "typing.List[typing.Dict[str, object]]":
-        result = []
-        for session_id in self._session_order:
-            session = self._sessions[session_id]
-            summary = session_summary(session)
-            result.append(
-                {
-                    "id": session_id,
-                    "title": summary.get("title") or "pycodex",
-                    "running": bool(summary.get("running")),
-                    "spinner": summary.get("spinner") or "",
-                    "turn_count": summary.get("turn_count") or 0,
-                    "last_assistant": summary.get("last_assistant") or "",
-                    "context_remaining_percent": summary.get(
-                        "context_remaining_percent"
-                    ),
-                }
-            )
-        return result
-
-
-@dataclass(frozen=True)
-class WorkspaceEntry:
-    definition: "WorkspaceDefinition"
-    manager: "WorkspaceSessionManager"
-
-    def to_dict(self) -> "typing.Dict[str, object]":
-        return {
-            "id": self.definition.workspace_id,
-            "board_path": (
-                ""
-                if self.definition.board_path is None
-                else str(self.definition.board_path)
-            ),
-            "work_dir": str(self.definition.work_dir),
-            "sessions": self.manager.list_sessions(),
-        }
-
-
-WorkspaceEntryFactory = typing.Callable[
-    [WorkspaceDefinition, "typing.Union[typing.Callable[[], None], None]"],
-    WorkspaceEntry,
-]
-
-
-class WorkspaceRegistry:
-    def __init__(
-        self,
-        entries: "typing.Iterable[WorkspaceEntry]",
-        config_path: "typing.Union[str, Path, None]" = None,
-        entry_factory: "typing.Union[WorkspaceEntryFactory, None]" = None,
-    ) -> None:
-        self._entries: "typing.Dict[str, WorkspaceEntry]" = {}
-        self._name_to_key: "typing.Dict[str, str]" = {}
-        self._order: "typing.List[str]" = []
-        self._started = False
-        self._config_path = (
-            None if config_path is None else Path(config_path).expanduser().resolve()
-        )
-        self._entry_factory = entry_factory
-        for entry in entries:
-            self._register(entry)
-        if not self._entries and self._entry_factory is None:
-            raise ValueError("at least one workspace is required")
-
-    def _register(self, entry: "WorkspaceEntry") -> None:
-        workspace_id = entry.definition.workspace_id
-        key = self._key_for_definition(entry.definition)
-        if key in self._entries:
-            raise ValueError(
-                "duplicate workspace board: {0}".format(entry.definition.board_path)
-            )
-        if workspace_id in self._name_to_key:
-            raise ValueError("duplicate workspace name: {0}".format(workspace_id))
-        self._entries[key] = entry
-        self._name_to_key[workspace_id] = key
-        self._order.append(key)
-        entry.manager.set_persist_callback(self.persist_definitions)
-
-    def _key_for_definition(self, definition: "WorkspaceDefinition") -> str:
-        if definition.board_path is None:
-            return "workspace:{0}".format(definition.workspace_id)
-        return str(definition.board_path.resolve())
-
-    async def start(self) -> None:
-        self._started = True
-        for key in self._order:
-            await self._entries[key].manager.start()
-
-    async def close(self) -> None:
-        for key in reversed(self._order):
-            await self._entries[key].manager.close()
-        self._started = False
-
-    def get(self, workspace_id: str) -> "WorkspaceEntry":
-        workspace_id = normalize_workspace_id(workspace_id)
-        if not workspace_id:
-            raise KeyError(workspace_id)
-        try:
-            return self._entries[self._name_to_key[workspace_id]]
-        except KeyError:
-            raise KeyError(workspace_id)
-
-    async def add_workspace(
-        self,
-        name: str,
-        work_dir: str = "./",
-        board: "typing.Union[str, bool, None]" = None,
-    ) -> "WorkspaceEntry":
-        if self._entry_factory is None:
-            raise ValueError("workspace creation is unavailable")
-
-        base_dir = (
-            self._config_path.parent if self._config_path is not None else Path.cwd()
-        )
-        resolved_work_dir = _resolve_workspace_path(str(work_dir or "./"), base_dir)
-        resolved_work_dir.mkdir(parents=True, exist_ok=True)
-
-        board_path = None
-        if board is not False:
-            if board is not None and not isinstance(board, str):
-                raise ValueError("board must be a path string or false")
-            board_path = (
-                _resolve_workspace_path(board, base_dir)
-                if board and board.strip()
-                else default_board_path()
-            )
-            board_path.parent.mkdir(parents=True, exist_ok=True)
-            board_key = str(board_path.resolve())
-            if board_key in self._entries:
-                raise ValueError(
-                    "workspace board already exists: {0}".format(board_path)
-                )
-
-        workspace_id = normalize_workspace_id(name)
-        if not workspace_id:
-            workspace_id = _next_workspace_id(self._name_to_key)
-        if workspace_id in self._name_to_key:
-            raise ValueError("workspace name already exists: {0}".format(workspace_id))
-
-        definition = WorkspaceDefinition(
-            workspace_id=workspace_id,
-            board_path=board_path,
-            work_dir=resolved_work_dir,
-        )
-        entry = self._entry_factory(definition, self.persist_definitions)
-        self._register(entry)
-        if self._started:
-            await entry.manager.start()
-        self.persist_definitions()
-        return entry
-
-    async def delete_workspace(self, name: str) -> "WorkspaceDefinition":
-        workspace_id = normalize_workspace_id(name)
-        if not workspace_id:
-            raise KeyError(workspace_id)
-        try:
-            key = self._name_to_key.pop(workspace_id)
-            entry = self._entries.pop(key)
-        except KeyError:
-            raise KeyError(workspace_id)
-        self._order = [item for item in self._order if item != key]
-        await entry.manager.close()
-        self.persist_definitions()
-        return entry.definition
-
-    def persist_definitions(self) -> None:
-        if self._config_path is None:
-            return
-        save_workspace_definitions(
-            self._config_path,
-            [self._entries[key].definition for key in self._order],
-        )
-
-    def list_workspaces(self) -> "typing.List[typing.Dict[str, object]]":
-        return [self._entries[key].to_dict() for key in self._order]
-
-
-def session_snapshot(session) -> "typing.Dict[str, object]":
-    return typing.cast("typing.Dict[str, object]", session.snapshot())
-
-
-def session_summary(session) -> "typing.Dict[str, object]":
-    summary = getattr(session, "summary", None)
-    if callable(summary):
-        return typing.cast("typing.Dict[str, object]", summary())
-    snapshot = session_snapshot(session)
-    return {
-        "title": snapshot.get("title") or "",
-        "running": bool(snapshot.get("running")),
-        "spinner": snapshot.get("spinner") or "",
-        "turn_count": len(snapshot.get("turns") or []),
-        "last_assistant": _last_assistant_text(snapshot.get("turns") or []),
-        "context_remaining_percent": snapshot.get("context_remaining_percent"),
-    }
-
-
-def _last_assistant_text(turns: "typing.Iterable[typing.Dict[str, object]]") -> str:
-    for turn in reversed(list(turns)):
-        if str(turn.get("kind") or "assistant") == "control":
-            continue
-        response = str(turn.get("response") or "").strip()
-        if response:
-            return response
-    return ""