From 61f1a07ba7c6b7ce032bf74fcf30e314c1bacdc5 Mon Sep 17 00:00:00 2001 From: Felix Martin Date: Wed, 10 Jun 2026 20:21:37 -0400 Subject: [PATCH] feat: worktree-aware session lifecycle in SessionService MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add worktree support to SessionService via the hqt.git.worktree module: - create_session(worktree_branch=...): creates the worktree before the row (ServiceError → no row, no spawn); spawn/capture use the worktree path; blank nickname defaults to the branch name. - _session_path / _repo_path_for_worktree helpers route every harness-facing path through the worktree when set, with a parent.parent repo fallback when the project row is gone. - attach_session rung 0: recreate a missing worktree from its branch, else raise "delete the session". - delete_session(remove_worktree=, force=): remove worktree before deleting the row; removal failure keeps the row. - worktree_state_for(session_id): None for plain/missing sessions, else the module's WorktreeState. Co-Authored-By: Claude Fable 5 --- src/hqt/sessions/service.py | 126 ++++++++++++-- tests/test_sessions.py | 327 ++++++++++++++++++++++++++++++++++++ 2 files changed, 441 insertions(+), 12 deletions(-) diff --git a/src/hqt/sessions/service.py b/src/hqt/sessions/service.py index 4289905..d696233 100644 --- a/src/hqt/sessions/service.py +++ b/src/hqt/sessions/service.py @@ -11,6 +11,8 @@ from sqlalchemy.orm import selectinload, sessionmaker from hqt.db.models import Harness, Project, Session from hqt.errors import ServiceError +from hqt.git import worktree +from hqt.git.worktree import WorktreeState from hqt.harnesses.base import HarnessConfigurator, SpawnConfig from hqt.status import window_label from hqt.tmux.manager import SpawnRequest, TmuxManager @@ -74,6 +76,30 @@ class SessionService: raise ServiceError(f"Project {project_id} no longer exists") return Path(project.path) + def _session_path(self, db: DBSession, sess: Session) -> Path: + """The harness-facing working directory for a session. + + Worktree sessions run in their own checkout; normal sessions run in + the project path. + """ + if sess.worktree_path: + return Path(sess.worktree_path) + return self._project_path(db, sess.project_id) + + def _repo_path_for_worktree(self, db: DBSession, sess: Session) -> Path: + """Resolve the repo path for a worktree session. + + Prefer the project row; if it is gone, derive the repo from the + worktree layout (/.worktrees/ → repo = parent.parent). + + Only meaningful for worktree sessions; ``sess.worktree_path`` must be set. + """ + project = db.get(Project, sess.project_id) + if project is not None: + return Path(project.path) + assert sess.worktree_path is not None + return Path(sess.worktree_path).parent.parent + async def _capture_session_id_with_retry( self, configurator: HarnessConfigurator, @@ -118,12 +144,18 @@ class SessionService: if sess.harness_session_id and sess.harness_session_id != placeholder_id: return - # Polling path: a deleted project just means we skip capture, never crash. - project = db.get(Project, sess.project_id) - if project is None: - return + if sess.worktree_path: + # Worktree sessions run in their own checkout; the project row is not + # needed to locate the working directory. + capture_path = Path(sess.worktree_path) + else: + # Polling path: a deleted project just means we skip capture, never crash. + project = db.get(Project, sess.project_id) + if project is None: + return + capture_path = Path(project.path) captured_id = configurator.capture_session_id( - Path(project.path), self._session_created_epoch(sess) + capture_path, self._session_created_epoch(sess) ) if captured_id and captured_id != sess.harness_session_id: sess.harness_session_id = captured_id @@ -135,11 +167,17 @@ class SessionService: harness_name: str, nickname: str | None = None, model: str | None = None, + worktree_branch: str | None = None, ) -> "CreateSessionResult": """Create a new session row, spawn the harness window, and return a CreateSessionResult with the Session row and spawn outcome. The capture-retry loop only runs when spawn succeeded. + + When ``worktree_branch`` is set, a dedicated git worktree is created + BEFORE the Session row is added: if worktree creation fails the + ServiceError propagates with no row and no spawn. The harness then runs + in the worktree path rather than the project path. """ log.info("Creating session: project=%s harness=%s", project_id, harness_name) configurator = self._configurator(harness_name) @@ -149,12 +187,26 @@ class SessionService: log.error("Harness not in DB: %s", harness_name) raise ServiceError(f"Unknown harness: {harness_name}") project_path = self._project_path(db, project_id) + + # Worktree sessions: create the worktree first; on failure no row is + # added and no spawn happens (ServiceError propagates). + worktree_path: str | None = None + run_path = project_path + if worktree_branch: + wt_path = await worktree.create_worktree(project_path, worktree_branch) + worktree_path = str(wt_path) + run_path = wt_path + if not (nickname or "").strip(): + nickname = worktree_branch + sess = Session( project_id=project_id, harness_id=harness_row.id, nickname=nickname, model=model, tmux_session_name="placeholder", + worktree_path=worktree_path, + worktree_branch=worktree_branch or None, archived=False, ) db.add(sess) @@ -163,7 +215,7 @@ class SessionService: harness_session_id = configurator.generate_session_id(sess.id) sess.harness_session_id = harness_session_id spawn_cfg = configurator.build_spawn_config( - project_path, harness_session_id, model + run_path, harness_session_id, model ) log.info( "Spawning window %s: cmd=%s cwd=%s", @@ -191,7 +243,7 @@ class SessionService: # the same cwd. if spawn_result.ok and configurator.captures_session_id: captured_id = await self._capture_session_id_with_retry( - configurator, project_path, since, sess.tmux_session_name + configurator, run_path, since, sess.tmux_session_name ) if captured_id: sess.harness_session_id = captured_id @@ -212,6 +264,21 @@ class SessionService: if sess is None: raise ServiceError("Session not found") window_name = sess.tmux_session_name + + # Rung 0: a worktree session whose directory was removed from disk + # but whose branch still exists can be recreated before anything + # else runs. If recreation fails, the branch is gone too. + if sess.worktree_path and not Path(sess.worktree_path).exists(): + project_path = self._project_path(db, sess.project_id) + try: + await worktree.add_worktree_for_branch( + project_path, sess.worktree_branch + ) + except ServiceError as exc: + raise ServiceError( + "Worktree and branch are gone; delete the session" + ) from exc + self._maybe_capture_missing_session_id(db, sess) if await self.tmux.is_alive(window_name): @@ -254,12 +321,12 @@ class SessionService: # Rung 2: fresh spawn (codex ignores the old session id; start fresh) configurator = self._configurator(sess.harness.name) - project_path = self._project_path(db, sess.project_id) + run_path = self._session_path(db, sess) harness_session_id = ( sess.harness_session_id or configurator.generate_session_id(sess.id) ) spawn_cfg = configurator.build_spawn_config( - project_path, harness_session_id, sess.model + run_path, harness_session_id, sess.model ) since = time.time() result = await self.tmux.respawn_verified( @@ -277,7 +344,7 @@ class SessionService: # it so that future resumes target this fresh conversation. if configurator.captures_session_id: new_id = await self._capture_session_id_with_retry( - configurator, project_path, since, window_name + configurator, run_path, since, window_name ) if new_id: sess.harness_session_id = new_id @@ -318,16 +385,51 @@ class SessionService: sess.nickname = (nickname or "").strip() or None db.commit() - async def delete_session(self, session_id: int) -> None: + async def delete_session( + self, session_id: int, remove_worktree: bool = False, force: bool = False + ) -> None: + """Delete a session: kill its tmux window, optionally remove its worktree, + then delete the DB row. + + When ``remove_worktree`` is True and the session has a worktree, the + worktree is removed first; if removal raises ServiceError, the row is + NOT deleted and the error propagates (the already-killed window is + acceptable). + """ with self.factory() as db: sess = db.get(Session, session_id) if sess is None: return if await self.tmux.window_exists(sess.tmux_session_name): await self.tmux.kill(sess.tmux_session_name) + + if remove_worktree and sess.worktree_path: + repo_path = self._repo_path_for_worktree(db, sess) + # Let ServiceError propagate: do NOT delete the row on failure. + await worktree.remove_worktree( + repo_path, + Path(sess.worktree_path), + sess.worktree_branch, + force, + ) + db.delete(sess) db.commit() + async def worktree_state_for(self, session_id: int) -> WorktreeState | None: + """Return the WorktreeState for a session's worktree, or None. + + None when the session does not exist or has no worktree. + """ + with self.factory() as db: + sess = db.get(Session, session_id) + if sess is None or not sess.worktree_path: + return None + repo_path = self._repo_path_for_worktree(db, sess) + return await worktree.worktree_state( + repo_path, Path(sess.worktree_path), sess.worktree_branch + ) + async def _status_for( self, sess: Session, wi: WindowInfo | None, now: float ) -> str: @@ -412,5 +514,5 @@ class SessionService: sess.harness_session_id or configurator.generate_session_id(sess.id) ) return configurator.build_resume_config( - self._project_path(db, sess.project_id), harness_session_id, sess.model + self._session_path(db, sess), harness_session_id, sess.model ) diff --git a/tests/test_sessions.py b/tests/test_sessions.py index 60eabc2..c0c847c 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -1106,3 +1106,330 @@ async def test_list_sessions_skips_capture_for_deleted_project( harnesses_capture["claude-code"].capture_session_id.return_value = "would-capture" infos = await service.list_sessions(1) assert len(infos) == 1 # no ServiceError escaped + + +# --------------------------------------------------------------------------- +# Task 3: worktree-aware session lifecycle +# --------------------------------------------------------------------------- + + +WT_PATH = Path("/tmp/myproj/.worktrees/feature-x") + + +@pytest.fixture +def fake_worktree(monkeypatch): + """Monkeypatch hqt.git.worktree functions with recording async fakes. + + Returns a namespace exposing the recorded call args and the WorktreeState + class so tests can assert ordering and arguments. + """ + import hqt.git.worktree as wt_mod + + class Recorder: + def __init__(self): + self.create_calls = [] + self.add_calls = [] + self.remove_calls = [] + self.state_calls = [] + self.create_path = WT_PATH + self.create_error: Exception | None = None + self.add_error: Exception | None = None + self.remove_error: Exception | None = None + self.state_result = wt_mod.WorktreeState( + exists=True, dirty=False, unique_commits=2 + ) + + async def create_worktree(self, repo, branch): + self.create_calls.append((repo, branch)) + if self.create_error is not None: + raise self.create_error + return self.create_path + + async def add_worktree_for_branch(self, repo, branch): + self.add_calls.append((repo, branch)) + if self.add_error is not None: + raise self.add_error + return self.create_path + + async def remove_worktree(self, repo, path, branch, force): + self.remove_calls.append((repo, path, branch, force)) + if self.remove_error is not None: + raise self.remove_error + + async def worktree_state(self, repo, path, branch): + self.state_calls.append((repo, path, branch)) + return self.state_result + + rec = Recorder() + monkeypatch.setattr(wt_mod, "create_worktree", rec.create_worktree) + monkeypatch.setattr(wt_mod, "add_worktree_for_branch", rec.add_worktree_for_branch) + monkeypatch.setattr(wt_mod, "remove_worktree", rec.remove_worktree) + monkeypatch.setattr(wt_mod, "worktree_state", rec.worktree_state) + return rec + + +@pytest.mark.asyncio +async def test_create_session_worktree_branch( + service, db, tmux, harnesses, fake_worktree +): + """worktree_branch set → create_worktree called with (project_path, branch); + spawn cwd and capture use the worktree path; row stores both fields.""" + harnesses["claude-code"].captures_session_id = True + harnesses["claude-code"].capture_session_id.return_value = "cap-id" + harnesses["claude-code"].build_spawn_config.return_value = MagicMock( + command=["claude"], env={}, cwd=WT_PATH + ) + + result = await service.create_session( + project_id=1, harness_name="claude-code", worktree_branch="feature-x" + ) + + # create_worktree called with project path + branch + assert fake_worktree.create_calls == [(Path("/tmp/myproj"), "feature-x")] + # spawn request cwd is the worktree path + spawn_req = tmux.spawn.call_args.args[0] + assert spawn_req.cwd == str(WT_PATH) + # build_spawn_config used the worktree path, not the project path + assert harnesses["claude-code"].build_spawn_config.call_args.args[0] == WT_PATH + # capture used the worktree path + assert harnesses["claude-code"].capture_session_id.call_args.args[0] == WT_PATH + # row carries both worktree fields + sess = result.session + assert sess.worktree_path == str(WT_PATH) + assert sess.worktree_branch == "feature-x" + + +@pytest.mark.asyncio +async def test_create_session_worktree_blank_nickname_defaults_to_branch( + service, db, tmux, harnesses, fake_worktree +): + """Blank nickname + worktree_branch → nickname defaults to branch name.""" + result = await service.create_session( + project_id=1, + harness_name="claude-code", + nickname="", + worktree_branch="feature-x", + ) + assert result.session.nickname == "feature-x" + + +@pytest.mark.asyncio +async def test_create_session_worktree_explicit_nickname_preserved( + service, db, tmux, harnesses, fake_worktree +): + """Explicit nickname + worktree_branch → nickname preserved.""" + result = await service.create_session( + project_id=1, + harness_name="claude-code", + nickname="my-work", + worktree_branch="feature-x", + ) + assert result.session.nickname == "my-work" + + +@pytest.mark.asyncio +async def test_create_session_worktree_create_fails_no_row_no_spawn( + factory, db, tmux, harnesses, fake_worktree +): + """create_worktree raising ServiceError → propagates, no Session row, no spawn.""" + from hqt.db.models import Session + + fake_worktree.create_error = ServiceError("git worktree add failed") + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + + with pytest.raises(ServiceError, match="worktree add failed"): + await service.create_session( + project_id=1, harness_name="claude-code", worktree_branch="feature-x" + ) + + # No Session row committed + db.expire_all() + assert db.query(Session).count() == 0 + # tmux spawn never called + tmux.spawn.assert_not_called() + + +@pytest.mark.asyncio +async def test_create_session_without_worktree_branch_regression( + service, db, tmux, harnesses, fake_worktree +): + """No worktree_branch → create_worktree not called, row worktree fields None.""" + result = await service.create_session(project_id=1, harness_name="claude-code") + assert fake_worktree.create_calls == [] + assert result.session.worktree_path is None + assert result.session.worktree_branch is None + + +def _seed_worktree_session(db, branch="feature-x", path=str(WT_PATH)): + """Seed a worktree session for project 1 / harness claude-code.""" + from hqt.db.models import Harness, Session + + harness = db.query(Harness).filter_by(name="claude-code").first() + sess = Session( + project_id=1, + harness_id=harness.id, + tmux_session_name="hqt-wt", + harness_session_id="sess-1", + worktree_path=path, + worktree_branch=branch, + archived=False, + ) + db.add(sess) + db.commit() + return sess + + +@pytest.mark.asyncio +async def test_attach_session_recreates_missing_worktree( + factory, db, tmux, harnesses, fake_worktree +): + """attach_session: worktree dir missing on disk → add_worktree_for_branch called.""" + # Use a path that does not exist on disk. + missing = "/tmp/does-not-exist/.worktrees/feature-x" + sess = _seed_worktree_session(db, path=missing) + + harnesses["claude-code"].build_resume_config = MagicMock( + return_value=MagicMock(command=["claude"], env={}, cwd=Path(missing)) + ) + tmux.is_alive = AsyncMock(return_value=False) + tmux.respawn_verified = AsyncMock( + return_value=SpawnResult(ok=True, window_id=None, error="") + ) + tmux.attach = AsyncMock(return_value=True) + + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + await service.attach_session(session_id=sess.id) + + assert fake_worktree.add_calls == [(Path("/tmp/myproj"), "feature-x")] + + +@pytest.mark.asyncio +async def test_attach_session_missing_worktree_recreate_fails( + factory, db, tmux, harnesses, fake_worktree +): + """attach_session: recreate raises → ServiceError mentioning deleting the session.""" + missing = "/tmp/does-not-exist/.worktrees/feature-x" + sess = _seed_worktree_session(db, path=missing) + fake_worktree.add_error = ServiceError("branch gone") + + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + with pytest.raises(ServiceError, match="delete the session"): + await service.attach_session(session_id=sess.id) + + +@pytest.mark.asyncio +async def test_attach_session_worktree_resume_uses_worktree_path( + factory, db, tmux, harnesses, fake_worktree +): + """Resume config for a worktree session uses the worktree path as cwd.""" + # WT_PATH parent exists (/tmp/myproj/.worktrees won't, but use a real existing dir) + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + sess = _seed_worktree_session(db, path=tmpdir) + + resume_cfg = MagicMock(command=["claude", "--resume"], env={}, cwd=Path(tmpdir)) + harnesses["claude-code"].build_resume_config = MagicMock( + return_value=resume_cfg + ) + tmux.is_alive = AsyncMock(return_value=False) + tmux.respawn_verified = AsyncMock( + return_value=SpawnResult(ok=True, window_id=None, error="") + ) + tmux.attach = AsyncMock(return_value=True) + + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + await service.attach_session(session_id=sess.id) + + # build_resume_config received the worktree path as the project path arg + assert harnesses["claude-code"].build_resume_config.call_args.args[0] == Path( + tmpdir + ) + # since the worktree exists on disk, recreate must not have been attempted + assert fake_worktree.add_calls == [] + + +@pytest.mark.asyncio +async def test_delete_session_remove_worktree_true( + factory, db, tmux, harnesses, fake_worktree +): + """delete_session(remove_worktree=True) calls remove_worktree and deletes the row.""" + from hqt.db.models import Session + + sess = _seed_worktree_session(db) + sess_id = sess.id + + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + await service.delete_session(sess_id, remove_worktree=True, force=True) + + assert fake_worktree.remove_calls == [ + (Path("/tmp/myproj"), WT_PATH, "feature-x", True) + ] + db.expire_all() + assert db.get(Session, sess_id) is None + + +@pytest.mark.asyncio +async def test_delete_session_remove_worktree_raises_keeps_row( + factory, db, tmux, harnesses, fake_worktree +): + """remove_worktree raising → ServiceError propagates, row still present.""" + from hqt.db.models import Session + + sess = _seed_worktree_session(db) + sess_id = sess.id + fake_worktree.remove_error = ServiceError("worktree dirty") + + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + with pytest.raises(ServiceError, match="worktree dirty"): + await service.delete_session(sess_id, remove_worktree=True) + + db.expire_all() + assert db.get(Session, sess_id) is not None + + +@pytest.mark.asyncio +async def test_delete_session_remove_worktree_false_skips_removal( + factory, db, tmux, harnesses, fake_worktree +): + """remove_worktree=False → remove_worktree not called, row deleted.""" + from hqt.db.models import Session + + sess = _seed_worktree_session(db) + sess_id = sess.id + + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + await service.delete_session(sess_id, remove_worktree=False) + + assert fake_worktree.remove_calls == [] + db.expire_all() + assert db.get(Session, sess_id) is None + + +@pytest.mark.asyncio +async def test_worktree_state_for_plain_session_returns_none( + service, db, tmux, harnesses, fake_worktree +): + """worktree_state_for: plain session → None.""" + sess = _seed_session(db) + assert await service.worktree_state_for(sess.id) is None + + +@pytest.mark.asyncio +async def test_worktree_state_for_missing_id_returns_none(service, db, fake_worktree): + """worktree_state_for: unknown session id → None.""" + assert await service.worktree_state_for(99999) is None + + +@pytest.mark.asyncio +async def test_worktree_state_for_worktree_session_returns_state( + factory, db, tmux, harnesses, fake_worktree +): + """worktree_state_for: worktree session → returns the module's WorktreeState.""" + sess = _seed_worktree_session(db) + service = SessionService(factory=factory, tmux=tmux, harnesses=harnesses) + + state = await service.worktree_state_for(sess.id) + + assert state is fake_worktree.state_result + assert fake_worktree.state_calls == [(Path("/tmp/myproj"), WT_PATH, "feature-x")]