python-codex 0.2.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

workspace_server/workspaces.py

@@ -9,7 +9,6 @@ from uuid import uuid4
 
 from pycodex.utils import uuid7_string
 
-
 SessionFactory = typing.Callable[[], object]
 
 
@@ -23,7 +22,7 @@ class WorkspaceStateStore:
             None if board_path is None else board_path.with_suffix(".pycodex-ws.json")
         )
 
-    def load_tabs(self) -> "typing.List[typing.Dict[str, str]]":
+    def load_tabs(self) -> "typing.List[typing.Dict[str, object]]":
         if self.path is None or not self.path.is_file():
             return []
 
@@ -44,11 +43,13 @@ class WorkspaceStateStore:
                 continue
             title = str(tab.get("title") or "").strip()
             rollout_path = str(tab.get("rollout_path") or "").strip()
-            if title or rollout_path:
-                result.append({"title": title, "rollout_path": rollout_path})
+            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, str]]") -> None:
+    def save_tabs(self, tabs: "typing.Iterable[typing.Dict[str, object]]") -> None:
         if self.path is None:
             return
 
@@ -56,28 +57,32 @@ class WorkspaceStateStore:
             {
                 "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"
+        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'
+    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]':
+    config_path: "typing.Union[str, Path]",
+) -> "typing.List[WorkspaceDefinition]":
     path = Path(config_path).expanduser().resolve()
     if not path.exists():
         return []
@@ -112,7 +117,9 @@ def load_workspace_definitions(
         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))
+            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(
@@ -127,7 +134,9 @@ def load_workspace_definitions(
         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)
+                    "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:
@@ -152,8 +161,8 @@ def load_workspace_definitions(
 
 
 def save_workspace_definitions(
-    config_path: 'typing.Union[str, Path]',
-    definitions: 'typing.Iterable[WorkspaceDefinition]',
+    config_path: "typing.Union[str, Path]",
+    definitions: "typing.Iterable[WorkspaceDefinition]",
 ) -> None:
     path = Path(config_path).expanduser().resolve()
     base_dir = path.parent
@@ -171,9 +180,9 @@ def save_workspace_definitions(
 
 
 def _workspace_definition_to_json(
-    definition: 'WorkspaceDefinition',
-    base_dir: 'Path',
-) -> 'typing.Dict[str, object]':
+    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),
@@ -188,7 +197,7 @@ def _workspace_definition_to_json(
     return result
 
 
-def _format_path_for_workspace_config(path: 'Path', base_dir: 'Path') -> 'str':
+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()))
@@ -201,14 +210,14 @@ def _format_path_for_workspace_config(path: 'Path', base_dir: 'Path') -> 'str':
     return relative
 
 
-def _resolve_workspace_path(value: str, base_dir: 'Path') -> 'Path':
+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:
+def _next_workspace_id(existing: "typing.Container[str]") -> str:
     index = 1
     while True:
         candidate = "workspace-{0}".format(index)
@@ -242,11 +251,16 @@ class WorkspaceSessionManager:
         self._sessions: "typing.Dict[str, object]" = {}
         self._session_order: "typing.List[str]" = []
         self._state_watchers: "typing.Dict[str, asyncio.Task]" = {}
-        self._persisted_titles: "typing.Dict[str, str]" = {}
-        self._lock = asyncio.Lock()
+        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]",
@@ -259,57 +273,89 @@ class WorkspaceSessionManager:
             await self.create_session()
             return
         for tab in state_tabs:
-            await self.create_session(
+            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())
-        self._sessions.clear()
-        self._session_order = []
-        self._state_watchers.clear()
-        self._persisted_titles.clear()
-        for watcher in watchers:
-            watcher.cancel()
-        if watchers:
-            await asyncio.gather(*watchers, return_exceptions=True)
-        for session in sessions:
-            await session.close()
+        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._lock:
+        async with self._get_lock():
             session_id = uuid7_string()
             session = self._session_factory()
             await session.start()
 
-            if rollout_path:
-                await session.restore_from_rollout(rollout_path, title=title)
+            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._persisted_titles[session_id] = str(
-                session_summary(session).get("title") or ""
-            )
-            self._state_watchers[session_id] = asyncio.create_task(
-                self._watch_session_title(session_id, session)
-            )
+            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._lock:
+        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._persisted_titles.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
             ]
@@ -319,20 +365,33 @@ class WorkspaceSessionManager:
         await session.close()
         self.persist_workspace_state()
 
-    async def _watch_session_title(self, session_id: str, session) -> None:
-        subscriber = session.subscribe()
+    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 event.get("type") != "title_changed":
+                if not isinstance(event, dict) or "snapshot" not in event:
                     continue
-                title = str(event.get("title") or "")
-                if title == self._persisted_titles.get(session_id, ""):
-                    continue
-                self._persisted_titles[session_id] = title
-                self.persist_workspace_state()
+                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)
 
@@ -342,12 +401,8 @@ class WorkspaceSessionManager:
             session = self._sessions.get(session_id)
             if session is None:
                 continue
-            summary = session_summary(session)
-            title = str(summary.get("title") or "").strip()
-            rollout_path = str(session.rollout_path() or "")
-            if not title and not rollout_path:
-                continue
-            tabs.append({"title": title, "rollout_path": rollout_path})
+            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()
@@ -389,14 +444,15 @@ class WorkspaceSessionManager:
 
 @dataclass(frozen=True)
 class WorkspaceEntry:
-    definition: 'WorkspaceDefinition'
-    manager: 'WorkspaceSessionManager'
+    definition: "WorkspaceDefinition"
+    manager: "WorkspaceSessionManager"
 
-    def to_dict(self) -> 'typing.Dict[str, object]':
+    def to_dict(self) -> "typing.Dict[str, object]":
         return {
             "id": self.definition.workspace_id,
             "board_path": (
-                "" if self.definition.board_path is None
+                ""
+                if self.definition.board_path is None
                 else str(self.definition.board_path)
             ),
             "work_dir": str(self.definition.work_dir),
@@ -413,13 +469,13 @@ WorkspaceEntryFactory = typing.Callable[
 class WorkspaceRegistry:
     def __init__(
         self,
-        entries: 'typing.Iterable[WorkspaceEntry]',
-        config_path: 'typing.Union[str, Path, None]' = None,
+        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._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()
@@ -430,7 +486,7 @@ class WorkspaceRegistry:
         if not self._entries and self._entry_factory is None:
             raise ValueError("at least one workspace is required")
 
-    def _register(self, entry: 'WorkspaceEntry') -> None:
+    def _register(self, entry: "WorkspaceEntry") -> None:
         workspace_id = entry.definition.workspace_id
         key = self._key_for_definition(entry.definition)
         if key in self._entries:
@@ -444,7 +500,7 @@ class WorkspaceRegistry:
         self._order.append(key)
         entry.manager.set_persist_callback(self.persist_definitions)
 
-    def _key_for_definition(self, definition: 'WorkspaceDefinition') -> str:
+    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())
@@ -459,7 +515,7 @@ class WorkspaceRegistry:
             await self._entries[key].manager.close()
         self._started = False
 
-    def get(self, workspace_id: str) -> 'WorkspaceEntry':
+    def get(self, workspace_id: str) -> "WorkspaceEntry":
         workspace_id = normalize_workspace_id(workspace_id)
         if not workspace_id:
             raise KeyError(workspace_id)
@@ -473,11 +529,13 @@ class WorkspaceRegistry:
         name: str,
         work_dir: str = "./",
         board: "typing.Union[str, bool, None]" = None,
-    ) -> 'WorkspaceEntry':
+    ) -> "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()
+        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)
 
@@ -493,7 +551,9 @@ class WorkspaceRegistry:
             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))
+                raise ValueError(
+                    "workspace board already exists: {0}".format(board_path)
+                )
 
         workspace_id = normalize_workspace_id(name)
         if not workspace_id:
@@ -513,7 +573,7 @@ class WorkspaceRegistry:
         self.persist_definitions()
         return entry
 
-    async def delete_workspace(self, name: str) -> 'WorkspaceDefinition':
+    async def delete_workspace(self, name: str) -> "WorkspaceDefinition":
         workspace_id = normalize_workspace_id(name)
         if not workspace_id:
             raise KeyError(workspace_id)
@@ -535,11 +595,8 @@ class WorkspaceRegistry:
             [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 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]":

pycodex/collaboration.py

@@ -1,20 +0,0 @@
-
-from .compat import Literal
-import typing
-
-CollaborationMode = Literal["default", "plan", "execute", "pair_programming"]
-
-DEFAULT_COLLABORATION_MODE: 'CollaborationMode' = "default"
-PLAN_COLLABORATION_MODE: 'CollaborationMode' = "plan"
-
-_MODE_DISPLAY_NAMES: 'typing.Dict[str, str]' = {
-    "default": "Default",
-    "plan": "Plan",
-    "execute": "Execute",
-    "pair_programming": "Pair Programming",
-}
-
-
-def collaboration_mode_display_name(mode: 'typing.Union[str, None]') -> 'str':
-    normalized = (mode or DEFAULT_COLLABORATION_MODE).strip().lower()
-    return _MODE_DISPLAY_NAMES.get(normalized, normalized.replace("_", " ").title())