Fix #2410: plumb slice_id through restart and uncommitted-change paths

orchestrator/kubernetes_spawner.py

@@ -323,11 +323,15 @@ def __init__(
         self._k8s = k8s_client
         self._gateway = gateway_client
         self._namespace = namespace
-        # Track restart counts per (pipeline_id, agent_role) pair
-        self._restart_counts: dict[tuple[str, str], int] = {}
-        # Per-(pipeline_id, agent_role) locks for serialising concurrent restarts.
-        # Protected by _restart_locks_lock (same pattern as state_store.py).
-        self._restart_locks: dict[tuple[str, str], threading.Lock] = {}
+        # Track restart counts per (pipeline_id, agent_role, slice_id) tuple.
+        # ``slice_id`` is ``None`` for pipeline-level agents and
+        # ``"slice-<N>"`` for slice-scoped agents (#2410), so concurrent
+        # slices each get an independent budget.
+        self._restart_counts: dict[tuple[str, str, str | None], int] = {}
+        # Per-(pipeline_id, agent_role, slice_id) locks for serialising
+        # concurrent restarts. Protected by _restart_locks_lock (same
+        # pattern as state_store.py).
+        self._restart_locks: dict[tuple[str, str, str | None], threading.Lock] = {}
         self._restart_locks_lock = threading.Lock()
 
     @property
@@ -356,8 +360,8 @@ def gateway(self) -> GatewayClient:
             self._gateway = get_gateway_client()
         return self._gateway
 
-    def _get_restart_lock(self, key: tuple[str, str]) -> threading.Lock:
-        """Get or create a per-(pipeline_id, agent_role) restart lock."""
+    def _get_restart_lock(self, key: tuple[str, str, str | None]) -> threading.Lock:
+        """Get or create a per-(pipeline_id, agent_role, slice_id) restart lock."""
         with self._restart_locks_lock:
             if key not in self._restart_locks:
                 self._restart_locks[key] = threading.Lock()
@@ -1044,6 +1048,7 @@ def restart_agent_job(
         reason: str = "",
         spawn_max_retries: int = DEFAULT_SPAWN_MAX_RETRIES,
         spawn_retry_initial_backoff_seconds: float = (DEFAULT_SPAWN_RETRY_INITIAL_BACKOFF_SECONDS),
+        slice_id: str | None = None,
     ) -> SpawnedContainer:
         """Restart an agent Job: delete and respawn preserving worktree.
 
@@ -1067,6 +1072,13 @@ def restart_agent_job(
                 during worktree creation (forwarded to ``spawn_agent_job``).
             spawn_retry_initial_backoff_seconds: Initial backoff for spawn
                 retries (forwarded to ``spawn_agent_job``).
+            slice_id: Optional slice scope (#2410). When supplied, the
+                slice-scoped Job name (``egg-agent-{pid}-{slice_id}-{role}``)
+                is the one deleted and respawned, the slice-scoped worktree
+                id is preserved, and ``EGG_SLICE_ID`` is propagated so the
+                restarted agent re-enters the per-slice consensus tracker.
+                The restart-budget key includes the slice scope so each
+                slice gets an independent budget.
 
         Returns:
             SpawnedContainer with new Job info.
@@ -1078,7 +1090,11 @@ def restart_agent_job(
         if mode is None:
             raise ValueError("mode must be explicitly provided ('public' or 'private')")
 
-        restart_key = (pipeline_id, agent_role.value)
+        # Slice scope is part of the restart key so concurrent slice-N
+        # and slice-M agents of the same role each get an independent
+        # restart budget and lock. ``reset_restart_counts(pipeline_id)``
+        # still clears all of them because it filters on ``k[0]``.
+        restart_key = (pipeline_id, agent_role.value, slice_id)
         lock = self._get_restart_lock(restart_key)
 
         # Timeout prevents indefinite blocking if a concurrent restart of the
@@ -1106,7 +1122,12 @@ def restart_agent_job(
             # spawn time (hyphenated, no JOB_PREFIX); ``actual_k8s_job_name``
             # is the real k8s Job name. Using the wrong form for either side
             # broke restart for every role with an underscore — see #2070.
-            job_name, actual_k8s_job_name = self._build_k8s_job_names(pipeline_id, agent_role)
+            # Slice scope (#2410) must be threaded through here so the
+            # delete + respawn target the slice-scoped Job name, not the
+            # pipeline-level one.
+            job_name, actual_k8s_job_name = self._build_k8s_job_names(
+                pipeline_id, agent_role, slice_id=slice_id
+            )
 
             logger.info(
                 "Restarting agent Job",
@@ -1148,7 +1169,10 @@ def restart_agent_job(
                     error=str(e),
                 )
 
-            # Respawn — gateway's create_worktrees() is idempotent
+            # Respawn — gateway's create_worktrees() is idempotent.
+            # ``slice_id`` is forwarded so spawn_agent_job builds the
+            # slice-scoped Job + worktree id and sets ``EGG_SLICE_ID``
+            # on the new Job (#2410).
             spawned = self.spawn_agent_job(
                 pipeline_id=pipeline_id,
                 agent_role=agent_role,
@@ -1167,6 +1191,7 @@ def restart_agent_job(
                 preserve_worktree_on_failure=True,
                 spawn_max_retries=spawn_max_retries,
                 spawn_retry_initial_backoff_seconds=spawn_retry_initial_backoff_seconds,
+                slice_id=slice_id,
             )
 
             logger.info(
@@ -1181,17 +1206,26 @@ def restart_agent_job(
         finally:
             lock.release()
 
-    def get_restart_count(self, pipeline_id: str, agent_role: str) -> int:
+    def get_restart_count(
+        self,
+        pipeline_id: str,
+        agent_role: str,
+        slice_id: str | None = None,
+    ) -> int:
         """Get the current restart count for an agent.
 
         Args:
             pipeline_id: Pipeline ID.
             agent_role: Agent role value string.
+            slice_id: Optional slice scope (#2410). Pipeline-level callers
+                pass ``None``; slice-aware callers pass the same
+                ``slice-<N>`` string they used at restart time so each
+                slice's budget is reported independently.
 
         Returns:
             Number of times the agent has been restarted.
         """
-        key = (pipeline_id, agent_role)
+        key = (pipeline_id, agent_role, slice_id)
         lock = self._get_restart_lock(key)
         with lock:
             return self._restart_counts.get(key, 0)
@@ -1218,19 +1252,33 @@ def detect_uncommitted_changes(
         self,
         pipeline_id: str,
         agent_role: str,
+        slice_id: str | None = None,
     ) -> dict | None:
         """Detect uncommitted changes in an agent's worktree after Job exit.
 
         Checks the agent's worktree directly on the filesystem for uncommitted
         changes. Per-agent worktrees are at:
-        /home/egg/.egg-worktrees/{pipeline_id}-{role}/{repo}/
+        /home/egg/.egg-worktrees/{pipeline_id}-{role}/{repo}/ (pipeline-level)
+        /home/egg/.egg-worktrees/{pipeline_id}-{slice_id}-{role}/{repo}/
+        (slice-scoped, #2410).
+
+        Args:
+            pipeline_id: Pipeline ID.
+            agent_role: Agent role value string.
+            slice_id: Optional slice scope. When supplied, the slice-scoped
+                worktree id is inspected; pipeline-level callers omit this.
 
         Returns:
             Dict with change info if uncommitted changes found, None otherwise.
         """
         import subprocess
 
-        agent_worktree_id = f"{pipeline_id}-{agent_role}"
+        # Mirrors ``_build_agent_worktree_id`` so a slice-scoped restart
+        # path can detect uncommitted work in the slice's worktree, not
+        # the (possibly absent) pipeline-level one.
+        agent_worktree_id = (
+            f"{pipeline_id}-{slice_id}-{agent_role}" if slice_id else f"{pipeline_id}-{agent_role}"
+        )
         worktree_base = WORKTREE_BASE_DIR / agent_worktree_id
 
         if not worktree_base.exists():
@@ -1269,13 +1317,15 @@ def detect_uncommitted_changes(
                         event_type="agent_uncommitted_changes",
                         pipeline_id=pipeline_id,
                         agent_role=agent_role,
+                        slice_id=slice_id,
                         worktree_path=str(repo_dir),
                         file_count=len(files),
                         changed_files=files[:20],
                     )
                     return {
                         "pipeline_id": pipeline_id,
                         "agent_role": agent_role,
+                        "slice_id": slice_id,
                         "worktree_id": agent_worktree_id,
                         "worktree_path": str(repo_dir),
                         "file_count": len(files),

orchestrator/routes/pipelines.py

@@ -107,6 +107,7 @@ def get_repo_checks(repo: str) -> list[dict[str, str]]:  # type: ignore[misc]
         PipelineStatus,
         ReviewVerdict,
     )
+    from ..slice_id_validation import extract_slice_id
     from ..state_store import (
         InvalidPipelineIdError,
         PipelineNotFoundError,
@@ -153,6 +154,7 @@ def get_repo_checks(repo: str) -> list[dict[str, str]]:  # type: ignore[misc]
         PipelineStatus,
         ReviewVerdict,
     )
+    from slice_id_validation import extract_slice_id  # type: ignore
     from state_store import (  # type: ignore
         InvalidPipelineIdError,
         PipelineNotFoundError,
@@ -2270,9 +2272,17 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]:
         pipeline_id: Pipeline ID
         agent_role: Agent role to restart (e.g. "coder", "tester")
 
+    Query string (optional):
+        slice_id: Slice scope (``slice-<N>``). When supplied, the
+            slice-scoped Job and worktree are restarted, ``EGG_SLICE_ID``
+            is propagated to the new Job, and consensus reset targets
+            the per-slice tracker. Pipeline-level agents omit it.
+            ``slice_id`` may also be supplied via the JSON body.
+
     Request body (optional):
         {
-            "reason": "Human-readable reason for the restart"
+            "reason": "Human-readable reason for the restart",
+            "slice_id": "slice-2"
         }
 
     Response:
@@ -2317,6 +2327,16 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]:
     body = request.get_json(silent=True) or {}
     reason = body.get("reason", "Manual restart via API")
 
+    # Slice scope (#2410): query param wins over body so the URL
+    # form is unambiguous; both forms validate against the canonical
+    # ``slice-<N>`` shape via ``extract_slice_id``.
+    raw_slice_id = request.args.get("slice_id")
+    slice_payload = {"slice_id": raw_slice_id} if raw_slice_id is not None else body
+    try:
+        slice_id = extract_slice_id(slice_payload)
+    except ValueError as e:
+        return make_error_response(str(e), status_code=400)
+
     # Restart the container via spawner
     spawner = _get_spawner()
 
@@ -2423,6 +2443,7 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]:
             reason=reason,
             spawn_max_retries=pipeline.config.spawn_max_retries,
             spawn_retry_initial_backoff_seconds=pipeline.config.spawn_retry_initial_backoff_seconds,
+            slice_id=slice_id,
         )
     except (ContainerSpawnError, KubernetesSpawnError) as e:
         # Revert early status update — the agent is not actually running.
@@ -2442,6 +2463,8 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]:
     # Spawn succeeded — now reset consensus state for this agent.
     # If consensus reset fails, log a warning but don't fail the restart:
     # the restarted agent will re-enter consensus on its own.
+    # Slice-scoped restarts (#2410) target the per-slice tracker; the
+    # pipeline-level tracker has no record of the slice agent.
     try:
         try:
             from peer_consensus import get_peer_consensus_tracker
@@ -2450,13 +2473,14 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]:
                 get_peer_consensus_tracker,  # type: ignore[import-not-found]
             )
 
-        tracker = get_peer_consensus_tracker(pipeline_id)
+        tracker = get_peer_consensus_tracker(pipeline_id, slice_id)
         if tracker:
             tracker.remove_agent(agent_role)
             logger.info(
                 "Reset consensus state for agent",
                 pipeline_id=pipeline_id,
                 agent_role=agent_role,
+                slice_id=slice_id,
             )
     except ImportError:
         pass
@@ -2465,6 +2489,7 @@ def restart_agent(pipeline_id: str, agent_role: str) -> tuple[Response, int]:
             "Failed to reset consensus state (agent will re-enter consensus)",
             pipeline_id=pipeline_id,
             agent_role=agent_role,
+            slice_id=slice_id,
             error=str(e),
         )
 

orchestrator/routes/signals.py

@@ -39,6 +39,19 @@ def get_logger(name: str, **kwargs) -> logging.Logger:  # type: ignore[misc]
 from egg_contracts.orchestrator import create_orchestrator
 from handoffs import AgentOutput, save_agent_output
 from models import AgentExecutionStatus, AgentRole, Pipeline, PipelineStatus
+
+# Slice-aware consensus routing (#2403): per-slice agents tag their
+# signals with a ``slice_id`` field on the request body so the
+# orchestrator routes each ``CONSENSUS_*`` to the per-slice tracker
+# (``peer_consensus._tracker_key`` composes ``{pipeline_id}/{slice_id}``).
+# The canonical extractor lives in ``slice_id_validation`` so the
+# operator-triggered restart route (#2410) and the gateway-bound branch
+# builders in ``concurrent_executor`` validate against the same shape.
+# The alias below preserves the existing private name for the many
+# handler call sites in this file.
+from slice_id_validation import (
+    extract_slice_id as _extract_slice_id,
+)
 from state_store import (
     InvalidPipelineIdError,
     PipelineNotFoundError,
@@ -65,53 +78,6 @@ def get_logger(name: str, **kwargs) -> logging.Logger:  # type: ignore[misc]
 
 _SIGTERM_PATTERN = re.compile(r"\b143\b")
 
-# Slice-aware consensus routing (#2403): per-slice agents tag their
-# signals with a ``slice_id`` field on the request body so the
-# orchestrator routes each ``CONSENSUS_*`` to the per-slice tracker
-# (``peer_consensus._tracker_key`` composes ``{pipeline_id}/{slice_id}``).
-# Defense-in-depth: validate the shape here even though the spawn-side
-# is already canonical — a future caller that forgets the upstream
-# regex must not be able to smuggle path separators or shell
-# metacharacters into a tracker registry key.
-#
-# Note: the contract-side ``Slice.id`` field
-# (``shared/egg_contracts/models.py``) accepts the broader pattern
-# ``^(?:slice|phase)-[0-9]+$`` for backward compatibility with
-# pre-#2137 contracts. The signal-side pattern below is intentionally
-# narrower — only the canonical ``slice-<N>`` shape. The
-# canonicalisation is performed by the parent ``Contract`` model
-# validator ``Contract._migrate_phases_to_slices`` (mode="wrap" — runs
-# before per-Slice field validation), which rewrites legacy
-# ``phase-<N>`` ids to ``slice-<N>`` whenever a contract is loaded via
-# ``Contract.model_validate(json_dict)``. Slices reaching the spawn /
-# signal path through the contract-loader path are therefore already
-# canonical. The guarantee does NOT extend to direct Slice
-# construction (e.g. ``Slice(id="phase-2", ...)``) — pydantic field
-# validation alone is permissive — so the regex below is doing real
-# work for any code path that constructs Slices outside Contract
-# loading (e.g. a future migration tool, a hand-built fixture, or
-# direct model_validate on a Slice dict). Such a caller's
-# ``phase-<N>`` slice will be rejected here as it should — the
-# registry key MUST be canonical so the per-slice tracker can be
-# looked up.
-_SLICE_ID_PATTERN = re.compile(r"^slice-[0-9]+$")
-
-
-def _extract_slice_id(data: dict[str, Any]) -> str | None:
-    """Return the ``slice_id`` from the signal payload, validated.
-
-    Returns ``None`` when no slice scope was supplied (the agent is
-    pipeline-level, not slice-scoped). Raises ``ValueError`` when the
-    caller supplied a non-empty value that does not match the
-    canonical ``slice-<N>`` shape.
-    """
-    raw = data.get("slice_id")
-    if raw is None or raw == "":
-        return None
-    if not isinstance(raw, str) or not _SLICE_ID_PATTERN.fullmatch(raw):
-        raise ValueError(f"Invalid slice_id {raw!r}: must match 'slice-<N>'")
-    return raw
-
 
 def _is_sigterm_after_completion(pipeline: Pipeline, error_message: str) -> bool:
     """Return True if this error is a SIGTERM exit on an already-complete pipeline."""