""" Integration tests for git worktree creation on ``POST /v1/sessions``. Drives the JSON create path with a `git` block through the full app and a fake host that auto-replies to the worktree control frames. Verifies that the request's branch_name + base_branch reach the host's ``host.create_worktree`` frame, and that the created worktree path and branch are persisted on the session. See designs/SESSION_GIT_WORKTREE.md. """ from __future__ import annotations import asyncio import contextlib from collections.abc import AsyncIterator, Callable from dataclasses import dataclass, field from typing import Any import httpx import pytest import pytest_asyncio from fastapi import FastAPI from omnigent.host.frames import ( HostCreateWorktreeFrame, HostHelloFrame, HostRemoveWorktreeFrame, HostStatFrame, decode_host_frame, ) from omnigent.server.auth import RESERVED_USER_LOCAL from omnigent.server.host_registry import HostConnection from omnigent.stores.host_store import HostStore from tests.server.helpers import create_test_agent pytestmark = pytest.mark.asyncio _HOST_ID = "host_wt_create" _SOURCE_REPO = "/Users/alice/myrepo" class _FakeWebSocket: """Minimal WebSocket stand-in (the registry only enqueues).""" async def send_text(self, data: str) -> None: """No-op send — frames flow through the outbound queue. :param data: JSON-encoded frame text (ignored). """ @dataclass class _HostCapture: """ Frames a fake host received during one ``POST /v1/sessions`` create. :param create: ``host.create_worktree`` frames received. :param remove: ``host.remove_worktree`` frames received (a non-empty list proves the create-rollback path fired). """ create: list[HostCreateWorktreeFrame] = field(default_factory=list) remove: list[HostRemoveWorktreeFrame] = field(default_factory=list) # Factory yielded by the ``register_worktree_host`` fixture: # register(*, create_status=, create_error=) -> _HostCapture. RegisterHost = Callable[..., _HostCapture] @pytest_asyncio.fixture() async def register_worktree_host( app: FastAPI, db_uri: str, ) -> AsyncIterator[RegisterHost]: """Yield a factory that registers a fake host with a replying drain. The drain answers ``host.stat`` (so workspace validation passes) and ``host.create_worktree`` (capturing each frame). Every drain started during the test is poisoned and awaited at teardown, so no background task leaks into the next test's event loop (mirrors the cleanup in ``test_host_worktree.py``). :param app: App whose ``host_registry`` to register into. :param db_uri: DB URI so the ``host_id`` FK target row exists. :returns: Async iterator yielding a ``register`` factory. Its kwargs: ``create_status`` (``"ok"`` returns a worktree path, ``"failed"`` simulates a host git failure such as a bad base ref) and ``create_error`` (the failure message). Returns a ``_HostCapture`` whose ``.create`` / ``.remove`` lists accumulate the create- and remove-worktree frames the host received. """ conns: list[HostConnection] = [] def _register(*, create_status: str = "ok", create_error: str | None = None) -> _HostCapture: HostStore(db_uri).upsert_on_connect(_HOST_ID, "wt-host", RESERVED_USER_LOCAL) conn = app.state.host_registry.register( host_id=_HOST_ID, ws=_FakeWebSocket(), # type: ignore[arg-type] — duck-typed hello=HostHelloFrame(version="0.1.0-test", frame_protocol_version=1, name="wt-host"), owner=RESERVED_USER_LOCAL, ) cap = _HostCapture() async def _drain() -> None: """Answer stat + create/remove-worktree frames; capture them.""" while True: frame_text = await conn.outbound_queue.get() if frame_text is None: return frame = decode_host_frame(frame_text) if isinstance(frame, HostStatFrame): fut = conn.pending_stats.pop(frame.request_id, None) if fut is not None and not fut.done(): fut.set_result( { "status": "ok", "exists": True, "type": "directory", "canonical_path": frame.path, "error": None, } ) elif isinstance(frame, HostCreateWorktreeFrame): cap.create.append(frame) fut = conn.pending_create_worktrees.pop(frame.request_id, None) if fut is not None and not fut.done(): if create_status == "ok": dirname = frame.branch_name.replace("/", "-") fut.set_result( { "status": "ok", "worktree_path": f"{frame.repo_path}-worktrees/{dirname}", "branch": frame.branch_name, "error": None, } ) else: fut.set_result( { "status": "failed", "worktree_path": None, "branch": None, "error": create_error, } ) elif isinstance(frame, HostRemoveWorktreeFrame): cap.remove.append(frame) fut = conn.pending_remove_worktrees.pop(frame.request_id, None) if fut is not None and not fut.done(): fut.set_result({"status": "ok", "error": None}) conn._drain_task_for_test = asyncio.create_task(_drain()) # type: ignore[attr-defined] conns.append(conn) return cap yield _register # Poison each queue so the drain returns, then await/cancel it. for conn in conns: conn.outbound_queue.put_nowait(None) task = conn._drain_task_for_test # type: ignore[attr-defined] with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError, Exception): await asyncio.wait_for(asyncio.shield(task), timeout=1.0) if not task.done(): task.cancel() async def _create_git_session( client: httpx.AsyncClient, agent_id: str, git: dict[str, Any], ) -> httpx.Response: """POST a JSON session-create with a ``git`` block. :param client: The test HTTP client. :param agent_id: Agent to bind. :param git: The ``git`` block, e.g. ``{"branch_name": "feature/x", "base_branch": "main"}``. :returns: The raw create response. """ return await client.post( "/v1/sessions", json={ "agent_id": agent_id, "host_id": _HOST_ID, "workspace": _SOURCE_REPO, "git": git, }, ) async def test_create_passes_branch_and_base_branch_to_host( register_worktree_host: RegisterHost, client: httpx.AsyncClient, ) -> None: """The request's branch_name + base_branch reach host.create_worktree, and the resulting worktree path + branch are persisted on the session. Proves the server route threads ``git.base_branch`` through ``_create_session_worktree`` → ``create_worktree_on_host`` → the frame. If base_branch were dropped on the route, the captured frame's base_branch would be ``None`` and this fails. """ cap = register_worktree_host() agent = await create_test_agent(client, name="wt-create-agent") resp = await _create_git_session( client, agent["id"], {"branch_name": "feature/login", "base_branch": "main"} ) assert resp.status_code == 201, resp.text # The host received exactly one create-worktree frame carrying both # the new branch and the requested base ref. assert len(cap.create) == 1, f"expected one create_worktree frame, got {len(cap.create)}" frame = cap.create[0] assert frame.repo_path == _SOURCE_REPO assert frame.branch_name == "feature/login" assert frame.base_branch == "main" # The returned worktree path becomes the session workspace, and the # branch is persisted (drives sidebar display + delete cleanup). body = resp.json() assert body["git_branch"] == "feature/login" assert body["workspace"] == f"{_SOURCE_REPO}-worktrees/feature-login" async def test_create_without_base_branch_sends_none( register_worktree_host: RegisterHost, client: httpx.AsyncClient, ) -> None: """Omitting base_branch sends ``None`` to the host (branch from HEAD). Pairs with the test above to pin both directions: a provided base threads through, an omitted one stays ``None`` so the host branches from the source repo's current HEAD. """ cap = register_worktree_host() agent = await create_test_agent(client, name="wt-create-agent-2") resp = await _create_git_session(client, agent["id"], {"branch_name": "wip"}) assert resp.status_code == 201, resp.text assert len(cap.create) == 1 assert cap.create[0].branch_name == "wip" assert cap.create[0].base_branch is None async def test_create_with_invalid_base_branch_fails_400( register_worktree_host: RegisterHost, client: httpx.AsyncClient, ) -> None: """An invalid base branch fails the create with 400 INVALID_INPUT. The host rejects the bad base ref (``host.create_worktree`` → ``status: failed``); the server maps that to INVALID_INPUT (400), NOT 500 — it's user-correctable input — and surfaces the host's reason. Worktree creation fails before ``create_conversation``, so no session row is created (the response carries no session id). """ register_worktree_host( create_status="failed", create_error="base branch does not exist: nope-not-a-branch", ) agent = await create_test_agent(client, name="wt-bad-base-agent") resp = await _create_git_session( client, agent["id"], {"branch_name": "feature/x", "base_branch": "nope-not-a-branch"}, ) # 400 (not 500): a bad base ref is user input, not a server fault. assert resp.status_code == 400, resp.text body = resp.json() assert body["error"]["code"] == "invalid_input" # The host's reason is surfaced verbatim so the UI can show it. assert "base branch does not exist" in body["error"]["message"] async def test_create_with_existing_worktree_persists_without_creating( register_worktree_host: RegisterHost, client: httpx.AsyncClient, ) -> None: """Starting in an existing worktree persists its branch, creates nothing. ``git.existing_worktree`` binds the session straight to a pre-existing worktree directory: no create-worktree frame is sent to the host, and ``branch_name`` is persisted as ``git_branch`` so the sidebar shows it and the opt-in delete flow can offer to remove it. """ cap = register_worktree_host() agent = await create_test_agent(client, name="wt-existing-agent") resp = await client.post( "/v1/sessions", json={ "agent_id": agent["id"], "host_id": _HOST_ID, "workspace": _SOURCE_REPO, "git": {"branch_name": "feature/existing", "existing_worktree": True}, }, ) assert resp.status_code == 201, resp.text # No worktree was created — the host received no create frame. assert len(cap.create) == 0, f"expected no create_worktree frame, got {len(cap.create)}" # The existing worktree's branch is persisted; the workspace is the # supplied directory verbatim (no worktree-path rewrite). body = resp.json() assert body["git_branch"] == "feature/existing" assert body["workspace"] == _SOURCE_REPO async def test_create_with_invalid_existing_worktree_branch_fails_400( register_worktree_host: RegisterHost, client: httpx.AsyncClient, ) -> None: """An invalid bind-mode ``branch_name`` fails the create with 400. The host never runs git for this path, so the server is the only gate on the name; a malformed branch is user-correctable input and maps to INVALID_INPUT (400), not 500. """ register_worktree_host() agent = await create_test_agent(client, name="wt-existing-bad-agent") resp = await client.post( "/v1/sessions", json={ "agent_id": agent["id"], "host_id": _HOST_ID, "workspace": _SOURCE_REPO, "git": {"branch_name": "bad..branch", "existing_worktree": True}, }, ) assert resp.status_code == 400, resp.text body = resp.json() assert body["error"]["code"] == "invalid_input" # The failed create returned an error, not a session. assert "id" not in body async def test_create_failure_never_removes_existing_worktree( register_worktree_host: RegisterHost, client: httpx.AsyncClient, monkeypatch: pytest.MonkeyPatch, ) -> None: """A create_conversation failure must NOT destroy the user's worktree. Regression: the ``existing_worktree`` bind path sets ``git_branch`` for a *pre-existing* worktree without Omnigent creating one. The create-rollback (``git worktree remove --force`` + ``git branch -D``) is gated on Omnigent having created a worktree, NOT on ``git_branch`` being set — otherwise a persistence failure would force-remove the user's own worktree and delete their branch. Assert no remove frame is sent when ``create_conversation`` raises on this path. """ from omnigent.stores.conversation_store.sqlalchemy_store import ( SqlAlchemyConversationStore, ) cap = register_worktree_host() agent = await create_test_agent(client, name="wt-no-destroy-agent") # Force the persistence step to fail after the bind path has already # set git_branch — the exact window the rollback guards. Patch the class # method (the store is a thin, stateless db_uri wrapper, and the route # uses its own instance) so the failure hits regardless of which # instance the router closed over. def _boom(*args: Any, **kwargs: Any) -> Any: raise RuntimeError("simulated create_conversation failure") monkeypatch.setattr(SqlAlchemyConversationStore, "create_conversation", _boom) # The in-process ASGI transport re-raises unhandled server errors, so # the simulated failure surfaces here rather than as a 500 response. # Either way the create failed; what matters is the side effect below. with pytest.raises(RuntimeError, match="simulated create_conversation failure"): await client.post( "/v1/sessions", json={ "agent_id": agent["id"], "host_id": _HOST_ID, "workspace": _SOURCE_REPO, "git": {"branch_name": "feature/existing", "existing_worktree": True}, }, ) # Critically, the user's worktree is left untouched: the create-rollback # did NOT fire, so no remove_worktree frame reached the host. assert cap.remove == [], ( f"create-rollback force-removed the user's existing worktree: {cap.remove}" ) async def test_create_failure_rolls_back_omnigent_created_worktree( register_worktree_host: RegisterHost, client: httpx.AsyncClient, monkeypatch: pytest.MonkeyPatch, ) -> None: """A create_conversation failure DOES clean up an Omnigent-made worktree. The counterpart to the data-loss guard: when Omnigent creates the worktree (the ``git`` path) and persistence then fails, the orphan worktree it just made must be force-removed. Proves the narrowed rollback guard (gated on Omnigent having created a worktree) still fires for the case it is meant to clean up. """ from omnigent.stores.conversation_store.sqlalchemy_store import ( SqlAlchemyConversationStore, ) cap = register_worktree_host() agent = await create_test_agent(client, name="wt-rollback-agent") def _boom(*args: Any, **kwargs: Any) -> Any: raise RuntimeError("simulated create_conversation failure") monkeypatch.setattr(SqlAlchemyConversationStore, "create_conversation", _boom) with pytest.raises(RuntimeError, match="simulated create_conversation failure"): await _create_git_session(client, agent["id"], {"branch_name": "feature/orphan"}) # Omnigent created the worktree, so the rollback force-removed it: one # remove frame for the worktree it just made, deleting the branch too. assert len(cap.create) == 1, cap.create assert len(cap.remove) == 1, f"expected a create-rollback remove frame, got {cap.remove}" assert cap.remove[0].branch == "feature/orphan" assert cap.remove[0].delete_branch is True