Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 59 additions & 21 deletions orchestrator/kubernetes_spawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,34 @@ class KubernetesSpawner:
DEFAULT_SANDBOX_IMAGE = os.environ.get("EGG_SANDBOX_IMAGE", "egg:latest")
JOB_NAME_FORMAT = "egg-agent-{pipeline_id}-{role}"

@classmethod
def _build_k8s_job_names(
cls,
pipeline_id: str,
agent_role: AgentRole,
) -> tuple[str, str]:
"""Build the two identifiers an agent Job is known by.

Returns a ``(job_name, actual_k8s_job_name)`` pair where:

- ``job_name`` is the unprefixed identifier used as the gateway
session ``container_id`` and in labels (e.g.
``egg-agent-issue-1962-task-planner``).
- ``actual_k8s_job_name`` is the real k8s Job name after
``KubernetesClient`` prepends ``JOB_PREFIX`` during
``create_container`` (e.g.
``egg-sandbox-egg-agent-issue-1962-task-planner``).

Underscores in ``agent_role.value`` (``task_planner``,
``reviewer_refine``, …) are converted to hyphens because k8s
resource names are RFC-1123 labels and reject underscores.
"""
job_name = cls.JOB_NAME_FORMAT.format(
pipeline_id=pipeline_id,
role=agent_role.value.replace("_", "-"),
)
return job_name, f"{KubernetesClient.JOB_PREFIX}{job_name}"

def __init__(
self,
k8s_client: KubernetesClient | None = None,
Expand Down Expand Up @@ -367,21 +395,9 @@ def spawn_agent_job(
Raises:
KubernetesSpawnError: If spawning fails
"""
job_name = self.JOB_NAME_FORMAT.format(
pipeline_id=pipeline_id,
# k8s names are RFC-1123 labels: no underscores allowed.
# Role enum values like "reviewer_refine" need hyphenation.
role=agent_role.value.replace("_", "-"),
)
job_name, actual_k8s_job_name = self._build_k8s_job_names(pipeline_id, agent_role)

# Clean up any existing Job with the same name.
# create_container() prepends JOB_PREFIX, so derive the actual k8s
# Job name that would have been created in a previous spawn.
actual_k8s_job_name = (
job_name
if job_name.startswith(KubernetesClient.JOB_PREFIX)
else f"{KubernetesClient.JOB_PREFIX}{job_name}"
)
try:
self.k8s.delete_job(actual_k8s_job_name, self._namespace)
logger.info(
Expand Down Expand Up @@ -1014,10 +1030,11 @@ def restart_agent_job(
# Increment count before spawn so failed attempts burn a restart budget slot
self._restart_counts[restart_key] = current_count + 1

job_name = self.JOB_NAME_FORMAT.format(
pipeline_id=pipeline_id,
role=agent_role.value,
)
# ``job_name`` matches the gateway session container_id used at
# 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)

logger.info(
"Restarting agent Job",
Expand All @@ -1028,12 +1045,33 @@ def restart_agent_job(
reason=reason,
)

# Delete the existing Job (best effort)
# Delete the existing Job (best effort) and clean up the
# gateway session. We can't go through ``remove_agent_job``
# here because it would route both the k8s and gateway calls
# through the same identifier, but k8s wants the prefixed form
# and the gateway session is keyed by the unprefixed form.
try:
self.remove_agent_job(job_name, force=True, cleanup_session=True)
except (PodNotFoundError, JobOperationError) as e:
logger.info(
self.k8s.delete_job(
actual_k8s_job_name,
self._namespace,
propagation_policy="Foreground",
)
except PodNotFoundError:
logger.debug(
"No existing Job found during restart (already removed)",
job_name=actual_k8s_job_name,
)
except JobOperationError as e:
logger.warning(
"Failed to delete existing Job during restart, continuing",
job_name=actual_k8s_job_name,
error=str(e),
)
try:
self.gateway.delete_session_by_container(job_name)
except GatewayError as e:
logger.warning(
"Failed to clean up gateway session during restart",
job_name=job_name,
error=str(e),
)
Expand Down
51 changes: 49 additions & 2 deletions orchestrator/tests/test_kubernetes_spawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -889,13 +889,13 @@ def test_restart_limit_exceeded(self, spawner):
)

def test_restart_removes_existing(self, spawner, mock_k8s_client):
"""Restart removes the existing Job before respawning."""
"""Restart deletes the existing Job before respawning."""
spawner.restart_agent_job(
pipeline_id="pipe-1",
agent_role=AgentRole.CODER,
repos=["owner/repo"],
)
mock_k8s_client.remove_container.assert_called()
mock_k8s_client.delete_job.assert_called()

def test_restart_preserves_worktree(self, spawner, mock_k8s_client):
"""Restart calls spawn_agent_job with preserve_worktree_on_failure=True."""
Expand All @@ -907,6 +907,53 @@ def test_restart_preserves_worktree(self, spawner, mock_k8s_client):
)
assert spawner.get_restart_count("pipe-1", "coder") == 1

@pytest.mark.parametrize(
"role",
[
AgentRole.TASK_PLANNER,
AgentRole.RISK_ANALYST,
AgentRole.REVIEWER_CODE,
AgentRole.REVIEWER_CONTRACT,
AgentRole.REVIEWER_AGENT_DESIGN,
AgentRole.REVIEWER_REFINE,
AgentRole.REVIEWER_PLAN,
AgentRole.CONFLICT_RESOLVER,
],
)
def test_restart_underscore_roles_use_hyphenated_k8s_name(
self, spawner, mock_k8s_client, mock_gateway, role
):
"""Restart must convert underscored roles to hyphenated k8s names (#2070).

K8s resource names are RFC-1123 labels and reject underscores, so a
role like ``task_planner`` must become ``task-planner`` in the Job
name. Independently, the call site must pass the prefixed
``egg-sandbox-`` name to ``delete_job`` (the actual k8s name) and
the unprefixed name to ``delete_session_by_container`` (which is
what the gateway session was registered under).
"""
spawner.restart_agent_job(
pipeline_id="issue-1962",
agent_role=role,
repos=["owner/repo"],
)

hyphen_role = role.value.replace("_", "-")
unprefixed = f"egg-agent-issue-1962-{hyphen_role}"
prefixed = f"egg-sandbox-{unprefixed}"

# k8s deletion uses the prefixed Job name.
delete_call = mock_k8s_client.delete_job.call_args_list[0]
assert delete_call.args[0] == prefixed
# No raw underscore must reach the k8s API call.
assert "_" not in delete_call.args[0]

# Gateway session cleanup uses the unprefixed name (matches what
# spawn_agent_job registered with).
gw_call = mock_gateway.delete_session_by_container.call_args_list[0]
assert gw_call.args[0] == unprefixed
assert "_" not in gw_call.args[0]


# ---------------------------------------------------------------------------
# TestRestartCounts
Expand Down
20 changes: 15 additions & 5 deletions orchestrator/tests/test_restart_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
)
from docker_client import ContainerNotFoundError, ContainerOperationError
from gateway_client import GatewayHealth, SessionInfo
from kubernetes_client import JobOperationError
from models import (
AgentExecution,
AgentExecutionStatus,
Expand Down Expand Up @@ -135,8 +136,11 @@ def test_restart_stops_existing_container(
mode="public",
)

# K8s restart calls remove_container (via remove_agent_job)
mock_docker_client.remove_container.assert_called()
# K8s restart calls delete_job directly (#2070): remove_agent_job
# would route both the k8s and gateway calls through one identifier,
# but k8s wants the prefixed form and the gateway session is keyed
# by the unprefixed form.
mock_docker_client.delete_job.assert_called()

def test_restart_spawns_new_container(self, spawner, mock_docker_client, mock_gateway_client):
"""Restart should create a new Job."""
Expand Down Expand Up @@ -192,8 +196,10 @@ def test_restart_custom_max_restarts(self, spawner, mock_docker_client, mock_gat
def test_restart_handles_stop_failure_gracefully(
self, spawner, mock_docker_client, mock_gateway_client
):
"""If removing the old Job fails, restart should still proceed."""
mock_docker_client.remove_container.side_effect = ContainerOperationError("timeout")
"""If deleting the old Job fails with a real k8s error, restart should still proceed."""
# Restart-side delete (line ~1054) fails first; spawn_agent_job's
# subsequent cleanup delete (line ~404) succeeds.
mock_docker_client.delete_job.side_effect = [JobOperationError("api timeout"), None]

result = spawner.restart_agent_container(
pipeline_id="issue-100",
Expand All @@ -202,8 +208,12 @@ def test_restart_handles_stop_failure_gracefully(
mode="public",
)

# Should still succeed — the method handles stop failures gracefully
# Should still succeed — the method swallows JobOperationError on best-effort cleanup
assert isinstance(result, SpawnedContainer)
# And the failing delete_job call really happened against the prefixed k8s name
assert mock_docker_client.delete_job.call_args_list[0].args[0] == (
"egg-sandbox-egg-agent-issue-100-coder"
)

def test_restart_handles_container_not_found(
self, spawner, mock_docker_client, mock_gateway_client
Expand Down
Loading