fix: prune stale tracking refs before pushing so re-split branches are not rejected as stale

pr_split/cli.py

@@ -64,6 +64,7 @@
     fetch_fork_pr,
     is_worktree_clean,
     merge_base,
+    prune_remote_tracking_refs,
     push_branch,
     remove_worktree,
 )
@@ -461,7 +462,15 @@ def _create_branches_and_commits(
     *,
     author: str | None = None,
     stacked: bool = False,
+    keep: dict[str, BranchRecord] | None = None,
 ) -> list[BranchRecord]:
+    """Create a branch and commit per group; groups in ``keep`` reuse their record.
+
+    ``keep`` holds groups whose PRs already exist from an earlier, partly
+    failed run: their branches are left as they are (recreating them would
+    rewrite the open PR's head).
+    """
+    kept = keep or {}
     worktree_base = Path(tempfile.mkdtemp(prefix="pr-split-worktrees-"))
     step = load_per_group_step()
     order = {gid: i for i, gid in enumerate(PlanDAG(groups).topological_order(), start=1)}
@@ -482,9 +491,10 @@ def _create_branches_and_commits(
         batches = iter([[(group, base_branch, merge_base_ref) for group in groups]])
 
     try:
-        results: dict[str, BranchRecord] = {}
+        results: dict[str, BranchRecord] = dict(kept)
         errors: list[tuple[str, Exception]] = []
         for batch_args in batches:
+            batch_args = [args for args in batch_args if args[0].id not in kept]
             with ThreadPoolExecutor(max_workers=_WORKTREE_MAX_WORKERS) as executor:
                 future_to_group_id = {
                     executor.submit(
@@ -514,7 +524,9 @@ def _create_branches_and_commits(
         if errors:
             # Failed groups already removed their own branch inside the
             # worker; only the successful ones remain to roll back.
-            for record in results.values():
+            for gid, record in results.items():
+                if gid in kept:
+                    continue
                 try:
                     delete_branch(record.branch_name)
                 except PRSplitError as exc:
@@ -625,17 +637,26 @@ def _push_and_create_prs(
     branch_records: list[BranchRecord],
     *,
     draft: bool = False,
+    existing_prs: dict[str, PRRecord] | None = None,
 ) -> list[PRRecord]:
+    """Push each group's branch and open its PR; groups in ``existing_prs`` are done."""
+    done = existing_prs or {}
     record_map = {r.group_id: r for r in branch_records}
     errors: list[tuple[str, Exception]] = []
 
+    # Branch names are reused across runs; drop tracking refs left by an
+    # earlier split whose remote branch was merged or deleted since, or
+    # force-with-lease rejects every push as stale.
+    prune_remote_tracking_refs()
+
     # Children target parent branches, so every branch is pushed before any PR opens.
     with ThreadPoolExecutor(max_workers=_PUSH_MAX_WORKERS) as executor:
         push_futures = {
             executor.submit(push_branch, record_map[group.id].branch_name): group.id
             for group in groups
+            if group.id not in done
         }
-        pushed: set[str] = set()
+        pushed: set[str] = set(done)
         for future in as_completed(push_futures):
             group_id = push_futures[future]
             try:
@@ -670,9 +691,9 @@ def _base_pushed(group: Group) -> bool:
                 _create_single_pr, group, record_map[group.id], groups, draft=draft
             ): group.id
             for group in groups
-            if group.id in pushed and _base_pushed(group)
+            if group.id not in done and group.id in pushed and _base_pushed(group)
         }
-        results: dict[str, PRRecord] = {}
+        results: dict[str, PRRecord] = dict(done)
         for future in as_completed(future_to_group_id):
             group_id = future_to_group_id[future]
             try:
@@ -1880,10 +1901,25 @@ def execute(
     if draft and not plan.draft:
         plan = plan.model_copy(update={"draft": True})
 
-    if plan_file.git_state.prs:
+    existing_prs = {r.group_id: r for r in plan_file.git_state.prs}
+    recorded = {r.group_id: r for r in plan_file.git_state.branches}
+    kept_branches = {gid: recorded[gid] for gid in existing_prs if gid in recorded}
+    group_ids = {g.id for g in plan.groups}
+    if existing_prs and (
+        set(existing_prs) >= group_ids
+        or not set(existing_prs) <= group_ids
+        or set(kept_branches) != set(existing_prs)
+    ):
+        # Every group already has a PR, or the recorded PRs no longer match
+        # the plan's groups: there is nothing safe to resume.
         console.print("[red]This plan already has PRs. Use 'pr-split clean' first.[/red]")
         raise typer.Exit(1)
-    if plan_file.git_state.branches:
+    if existing_prs:
+        console.print(
+            f"[yellow]A previous run opened {len(existing_prs)} of {len(plan.groups)} PR(s)."
+            " Keeping those and creating the rest.[/yellow]"
+        )
+    elif plan_file.git_state.branches:
         # A previous execute created branches but no PRs - a failed push or PR
         # creation. Branch creation is idempotent (add_worktree recreates an
         # existing branch), so retry instead of forcing 'clean' + a full
@@ -1942,6 +1978,16 @@ def execute(
         console.print(f"[red]{exc}[/red]")
         raise typer.Exit(1) from exc
 
+    # Any plan with dependency edges is laid out along its DAG (stacked or not),
+    # so a kept layer's branch sits on its ancestors' current commits; rebuilding
+    # an ancestor would rewrite the base under the kept layer's open PR, which
+    # would then show the ancestor's changes again. Their PRs are still opened.
+    if plan.stacked or any(g.depends_on for g in plan.groups):
+        for gid in list(kept_branches):
+            for ancestor in dag.ancestors(gid):
+                if ancestor in recorded:
+                    kept_branches.setdefault(ancestor, recorded[ancestor])
+
     _present_plan(plan.groups)
     _report_oversized_groups(
         plan.groups, plan.max_loc, {pf.path: len(pf) for pf in parsed_diff.patch_set}
@@ -1960,12 +2006,15 @@ def execute(
             namespace,
             author=plan.author,
             stacked=plan.stacked,
+            keep=kept_branches,
         )
     except PRSplitError as exc:
         console.print(f"[red]{exc}[/red]")
         raise typer.Exit(1) from exc
     try:
-        pr_records = _push_and_create_prs(plan.groups, branch_records, draft=plan.draft)
+        pr_records = _push_and_create_prs(
+            plan.groups, branch_records, draft=plan.draft, existing_prs=existing_prs
+        )
     except PRCreationError as exc:
         save_plan(
             PlanFile(

pr_split/exceptions.py

@@ -102,6 +102,11 @@ class ErrorMsg(StrEnum):
         "Cannot tell which base branch the stack under '{prefix}' targets ({bases}); pass --base"
     )
     RECOVER_PR_LIST_FAILED = "Cannot list PRs to rebuild the plan ({detail})"
+    LLM_TOKEN_COUNT_FAILED = "Could not count prompt tokens with the LLM API: {detail}"
+    LLM_KEY_REJECTED = (
+        "{variable} was rejected by the LLM API; set a valid key, or plan without one"
+        " using --partition-strategy graph"
+    )
     BRANCH_CREATE_FAILED = "Failed to create branch '{branch}': {detail}"
     PR_CREATE_FAILED = "Failed to create PR for group '{group}': {detail}"
     MERGE_FAILED = "Merge of '{source}' into '{target}' failed: {detail}"

pr_split/git_ops/__init__.py

@@ -25,6 +25,9 @@
 from .branches import (
     merge_base as merge_base,
 )
+from .branches import (
+    prune_remote_tracking_refs as prune_remote_tracking_refs,
+)
 from .branches import (
     push_branch as push_branch,
 )

pr_split/git_ops/branches.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import contextlib
 import os
 import re
 import subprocess
+import time
 
 from loguru import logger
 
@@ -93,9 +95,50 @@ def is_worktree_clean() -> bool:
     return all(line.startswith("??") for line in output.splitlines())
 
 
+# Server-side failures GitHub reports for a push that can succeed on retry.
+_TRANSIENT_PUSH_ERRORS = (
+    "fatal error in commit_refs",
+    "the remote end hung up unexpectedly",
+    "internal server error",
+    "http 500",
+    "http 502",
+    "http 503",
+    "connection reset",
+    "operation timed out",
+)
+_PUSH_ATTEMPTS = 3
+_PUSH_RETRY_DELAY = 2.0
+
+
+def _remote_has_local_head(branch: str) -> bool:
+    try:
+        remote = run_git("ls-remote", "origin", f"refs/heads/{branch}").split()
+        return bool(remote) and remote[0] == run_git("rev-parse", branch)
+    except GitOperationError:
+        return False
+
+
 def push_branch(branch: str) -> None:
     logger.info(logs.PUSHING_BRANCH.format(branch=branch))
-    run_git("push", "--force-with-lease", "-u", "origin", branch)
+    retried = False
+    for attempt in range(1, _PUSH_ATTEMPTS + 1):
+        try:
+            run_git("push", "--force-with-lease", "-u", "origin", branch)
+            return
+        except GitOperationError as exc:
+            transient = any(marker in str(exc).lower() for marker in _TRANSIENT_PUSH_ERRORS)
+            # A push reported as failed (hung-up remote, timeout) may still have
+            # landed; the retry is then rejected as a stale lease although the
+            # branch is already on the remote.
+            if retried and not transient and _remote_has_local_head(branch):
+                return
+            if not transient or attempt == _PUSH_ATTEMPTS:
+                raise
+            retried = True
+            logger.warning(
+                logs.PUSH_RETRY.format(branch=branch, attempt=attempt, error=str(exc).strip())
+            )
+            time.sleep(_PUSH_RETRY_DELAY * attempt)
 
 
 _LOCAL_BRANCH_MISSING = "not found"
@@ -129,10 +172,36 @@ def delete_branch(branch: str, *, remote: bool = False) -> None:
             if _REMOTE_REF_MISSING not in str(exc):
                 raise
             logger.info(logs.BRANCH_ALREADY_GONE.format(branch=branch, where=" on origin"))
+        # Either way the remote branch is gone; drop the local tracking ref
+        # too, or the next `push --force-with-lease` of a reused branch name
+        # is rejected as "stale info" against a ref origin no longer has.
+        forget_remote_tracking_ref(branch)
     if local_error is not None:
         raise local_error
 
 
+def forget_remote_tracking_ref(branch: str) -> None:
+    with contextlib.suppress(GitOperationError):
+        run_git("update-ref", "-d", f"refs/remotes/origin/{branch}")
+
+
+def prune_remote_tracking_refs() -> None:
+    """Drop tracking refs whose remote branch no longer exists.
+
+    Split branch names are reused across runs; after a merge or `clean` the
+    remote branch is gone but `refs/remotes/origin/pr-split/...` may still
+    point at the old head, and every `push --force-with-lease` would be
+    rejected as stale. Only *gone* refs are dropped (`git remote prune`, no
+    fetch): refreshing the refs that still exist would make the lease
+    compare against origin's current tip and silently force over commits
+    someone else pushed to a reused branch.
+    """
+    try:
+        run_git("remote", "prune", "origin")
+    except GitOperationError as exc:
+        logger.warning(logs.PRUNE_FAILED.format(error=exc))
+
+
 def merge_base(ref_a: str, ref_b: str) -> str:
     return run_git("merge-base", ref_a, ref_b)
 

pr_split/logs.py

@@ -148,6 +148,7 @@
     "CP-SAT stopped at the {timeout:g}s limit with a feasible but unproven-optimal plan "
     "({units} units, {groups} groups); raise --cp-sat-timeout for a better partition"
 )
+PUSH_RETRY = "Push of {branch} failed (attempt {attempt}): {error}; retrying"
 BASE_ALREADY_MERGED = (
     "The split's base branch {base} has already merged; run 'pr-split retarget' to move"
     " the split onto the new base before merging"
@@ -161,3 +162,7 @@
     " its sub-PR will not compile or import on its own"
 )
 ADOPTED_REMOTE_BRANCH = "Created local branch {branch} from {remote_ref}"
+PRUNE_FAILED = (
+    "Could not prune origin's tracking refs before pushing ({error}); "
+    "a reused branch name may be rejected as stale"
+)

pr_split/planner/client.py

@@ -105,12 +105,20 @@ def _extract_raw_output(block_input: dict[str, object]) -> list[_RawGroup]:
 
 def _count_tokens_anthropic(system: str, user: str, *, settings: Settings) -> int:
     client = anthropic.Anthropic(api_key=settings.api_key)
-    response = client.messages.count_tokens(
-        model=settings.model,
-        system=system,
-        messages=[{"role": "user", "content": user}],
-        tools=[_ANTHROPIC_TOOL_DEF],
-    )
+    # This is the first API call of a run, so a bad key, a network blip or a
+    # rate limit surfaces here; it must be an LLMError like every other
+    # provider failure rather than an SDK traceback.
+    try:
+        response = client.messages.count_tokens(
+            model=settings.model,
+            system=system,
+            messages=[{"role": "user", "content": user}],
+            tools=[_ANTHROPIC_TOOL_DEF],
+        )
+    except anthropic.AuthenticationError as exc:
+        raise LLMError(ErrorMsg.LLM_KEY_REJECTED(variable="ANTHROPIC_API_KEY")) from exc
+    except anthropic.APIError as exc:
+        raise LLMError(ErrorMsg.LLM_TOKEN_COUNT_FAILED(detail=str(exc))) from exc
     return response.input_tokens
 
 

tests/network_guards.py

@@ -13,6 +13,14 @@ def _no_base_fetch() -> Iterator[None]:
         yield
 
 
+@pytest.fixture(autouse=True)
+def _no_remote_prune() -> Iterator[None]:
+    """Pushing prunes origin's tracking refs first; keep that off the network
+    (and off this checkout's refs). Tests that exercise it patch it again."""
+    with patch("pr_split.cli.prune_remote_tracking_refs"):
+        yield
+
+
 @pytest.fixture(autouse=True)
 def _no_selected_plan() -> Iterator[None]:
     """A command's plan selection is process-wide; start each test without one."""