chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,445 @@
|
||||
"""
|
||||
Integration tests for git worktree creation on the dedicated per-session
|
||||
bind endpoint ``POST /v1/hosts/{host_id}/runners`` (``launch_runner``).
|
||||
|
||||
This is the endpoint the fork-resume flow uses to bind an already-existing
|
||||
(unbound) session to a host + directory. Unlike ``POST /v1/sessions`` —
|
||||
which creates the worktree before the conversation row exists —
|
||||
``launch_runner`` operates on a row that already exists, so it must create
|
||||
the worktree at bind time and roll it back if the bind/launch fails.
|
||||
|
||||
Drives the endpoint through the full app and a fake host that auto-replies
|
||||
to the host control frames (``host.stat`` for workspace validation,
|
||||
``host.create_worktree``, ``host.launch_runner``, and
|
||||
``host.remove_worktree`` for rollback). 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 pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from fastapi import FastAPI
|
||||
|
||||
from omnigent.host.frames import (
|
||||
HostCreateWorktreeFrame,
|
||||
HostHelloFrame,
|
||||
HostLaunchRunnerFrame,
|
||||
HostRemoveWorktreeFrame,
|
||||
HostStatFrame,
|
||||
decode_host_frame,
|
||||
)
|
||||
from omnigent.runtime.agent_cache import AgentCache
|
||||
from omnigent.server.app import create_app
|
||||
from omnigent.server.auth import RESERVED_USER_LOCAL
|
||||
from omnigent.server.host_registry import HostConnection
|
||||
from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore
|
||||
from omnigent.stores.artifact_store.local import LocalArtifactStore
|
||||
from omnigent.stores.comment_store.sqlalchemy_store import SqlAlchemyCommentStore
|
||||
from omnigent.stores.conversation_store.sqlalchemy_store import (
|
||||
SqlAlchemyConversationStore,
|
||||
)
|
||||
from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore
|
||||
from omnigent.stores.host_store import HostStore
|
||||
from tests.server.helpers import create_test_agent
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
_HOST_ID = "host_runner_wt"
|
||||
_SOURCE_REPO = "/Users/alice/myrepo"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def app(runtime_init: None, db_uri: str, tmp_path: Path) -> FastAPI:
|
||||
"""FastAPI app wired WITH ``host_store`` so ``launch_runner`` can
|
||||
resolve host ownership and launch a runner.
|
||||
|
||||
Overrides the shared ``app`` fixture (which passes
|
||||
``host_store=None`` and so can't run the dedicated launch endpoint).
|
||||
The shared ``client`` fixture depends on this ``app``.
|
||||
|
||||
:param runtime_init: Initializes the runtime + mock LLM.
|
||||
:param db_uri: SQLite database URI.
|
||||
:param tmp_path: Pytest temp dir for artifacts and cache.
|
||||
:returns: A configured FastAPI app with host routes mounted.
|
||||
"""
|
||||
artifact_store = LocalArtifactStore(str(tmp_path / "artifacts"))
|
||||
return create_app(
|
||||
agent_store=SqlAlchemyAgentStore(db_uri),
|
||||
file_store=SqlAlchemyFileStore(db_uri),
|
||||
conversation_store=SqlAlchemyConversationStore(db_uri),
|
||||
artifact_store=artifact_store,
|
||||
agent_cache=AgentCache(
|
||||
artifact_store=artifact_store,
|
||||
cache_dir=tmp_path / "cache",
|
||||
),
|
||||
comment_store=SqlAlchemyCommentStore(db_uri),
|
||||
host_store=HostStore(db_uri),
|
||||
)
|
||||
|
||||
|
||||
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 ``launch_runner`` call.
|
||||
|
||||
:param create: ``host.create_worktree`` frames received.
|
||||
:param launch: ``host.launch_runner`` frames received.
|
||||
:param remove: ``host.remove_worktree`` frames received (a non-empty
|
||||
list proves the rollback path fired).
|
||||
"""
|
||||
|
||||
create: list[HostCreateWorktreeFrame] = field(default_factory=list)
|
||||
launch: list[HostLaunchRunnerFrame] = field(default_factory=list)
|
||||
remove: list[HostRemoveWorktreeFrame] = field(default_factory=list)
|
||||
|
||||
|
||||
# register(*, create_status=, create_error=, launch_status=) -> _HostCapture
|
||||
RegisterHost = Callable[..., _HostCapture]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture()
|
||||
async def register_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`` (workspace validation passes),
|
||||
``host.create_worktree``, ``host.launch_runner``, and
|
||||
``host.remove_worktree`` — capturing each into a :class:`_HostCapture`.
|
||||
Every drain is poisoned and awaited at teardown so no background task
|
||||
leaks into the next test's event loop.
|
||||
|
||||
: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. Kwargs:
|
||||
``create_status`` (``"ok"``/``"failed"``), ``create_error``
|
||||
(host failure detail), ``launch_status``
|
||||
(``"launched"``/``"failed"``). Returns the :class:`_HostCapture`
|
||||
accumulating frames the host received.
|
||||
"""
|
||||
conns: list[HostConnection] = []
|
||||
|
||||
def _register(
|
||||
*,
|
||||
create_status: str = "ok",
|
||||
create_error: str | None = None,
|
||||
launch_status: str = "launched",
|
||||
) -> _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/launch/remove 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, HostLaunchRunnerFrame):
|
||||
cap.launch.append(frame)
|
||||
fut = conn.pending_launches.pop(frame.request_id, None)
|
||||
if fut is not None and not fut.done():
|
||||
fut.set_result(
|
||||
{
|
||||
"status": launch_status,
|
||||
"runner_id": (
|
||||
"runner_from_host" if launch_status == "launched" else None
|
||||
),
|
||||
"error": None if launch_status == "launched" else "boom",
|
||||
}
|
||||
)
|
||||
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
|
||||
|
||||
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 _bare_session(client: httpx.AsyncClient, name: str) -> str:
|
||||
"""Create an unbound session (agent only, no host/workspace).
|
||||
|
||||
:param client: The test HTTP client.
|
||||
:param name: Agent name to create.
|
||||
:returns: The new session id.
|
||||
"""
|
||||
agent = await create_test_agent(client, name=name)
|
||||
resp = await client.post("/v1/sessions", json={"agent_id": agent["id"]})
|
||||
assert resp.status_code == 201, resp.text
|
||||
return resp.json()["id"]
|
||||
|
||||
|
||||
async def _launch(
|
||||
client: httpx.AsyncClient,
|
||||
session_id: str,
|
||||
*,
|
||||
git: dict[str, object] | None = None,
|
||||
) -> httpx.Response:
|
||||
"""POST the dedicated per-session bind+launch endpoint.
|
||||
|
||||
:param client: The test HTTP client.
|
||||
:param session_id: Existing session to bind.
|
||||
:param git: Optional ``git`` block. Create mode, e.g.
|
||||
``{"branch_name": "feature/x"}``; bind mode, e.g.
|
||||
``{"branch_name": "feature/x", "existing_worktree": True}``.
|
||||
:returns: The raw HTTP response.
|
||||
"""
|
||||
body: dict[str, object] = {"session_id": session_id, "workspace": _SOURCE_REPO}
|
||||
if git is not None:
|
||||
body["git"] = git
|
||||
return await client.post(f"/v1/hosts/{_HOST_ID}/runners", json=body)
|
||||
|
||||
|
||||
async def test_launch_runner_with_git_creates_worktree_and_persists_branch(
|
||||
register_host: RegisterHost,
|
||||
client: httpx.AsyncClient,
|
||||
db_uri: str,
|
||||
) -> None:
|
||||
"""``launch_runner`` with a ``git`` block creates a worktree off the
|
||||
source repo and binds the session to the worktree path + branch.
|
||||
|
||||
Proves the new worktree step on the dedicated bind endpoint: the
|
||||
request's branch reaches ``host.create_worktree``, and the resulting
|
||||
worktree path + branch are persisted on the (previously unbound)
|
||||
session row via the extended ``set_host_id``. Without the new code the
|
||||
session would bind to the source repo with ``git_branch=NULL``.
|
||||
"""
|
||||
cap = register_host()
|
||||
session_id = await _bare_session(client, "wt-launch-agent")
|
||||
|
||||
resp = await _launch(
|
||||
client, session_id, git={"branch_name": "feature/login", "base_branch": "main"}
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["runner_id"] # a runner was bound
|
||||
|
||||
# The host received exactly one create-worktree frame off the source
|
||||
# repo, carrying the requested branch + base ref.
|
||||
assert len(cap.create) == 1, f"expected one create_worktree frame, got {len(cap.create)}"
|
||||
assert cap.create[0].repo_path == _SOURCE_REPO
|
||||
assert cap.create[0].branch_name == "feature/login"
|
||||
assert cap.create[0].base_branch == "main"
|
||||
# Success path: no rollback.
|
||||
assert cap.remove == [], "worktree was rolled back on a successful launch"
|
||||
|
||||
# Persisted row: workspace is the worktree path (NOT the source repo),
|
||||
# git_branch is the new branch, host_id is bound. A NULL git_branch
|
||||
# here means set_host_id didn't receive/persist the branch.
|
||||
conv = SqlAlchemyConversationStore(db_uri).get_conversation(session_id)
|
||||
assert conv is not None
|
||||
assert conv.workspace == f"{_SOURCE_REPO}-worktrees/feature-login"
|
||||
assert conv.git_branch == "feature/login"
|
||||
assert conv.host_id == _HOST_ID
|
||||
|
||||
|
||||
async def test_launch_runner_without_git_binds_source_dir_no_worktree(
|
||||
register_host: RegisterHost,
|
||||
client: httpx.AsyncClient,
|
||||
db_uri: str,
|
||||
) -> None:
|
||||
"""Without a ``git`` block the endpoint binds the source directory
|
||||
directly and creates no worktree (the same-directory resume path).
|
||||
|
||||
Pins that the new worktree code is inert when ``git`` is omitted:
|
||||
no ``host.create_worktree`` frame, workspace stays the source repo,
|
||||
and ``git_branch`` stays NULL.
|
||||
"""
|
||||
cap = register_host()
|
||||
session_id = await _bare_session(client, "no-wt-agent")
|
||||
|
||||
resp = await _launch(client, session_id, git=None)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
assert cap.create == [], "no worktree should be created without a git block"
|
||||
conv = SqlAlchemyConversationStore(db_uri).get_conversation(session_id)
|
||||
assert conv is not None
|
||||
assert conv.workspace == _SOURCE_REPO
|
||||
assert conv.git_branch is None
|
||||
assert conv.host_id == _HOST_ID
|
||||
|
||||
|
||||
async def test_launch_runner_with_existing_worktree_persists_without_creating(
|
||||
register_host: RegisterHost,
|
||||
client: httpx.AsyncClient,
|
||||
db_uri: str,
|
||||
) -> None:
|
||||
"""``git.existing_worktree`` binds the existing worktree dir and records
|
||||
its branch without creating a worktree (the existing-worktree resume path).
|
||||
|
||||
The workspace is already a worktree, so no ``host.create_worktree``
|
||||
frame is sent; ``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_host()
|
||||
session_id = await _bare_session(client, "existing-wt-agent")
|
||||
|
||||
resp = await _launch(
|
||||
client,
|
||||
session_id,
|
||||
git={"branch_name": "feature/existing", "existing_worktree": True},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
assert cap.create == [], "no worktree should be created for an existing worktree"
|
||||
conv = SqlAlchemyConversationStore(db_uri).get_conversation(session_id)
|
||||
assert conv is not None
|
||||
assert conv.workspace == _SOURCE_REPO
|
||||
assert conv.git_branch == "feature/existing"
|
||||
assert conv.host_id == _HOST_ID
|
||||
|
||||
|
||||
async def test_launch_runner_rolls_back_worktree_on_launch_failure(
|
||||
register_host: RegisterHost,
|
||||
client: httpx.AsyncClient,
|
||||
db_uri: str,
|
||||
) -> None:
|
||||
"""When the host fails the launch, the just-created worktree is
|
||||
rolled back AND the runner binding is cleared so the picker can retry.
|
||||
|
||||
The worktree is created (status ok) but the launch reports
|
||||
``failed`` → the endpoint returns 502, sends a
|
||||
``host.remove_worktree`` for the created worktree, and clears the
|
||||
session's ``runner_id``. If the rollback were missing, ``cap.remove``
|
||||
would be empty; if the binding weren't cleared, ``runner_id`` would
|
||||
stay set and a retry would dead-end on the atomic ``set_runner_id``
|
||||
CAS with "session already has a runner bound" (the whole point of the
|
||||
fork-resume picker is that the user can retry after a failed bind).
|
||||
"""
|
||||
cap = register_host(launch_status="failed")
|
||||
session_id = await _bare_session(client, "wt-rollback-agent")
|
||||
|
||||
resp = await _launch(client, session_id, git={"branch_name": "feature/x"})
|
||||
|
||||
# Launch failed → 502 (host fault), as the no-git path already does.
|
||||
assert resp.status_code == 502, resp.text
|
||||
# The worktree was created, then removed (rollback fired) for the same path.
|
||||
assert len(cap.create) == 1
|
||||
assert len(cap.remove) == 1, "expected a rollback remove_worktree frame after launch failure"
|
||||
created_path = f"{_SOURCE_REPO}-worktrees/feature-x"
|
||||
assert cap.remove[0].worktree_path == created_path
|
||||
assert cap.remove[0].delete_branch is True # orphan branch also cleaned up
|
||||
|
||||
# The session is fully unbound so the DB matches the host's actual
|
||||
# state (worktree removed) and a retry starts clean. A non-None
|
||||
# runner_id would stick the session as "already bound"; a leftover
|
||||
# workspace/git_branch would point at the deleted worktree/branch and
|
||||
# could wrongly trigger worktree-cleanup paths (git_branch IS NOT NULL).
|
||||
conv = SqlAlchemyConversationStore(db_uri).get_conversation(session_id)
|
||||
assert conv is not None
|
||||
assert conv.runner_id is None, (
|
||||
"runner_id should be cleared after a failed launch so the picker can "
|
||||
f"rebind; got {conv.runner_id!r} (retry would 400 'already has a runner bound')"
|
||||
)
|
||||
assert conv.host_id is None, f"host_id should be cleared on rollback; got {conv.host_id!r}"
|
||||
assert conv.workspace is None, (
|
||||
f"workspace should be cleared (worktree was removed); got {conv.workspace!r}"
|
||||
)
|
||||
assert conv.git_branch is None, (
|
||||
f"git_branch should be cleared (branch was removed); got {conv.git_branch!r}"
|
||||
)
|
||||
|
||||
|
||||
async def test_launch_runner_retry_succeeds_after_failed_launch(
|
||||
register_host: RegisterHost,
|
||||
client: httpx.AsyncClient,
|
||||
db_uri: str,
|
||||
) -> None:
|
||||
"""A second bind succeeds after the first launch failed.
|
||||
|
||||
End-to-end proof of the cleared-binding fix: a failed launch (502)
|
||||
must leave the session re-bindable. The retry creates a fresh
|
||||
worktree and binds the runner. Without clearing ``runner_id`` on the
|
||||
first failure, this retry returns 400 "session already has a runner
|
||||
bound" — the dead-end the fork-resume picker would otherwise hit.
|
||||
"""
|
||||
register_host(launch_status="failed")
|
||||
session_id = await _bare_session(client, "wt-retry-agent")
|
||||
|
||||
first = await _launch(client, session_id, git={"branch_name": "feature/x"})
|
||||
assert first.status_code == 502, first.text
|
||||
|
||||
# Re-register the host to launch successfully this time (newest-wins
|
||||
# replaces the failing connection), then retry the bind.
|
||||
cap_ok = register_host(launch_status="launched")
|
||||
second = await _launch(client, session_id, git={"branch_name": "feature/y"})
|
||||
|
||||
# Retry binds cleanly — proves runner_id was released by the failure.
|
||||
assert second.status_code == 200, second.text
|
||||
assert second.json()["runner_id"]
|
||||
assert len(cap_ok.create) == 1, "retry created a fresh worktree off the source repo"
|
||||
conv = SqlAlchemyConversationStore(db_uri).get_conversation(session_id)
|
||||
assert conv is not None
|
||||
assert conv.workspace == f"{_SOURCE_REPO}-worktrees/feature-y"
|
||||
assert conv.git_branch == "feature/y"
|
||||
Reference in New Issue
Block a user