python-codex 0.1.14__py3-none-any.whl → 0.1.15__py3-none-any.whl

workspace_server/app.py

@@ -3,6 +3,7 @@ import asyncio
 import html
 import json
 import os
+import threading
 import tempfile
 from uuid import uuid4
 from dataclasses import asdict, is_dataclass
@@ -19,6 +20,10 @@ from pycodex.cli import build_agent, build_cli_queue, build_model, configure_log
 from pycodex.interactive_session import run_interactive_session
 from pycodex.model import DEFAULT_CODEX_CONFIG_PATH
 from pycodex.protocol import AgentEvent, ToolCall
+from pycodex.utils.session_persist import (
+    SessionRolloutRecorder,
+    load_resumed_session_path,
+)
 from pycodex.utils import uuid7_string
 from pycodex.utils.visualize import IDLE_LISTENING_STATUS, shorten_title, tool_summary
 import typing
@@ -35,6 +40,55 @@ JSONValue = typing.Union[
 ]
 
 
+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, str]]":
+        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()
+            if title or rollout_path:
+                result.append({"title": title, "rollout_path": rollout_path})
+        return result
+
+    def save_tabs(self, tabs: "typing.Iterable[typing.Dict[str, str]]") -> None:
+        if self.path is None:
+            return
+
+        state_tabs = [
+            {
+                "title": str(tab.get("title") or ""),
+                "rollout_path": str(tab.get("rollout_path") or ""),
+            }
+            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")
+
+
 def build_parser() -> "argparse.ArgumentParser":
     parser = argparse.ArgumentParser(
         prog="pycodex-ws",
@@ -109,8 +163,10 @@ def parse_target(
     return host, port, board_path
 
 
-SessionFactory = typing.Callable[[], "WorkspaceInteractiveSession"]
+SessionFactory = typing.Callable[[], object]
+ThreadedSessionFactory = typing.Callable[[], "WorkspaceInteractiveSession"]
 SESSION_CLOSE_TIMEOUT_SECONDS = 2.0
+SPINNER_STATUS_PREVIEW_LIMIT = 180
 
 
 class WebSessionView:
@@ -125,12 +181,21 @@ class WebSessionView:
         self._spinner_status = ""
         self._stream_buffer = ""
         self._closed = False
+        self._server_loop: "typing.Union[asyncio.AbstractEventLoop, None]" = None
+        self._worker_loop: "typing.Union[asyncio.AbstractEventLoop, None]" = None
+        self._lock = threading.RLock()
+
+    def attach_server_loop(self, loop: "asyncio.AbstractEventLoop") -> None:
+        self._server_loop = loop
+
+    def attach_worker_loop(self, loop: "asyncio.AbstractEventLoop") -> None:
+        self._worker_loop = loop
 
     async def submit(self, prompt: str) -> "typing.Dict[str, object]":
         prompt = str(prompt or "").strip()
         if not prompt:
             return {"ok": False, "error": "prompt is empty"}
-        await self._input_queue.put(prompt)
+        await self._put_input(prompt)
         await self._publish(
             {
                 "type": "input",
@@ -140,6 +205,18 @@ class WebSessionView:
         )
         return {"ok": True, "type": "submitted", "snapshot": self.snapshot()}
 
+    async def _put_input(self, item: object) -> None:
+        worker_loop = self._worker_loop
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+        if worker_loop is None or worker_loop is running_loop:
+            await self._input_queue.put(item)
+            return
+        future = asyncio.run_coroutine_threadsafe(self._input_queue.put(item), worker_loop)
+        await asyncio.wrap_future(future)
+
     async def poll_prompt(self, prompt: "typing.Union[str, None]" = None) -> "typing.Union[str, None]":
         del prompt
         if self._closed and self._input_queue.empty():
@@ -161,45 +238,51 @@ class WebSessionView:
         return str(item)
 
     def handle_event(self, event: "AgentEvent") -> None:
-        self._apply_runtime_event(event)
-        payload = {
-            "type": "event",
-            "kind": str(getattr(event, "kind", "")),
-            "turn_id": str(getattr(event, "turn_id", "")),
-            "payload": _json_safe(getattr(event, "payload", {})),
-            "snapshot": self.snapshot(),
-        }
-        if payload["kind"] == "tool_completed":
-            payload["summary"] = tool_summary(getattr(event, "payload", {}))
+        with self._lock:
+            self._apply_runtime_event(event)
+            payload = {
+                "type": "event",
+                "kind": str(getattr(event, "kind", "")),
+                "turn_id": str(getattr(event, "turn_id", "")),
+                "payload": _json_safe(getattr(event, "payload", {})),
+                "snapshot": self.snapshot(),
+            }
+            if payload["kind"] == "tool_completed":
+                payload["summary"] = tool_summary(getattr(event, "payload", {}))
         self._publish_nowait(payload)
 
     def finish_stream(self) -> None:
-        if not self._stream_buffer:
-            return
-        active_turn = self._last_active_turn()
-        if active_turn is not None and not active_turn.get("response"):
-            active_turn["response"] = self._stream_buffer
-            active_turn["thinking"] = ""
-            active_turn["_thinking_active"] = False
-        self._stream_buffer = ""
-        self._publish_nowait({"type": "snapshot", "snapshot": self.snapshot()})
+        with self._lock:
+            if not self._stream_buffer:
+                return
+            active_turn = self._last_active_turn()
+            if active_turn is not None and not active_turn.get("response"):
+                active_turn["response"] = self._stream_buffer
+                active_turn["thinking"] = ""
+                active_turn["_thinking_active"] = False
+            self._stream_buffer = ""
+            event = {"type": "snapshot", "snapshot": self.snapshot()}
+        self._publish_nowait(event)
 
     def write_line(self, text: str) -> None:
-        text = str(text or "")
-        turn = self._new_control_turn(text)
-        turn["response"] = text
-        turn["status"] = "completed"
-        self._publish_nowait({"type": "snapshot", "snapshot": self.snapshot()})
+        with self._lock:
+            text = str(text or "")
+            turn = self._new_control_turn(text)
+            turn["response"] = text
+            turn["status"] = "completed"
+            event = {"type": "snapshot", "snapshot": self.snapshot()}
+        self._publish_nowait(event)
 
     def show_error(self, text: str) -> None:
         self.finish_stream()
-        turn = self._new_control_turn("")
-        turn["error"] = str(text or "")
-        turn["status"] = "error"
-        self._publish_nowait({"type": "snapshot", "snapshot": self.snapshot()})
+        with self._lock:
+            turn = self._new_control_turn("")
+            turn["error"] = str(text or "")
+            turn["status"] = "error"
+            event = {"type": "snapshot", "snapshot": self.snapshot()}
+        self._publish_nowait(event)
 
     def show_history(self) -> None:
-        self.finish_stream()
         assistant_turns = [turn for turn in self._turns if turn.get("kind") != "control"]
         if not assistant_turns:
             self.write_line("No history yet.")
@@ -214,17 +297,23 @@ class WebSessionView:
         self.write_line("\n".join(lines))
 
     def show_title(self) -> None:
-        self.finish_stream()
         self.write_line("Session: {0}".format(self._title or "untitled"))
 
     def set_session_title(self, title: str) -> None:
-        self.finish_stream()
-        self._title = str(title or "").strip()
-        self._publish_nowait({"type": "snapshot", "snapshot": self.snapshot()})
+        with self._lock:
+            self._set_title(title)
+            event = {
+                "type": "title_changed",
+                "title": self._title,
+                "snapshot": self.snapshot(),
+            }
+        self._publish_nowait(event)
 
     def show_resumed_session(self, title: str) -> None:
-        self._title = str(title or "")
-        self._publish_nowait({"type": "snapshot", "snapshot": self.snapshot()})
+        with self._lock:
+            self._set_title(title)
+            event = {"type": "snapshot", "snapshot": self.snapshot()}
+        self._publish_nowait(event)
 
     def load_session_history(
         self,
@@ -232,26 +321,27 @@ class WebSessionView:
         history: "typing.Iterable[typing.Tuple[str, str]]",
     ) -> None:
         self.finish_stream()
-        self._title = title or ""
-        self._turns = []
-        self._turns_by_submission_id = {}
-        self._turns_by_turn_id = {}
-        self._events = []
-        for prompt, response in history:
-            submission_id = uuid7_string()
-            turn = self._ensure_turn(submission_id, submission_id, str(prompt or ""))
-            turn["response"] = str(response or "")
-            turn["status"] = "completed"
-            turn["queue"] = "history"
-            turn["sender"] = "resume"
-        self._publish_nowait({"type": "snapshot", "snapshot": self.snapshot()})
+        with self._lock:
+            self._set_title(title)
+            self._turns = []
+            self._turns_by_submission_id = {}
+            self._turns_by_turn_id = {}
+            self._events = []
+            for prompt, response in history:
+                submission_id = uuid7_string()
+                turn = self._ensure_turn(submission_id, submission_id, str(prompt or ""))
+                turn["response"] = str(response or "")
+                turn["status"] = "completed"
+                turn["queue"] = "history"
+                turn["sender"] = "resume"
+            event = {"type": "snapshot", "snapshot": self.snapshot()}
+        self._publish_nowait(event)
 
     def show_steer_queued(self, turn_id: str, prompt: str) -> None:
-        self.write_line("[steer] queued: {0}".format(shorten_title(prompt, limit=72)))
+        del turn_id, prompt
 
     def schedule_steer_inserted(self, turn_id: str, prompt: str) -> None:
-        del turn_id
-        self.write_line("[steer] inserted: {0}".format(shorten_title(prompt, limit=72)))
+        del turn_id, prompt
 
     def set_context_window_tokens(
         self,
@@ -261,36 +351,52 @@ class WebSessionView:
 
     def subscribe(self) -> "asyncio.Queue":
         queue: "asyncio.Queue" = asyncio.Queue()
-        self._subscribers.add(queue)
-        queue.put_nowait(
-            {
+        with self._lock:
+            self._subscribers.add(queue)
+            event = {
                 "type": "hello",
                 "events": list(self._events[-200:]),
                 "snapshot": self.snapshot(),
             }
-        )
+        queue.put_nowait(event)
         return queue
 
     def unsubscribe(self, queue: "asyncio.Queue") -> None:
-        self._subscribers.discard(queue)
+        with self._lock:
+            self._subscribers.discard(queue)
 
     def close(self) -> None:
-        self._closed = True
-        self._input_queue.put_nowait(None)
-        for subscriber in tuple(self._subscribers):
-            subscriber.put_nowait(None)
-        self._subscribers.clear()
+        with self._lock:
+            self._closed = True
+            subscribers = tuple(self._subscribers)
+            self._subscribers.clear()
+        worker_loop = self._worker_loop
+        if worker_loop is None:
+            self._input_queue.put_nowait(None)
+        else:
+            asyncio.run_coroutine_threadsafe(self._input_queue.put(None), worker_loop)
+        self._publish_to_queues(subscribers, None)
 
     def snapshot(self) -> "typing.Dict[str, object]":
-        return {
-            "running": bool(self._spinner_status),
-            "status": self._spinner_status,
-            "status_kind": "spinner" if self._spinner_status else "idle",
-            "spinner": self._spinner_status,
-            "model": "pycodex",
-            "title": self._title,
-            "turns": [_public_turn(turn) for turn in self._turns[-80:]],
-        }
+        with self._lock:
+            return {
+                "running": bool(self._spinner_status),
+                "status": self._spinner_status,
+                "status_kind": "spinner" if self._spinner_status else "idle",
+                "spinner": self._spinner_status,
+                "model": "pycodex",
+                "title": self._title,
+                "turns": [_public_turn(turn) for turn in self._turns[-80:]],
+            }
+
+    def summary(self) -> "typing.Dict[str, object]":
+        with self._lock:
+            return {
+                "running": bool(self._spinner_status),
+                "spinner": self._spinner_status,
+                "title": self._title,
+                "turn_count": len(self._turns),
+            }
 
     def _apply_runtime_event(self, event: "AgentEvent") -> None:
         kind = str(getattr(event, "kind", "") or "")
@@ -300,7 +406,7 @@ class WebSessionView:
         turn_id = str(payload.get("turn_id") or getattr(event, "turn_id", "") or "")
         submission_id = str(payload.get("submission_id") or turn_id or "")
         turn = self._turns_by_submission_id.get(submission_id)
-        if turn is None and turn_id:
+        if turn is None and turn_id and not submission_id:
             turn = self._turns_by_turn_id.get(turn_id)
 
         if kind == "turn_started":
@@ -309,7 +415,7 @@ class WebSessionView:
                 str(item) for item in payload.get("user_texts", []) or []
             )
             if not self._title and str(prompt or "").strip():
-                self._title = shorten_title(str(prompt or ""))
+                self._set_title(shorten_title(str(prompt or "")))
             turn = self._ensure_turn(submission_id, turn_id, str(prompt or ""))
             turn["status"] = "running"
             turn["thinking"] = ""
@@ -340,11 +446,7 @@ class WebSessionView:
 
         if kind == "tool_completed":
             turn["_thinking_active"] = False
-            if payload.get("is_error"):
-                turn["status"] = "error"
-                turn["error"] = str(payload.get("summary") or payload.get("tool_name") or "tool failed")
-            else:
-                turn["status"] = "running"
+            turn["status"] = "running"
             return
 
         if kind == "turn_completed":
@@ -397,7 +499,7 @@ class WebSessionView:
                 self._set_spinner_status(
                     shorten_title(
                         "calling {0}({1})".format(tool_name, call.arguments),
-                        limit=72,
+                        limit=SPINNER_STATUS_PREVIEW_LIMIT,
                     )
                 )
             elif tool_name:
@@ -438,7 +540,7 @@ class WebSessionView:
         submission_id = str(submission_id or "").strip()
         turn_id = str(turn_id or submission_id).strip()
         turn = self._turns_by_submission_id.get(submission_id)
-        if turn is None and turn_id:
+        if turn is None and turn_id and not submission_id:
             turn = self._turns_by_turn_id.get(turn_id)
         if turn is None:
             turn = {
@@ -478,6 +580,9 @@ class WebSessionView:
         turn["prompt"] = ""
         return turn
 
+    def _set_title(self, title: "typing.Union[str, None]") -> None:
+        self._title = str(title or "").strip()
+
     def _last_active_turn(self) -> "typing.Union[typing.Dict[str, object], None]":
         for turn in reversed(self._turns):
             if turn.get("kind") != "control" and turn.get("status") not in {
@@ -489,11 +594,29 @@ class WebSessionView:
         return None
 
     def _publish_nowait(self, event: "typing.Dict[str, object]") -> None:
-        self._events.append(event)
-        if len(self._events) > 500:
-            del self._events[:-500]
-        for subscriber in tuple(self._subscribers):
-            subscriber.put_nowait(event)
+        with self._lock:
+            self._events.append(event)
+            if len(self._events) > 500:
+                del self._events[:-500]
+            subscribers = tuple(self._subscribers)
+        self._publish_to_queues(subscribers, event)
+
+    def _publish_to_queues(
+        self,
+        queues: "typing.Iterable[asyncio.Queue]",
+        event: "typing.Union[typing.Dict[str, object], None]",
+    ) -> None:
+        loop = self._server_loop
+        if loop is None:
+            for queue in queues:
+                queue.put_nowait(event)
+            return
+
+        def publish() -> None:
+            for queue in queues:
+                queue.put_nowait(event)
+
+        loop.call_soon_threadsafe(publish)
 
     async def _publish(self, event: "typing.Dict[str, object]") -> None:
         self._publish_nowait(event)
@@ -513,7 +636,10 @@ class WorkspaceInteractiveSession:
         self._task: "typing.Union[asyncio.Task[int], None]" = None
         self._initial_prompt_submitted = False
 
-    async def start(self) -> "WorkspaceInteractiveSession":
+    async def start(
+        self,
+        submit_initial_prompt: bool = True,
+    ) -> "WorkspaceInteractiveSession":
         if self._task is None:
             self._task = asyncio.create_task(
                 run_interactive_session(
@@ -524,7 +650,11 @@ class WorkspaceInteractiveSession:
                     show_banner=False,
                 )
             )
-        if self.initial_prompt and not self._initial_prompt_submitted:
+        if (
+            submit_initial_prompt
+            and self.initial_prompt
+            and not self._initial_prompt_submitted
+        ):
             self._initial_prompt_submitted = True
             await self.view.submit(self.initial_prompt)
         return self
@@ -566,31 +696,227 @@ class WorkspaceInteractiveSession:
         snapshot["model"] = getattr(getattr(agent, "_model_client", None), "model", "pycodex")
         return snapshot
 
+    def summary(self) -> "typing.Dict[str, object]":
+        summary = self.view.summary()
+        agent = getattr(self.queue, "_agent", None)
+        summary["model"] = getattr(getattr(agent, "_model_client", None), "model", "pycodex")
+        return summary
+
+    def rollout_path(self) -> str:
+        recorder = getattr(getattr(self.queue, "_agent", None), "_rollout_recorder", None)
+        path = getattr(recorder, "rollout_path", None)
+        return "" if path is None else str(path)
+
+    async def restore_from_rollout(self, rollout_path: str, title: str = "") -> None:
+        resumed = load_resumed_session_path(rollout_path, thread_name=title or None)
+        agent = self.queue._agent
+        agent.replace_history(resumed["history"])
+        model_client = getattr(agent, "_model_client", None)
+        if hasattr(model_client, "_session_id"):
+            model_client._session_id = str(resumed["session_id"])
+        agent.set_rollout_recorder(SessionRolloutRecorder.resume(resumed["rollout_path"]))
+        self.view.load_session_history(
+            str(title or resumed["title"]),
+            tuple(resumed["turns"]),
+        )
+
+
+class ThreadedWorkspaceInteractiveSession:
+    def __init__(
+        self,
+        session_factory: "ThreadedSessionFactory",
+        server_loop: "asyncio.AbstractEventLoop",
+    ) -> None:
+        self._session_factory = session_factory
+        self._server_loop = server_loop
+        self._view = WebSessionView()
+        self._view.attach_server_loop(server_loop)
+        self._thread: "typing.Union[threading.Thread, None]" = None
+        self._worker_loop: "typing.Union[asyncio.AbstractEventLoop, None]" = None
+        self._ready = threading.Event()
+        self._closed = threading.Event()
+        self._startup_error: "typing.Union[BaseException, None]" = None
+        self._session: "typing.Union[WorkspaceInteractiveSession, None]" = None
+        self._submit_initial_prompt = True
+
+    async def start(
+        self,
+        submit_initial_prompt: bool = True,
+    ) -> "ThreadedWorkspaceInteractiveSession":
+        if self._thread is not None:
+            return self
+        self._submit_initial_prompt = submit_initial_prompt
+        self._thread = threading.Thread(
+            target=self._thread_main,
+            name="pycodex-workspace-session",
+            daemon=True,
+        )
+        self._thread.start()
+        await asyncio.to_thread(self._ready.wait)
+        if self._startup_error is not None:
+            raise RuntimeError("workspace session thread failed to start") from self._startup_error
+        return self
+
+    def _thread_main(self) -> None:
+        loop = asyncio.new_event_loop()
+        self._worker_loop = loop
+        self._view.attach_worker_loop(loop)
+        asyncio.set_event_loop(loop)
+        try:
+            session = self._session_factory()
+            session.view = self._view
+            self._session = session
+            loop.run_until_complete(
+                session.start(submit_initial_prompt=self._submit_initial_prompt)
+            )
+            self._ready.set()
+            loop.run_forever()
+        except BaseException as exc:
+            self._startup_error = exc
+            self._ready.set()
+        finally:
+            session = self._session
+            if session is not None:
+                try:
+                    loop.run_until_complete(session.close())
+                except BaseException:
+                    pass
+            pending = asyncio.all_tasks(loop)
+            for task in pending:
+                task.cancel()
+            if pending:
+                loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+            asyncio.set_event_loop(None)
+            loop.close()
+            self._closed.set()
+
+    async def close(self) -> None:
+        session = self._session
+        loop = self._worker_loop
+        if session is not None and loop is not None and loop.is_running():
+            future = asyncio.run_coroutine_threadsafe(session.close(), loop)
+            try:
+                await asyncio.wait_for(
+                    asyncio.wrap_future(future),
+                    timeout=SESSION_CLOSE_TIMEOUT_SECONDS + 1.0,
+                )
+            except (asyncio.TimeoutError, RuntimeError):
+                cancel_current = getattr(getattr(session, "queue", None), "cancel_current", None)
+                if callable(cancel_current):
+                    cancel_current()
+        if loop is not None and loop.is_running():
+            loop.call_soon_threadsafe(loop.stop)
+        thread = self._thread
+        if thread is not None:
+            await asyncio.to_thread(thread.join, SESSION_CLOSE_TIMEOUT_SECONDS + 1.0)
+        self._thread = None
+
+    async def submit(self, prompt: str, sender: str = "web") -> "typing.Dict[str, object]":
+        del sender
+        result = await self._view.submit(prompt)
+        result["snapshot"] = self.snapshot()
+        return result
+
+    def subscribe(self) -> "asyncio.Queue":
+        return self._view.subscribe()
+
+    def unsubscribe(self, queue: "asyncio.Queue") -> None:
+        self._view.unsubscribe(queue)
+
+    def snapshot(self) -> "typing.Dict[str, object]":
+        snapshot = self._view.snapshot()
+        session = self._session
+        queue = getattr(session, "queue", None)
+        agent = getattr(queue, "_agent", None)
+        snapshot["model"] = getattr(getattr(agent, "_model_client", None), "model", "pycodex")
+        return snapshot
+
+    def summary(self) -> "typing.Dict[str, object]":
+        summary = self._view.summary()
+        session = self._session
+        queue = getattr(session, "queue", None)
+        agent = getattr(queue, "_agent", None)
+        summary["model"] = getattr(getattr(agent, "_model_client", None), "model", "pycodex")
+        return summary
+
+    def rollout_path(self) -> str:
+        if self._session is None:
+            return ""
+        return self._session.rollout_path()
+
+    async def restore_from_rollout(self, rollout_path: str, title: str = "") -> None:
+        session = self._session
+        loop = self._worker_loop
+        if session is None or loop is None:
+            return
+
+        future = asyncio.run_coroutine_threadsafe(
+            session.restore_from_rollout(rollout_path, title=title),
+            loop,
+        )
+        await asyncio.wrap_future(future)
+
 
 class WorkspaceSessionManager:
-    def __init__(self, session_factory: "SessionFactory") -> None:
+    def __init__(
+        self,
+        session_factory: "SessionFactory",
+        board_path: "typing.Union[Path, None]" = None,
+    ) -> None:
         self._session_factory = session_factory
         self._sessions: "typing.Dict[str, WorkspaceInteractiveSession]" = {}
         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._state_store = WorkspaceStateStore(board_path)
 
     async def start(self) -> None:
-        await self.create_session()
+        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 ""),
+            )
 
     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()
 
-    async def create_session(self) -> str:
+    async def create_session(
+        self,
+        title: str = "",
+        rollout_path: str = "",
+    ) -> str:
         async with self._lock:
             session_id = uuid7_string()
             session = self._session_factory()
-            await session.start()
+            await session.start(submit_initial_prompt=not bool(rollout_path))
+
+            if rollout_path:
+                await session.restore_from_rollout(rollout_path, title=title)
+
             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)
+            )
             return session_id
 
     async def close_session(self, session_id: str) -> None:
@@ -600,10 +926,47 @@ class WorkspaceSessionManager:
             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._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_title(self, session_id: str, session) -> None:
+        subscriber = session.subscribe()
+        try:
+            while True:
+                event = await subscriber.get()
+                if event is None:
+                    return
+                if not isinstance(event, dict) or event.get("type") != "title_changed":
+                    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()
+        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
+            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})
+        self._state_store.save_tabs(tabs)
 
     def get(self, session_id: "typing.Union[str, None]" = None) -> "WorkspaceInteractiveSession":
         resolved_id = self.resolve_session_id(session_id)
@@ -623,14 +986,14 @@ class WorkspaceSessionManager:
         result = []
         for session_id in self._session_order:
             session = self._sessions[session_id]
-            snapshot = _session_snapshot(session)
+            summary = _session_summary(session)
             result.append(
                 {
                     "id": session_id,
-                    "title": snapshot.get("title") or "pycodex",
-                    "running": bool(snapshot.get("running")),
-                    "spinner": snapshot.get("spinner") or "",
-                    "turn_count": len(snapshot.get("turns") or []),
+                    "title": summary.get("title") or "pycodex",
+                    "running": bool(summary.get("running")),
+                    "spinner": summary.get("spinner") or "",
+                    "turn_count": summary.get("turn_count") or 0,
                 }
             )
         return result
@@ -643,7 +1006,7 @@ def create_app(
     manager = (
         session_source
         if isinstance(session_source, WorkspaceSessionManager)
-        else WorkspaceSessionManager(session_source)
+        else WorkspaceSessionManager(session_source, board_path)
     )
 
     if asynccontextmanager is not None:
@@ -742,7 +1105,9 @@ def create_app(
         return JSONResponse({"ok": True, "sessions": manager.list_sessions()})
 
     @app.get("/api/session")
-    async def session(session_id: "typing.Union[str, None]" = None) -> JSONResponse:
+    async def session(
+        session_id: "typing.Union[str, None]" = None,
+    ) -> JSONResponse:
         try:
             resolved_id = manager.resolve_session_id(session_id)
             link = manager.get(resolved_id)
@@ -757,7 +1122,9 @@ def create_app(
         )
 
     @app.post("/api/session/message")
-    async def message(payload: "typing.Dict[str, object]") -> JSONResponse:
+    async def message(
+        payload: "typing.Dict[str, object]",
+    ) -> JSONResponse:
         session_id = str(payload.get("session_id") or "")
         try:
             link = manager.get(session_id or None)
@@ -821,6 +1188,19 @@ 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 []),
+    }
+
+
 def _public_turn(turn: "typing.Dict[str, object]") -> "typing.Dict[str, object]":
     return typing.cast(
         "typing.Dict[str, object]",
@@ -884,7 +1264,7 @@ def run_serve_cli(args: "argparse.Namespace") -> int:
         raise ValueError("board parent directory does not exist: {0}".format(board_path.parent))
 
     configure_loguru()
-    def session_factory() -> "WorkspaceInteractiveSession":
+    def build_session() -> "WorkspaceInteractiveSession":
         model = build_model(
             config_path=args.config,
             profile=args.profile,
@@ -907,7 +1287,10 @@ def run_serve_cli(args: "argparse.Namespace") -> int:
             initial_prompt=initial_prompt,
         )
 
-    app = create_app(WorkspaceSessionManager(session_factory), board_path)
+    def session_factory() -> "ThreadedWorkspaceInteractiveSession":
+        return ThreadedWorkspaceInteractiveSession(build_session, asyncio.get_running_loop())
+
+    app = create_app(WorkspaceSessionManager(session_factory, board_path), board_path)
     print(
         "pycodex workspace listening on http://{0}:{1}".format(host, port),
         flush=True,