From bc29425a3fce67ce00a8d614dbbb75abff9ade0c Mon Sep 17 00:00:00 2001 From: Panopticon Agent Date: Wed, 8 Jul 2026 22:23:55 +0000 Subject: [PATCH] =?UTF-8?q?feat(m5.6):=20preferred=5Frunner=5Fid=20?= =?UTF-8?q?=E2=80=94=20pin=20tasks=20to=20a=20specific=20runner=20at=20cre?= =?UTF-8?q?ation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `preferred_runner_id` to `Task`: when set (or `None`), only the matching runner may claim the task. `None` and `"local"` are equivalent — both mean the local runner only. Remote runners are never implicitly handed unaddressed tasks. - `Task` + `TaskRow`: new nullable `preferred_runner_id` field; migration - `service.claim()`: enforces the preference (`effective = pref or "local"`) - `Spawner.spawn_one()` + `spawnable_tasks()`: pre-filter so runners skip tasks reserved for another runner without a wasted claim attempt - `POST /tasks` + `TaskSummaryOut` + `TaskOut`: accept and expose the field - `TaskServiceClient.create_task()`: passes `preferred_runner_id` through - Dashboard `action_new_task`: calls `live_runners()` after memo entry; 0 or 1 runners → auto-select, no selector shown; ≥2 runners → `ChoiceScreen` (local first) for the operator to pick Co-Authored-By: Claude Sonnet 4.6 --- ...dfdbe1_add_preferred_runner_id_to_tasks.py | 36 +++++++++++ src/panopticon/client.py | 3 + src/panopticon/core/models.py | 4 ++ src/panopticon/sessionservice/spawner.py | 22 ++++--- src/panopticon/taskservice/api.py | 4 ++ src/panopticon/taskservice/service.py | 10 +++ .../taskservice/store_sqlalchemy.py | 4 ++ src/panopticon/terminal/dashboard.py | 33 ++++++++-- tests/test_api.py | 9 +-- tests/test_dashboard.py | 4 ++ tests/test_host.py | 2 +- tests/test_runner_liveness.py | 4 +- tests/test_service.py | 50 ++++++++++----- tests/test_skeleton.py | 44 ++++++++++++- tests/test_spawner.py | 63 +++++++++++++------ tests/test_store.py | 1 + 16 files changed, 239 insertions(+), 54 deletions(-) create mode 100644 migrations/versions/20260708_ecf318dfdbe1_add_preferred_runner_id_to_tasks.py diff --git a/migrations/versions/20260708_ecf318dfdbe1_add_preferred_runner_id_to_tasks.py b/migrations/versions/20260708_ecf318dfdbe1_add_preferred_runner_id_to_tasks.py new file mode 100644 index 00000000..0c04012d --- /dev/null +++ b/migrations/versions/20260708_ecf318dfdbe1_add_preferred_runner_id_to_tasks.py @@ -0,0 +1,36 @@ +"""add preferred_runner_id to tasks + +Revision ID: ecf318dfdbe1 +Revises: 596be4c1bf2a +Create Date: 2026-07-08 22:07:23.311486 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'ecf318dfdbe1' +down_revision: str | None = '596be4c1bf2a' +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('task', schema=None) as batch_op: + batch_op.add_column(sa.Column('preferred_runner_id', sa.String(), nullable=True)) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table('task', schema=None) as batch_op: + batch_op.drop_column('preferred_runner_id') + + # ### end Alembic commands ### diff --git a/src/panopticon/client.py b/src/panopticon/client.py index 422c6276..583734bb 100644 --- a/src/panopticon/client.py +++ b/src/panopticon/client.py @@ -148,10 +148,13 @@ def create_task( memo: str | None = None, *, initial_prompt: str | None = None, + preferred_runner_id: str | None = None, ) -> JsonObj: body: JsonObj = {"repo_id": repo_id, "workflow": workflow, "memo": memo} if initial_prompt is not None: body["initial_prompt"] = initial_prompt + if preferred_runner_id is not None: + body["preferred_runner_id"] = preferred_runner_id return cast(JsonObj, self._json(self._http.post("/tasks", json=body))) def set_slug(self, task_id: str, slug: str) -> JsonObj: diff --git a/src/panopticon/core/models.py b/src/panopticon/core/models.py index 63ba94bd..46286330 100644 --- a/src/panopticon/core/models.py +++ b/src/panopticon/core/models.py @@ -258,6 +258,10 @@ class Task: #: or have it respawned. Distinct from liveness — a claimed task whose container died is #: "claimed but down". claimed_by: str | None = None + #: The runner that should claim this task. ``None`` and ``"local"`` are equivalent — only the + #: ``"local"`` runner may claim the task. Any other value restricts claiming to that runner id. + #: Set at creation time (from the dashboard's runner selector); ``None`` until then. + preferred_runner_id: str | None = None #: Cumulative **cost-weighted** tokens the ``claude`` agent in this task's container has used, #: expressed in input-equivalent units (cache-reads ≈0.1×, output ≈5×). The container's Stop #: hook reports it via :meth:`TaskService.set_tokens_used`, recomputing the session total each diff --git a/src/panopticon/sessionservice/spawner.py b/src/panopticon/sessionservice/spawner.py index 7e64646c..a9f73d27 100644 --- a/src/panopticon/sessionservice/spawner.py +++ b/src/panopticon/sessionservice/spawner.py @@ -124,6 +124,9 @@ def spawn_one(self, task: JsonObj) -> str | None: host daemon's per-task isolation still applies but the failure is visible, not silent.""" if task["state"] in TERMINAL_LABELS or task.get("claimed_by"): return None + effective_runner = task.get("preferred_runner_id") or "local" + if effective_runner != self._runner_id: + return None # task is reserved for a different runner try: self._client.claim(task["id"], self._runner_id) except httpx.HTTPStatusError as exc: @@ -325,12 +328,17 @@ def _compose_image(self, workflow: str, repo: JsonObj) -> str | None: return self._images.build(workflow, repo["id"], layers, verbose=True) -def spawnable_tasks(client: TaskServiceClient) -> Callable[[], list[JsonObj]]: - """This host's spawn candidates: unclaimed, non-terminal tasks (the runner claims-then-spawns). +def spawnable_tasks(client: TaskServiceClient, runner_id: str) -> Callable[[], list[JsonObj]]: + """This runner's spawn candidates: unclaimed, non-terminal tasks whose effective preferred + runner matches ``runner_id``. - For M1 (single host) that's every such task the service knows; scoping to this runner's own - assignments is an M5 refinement. + ``preferred_runner_id=None`` is treated as ``"local"`` — only the local runner claims + unaddressed tasks; a remote runner only picks up tasks explicitly assigned to it. """ - return lambda: [ - t for t in client.list_tasks() if not t["claimed_by"] and t["state"] not in TERMINAL_LABELS - ] + def _is_spawnable(t: JsonObj) -> bool: + if t["claimed_by"] or t["state"] in TERMINAL_LABELS: + return False + effective = t.get("preferred_runner_id") or "local" + return effective == runner_id + + return lambda: [t for t in client.list_tasks() if _is_spawnable(t)] diff --git a/src/panopticon/taskservice/api.py b/src/panopticon/taskservice/api.py index 7f52fbf3..f69b5206 100644 --- a/src/panopticon/taskservice/api.py +++ b/src/panopticon/taskservice/api.py @@ -76,6 +76,7 @@ class TaskSummaryOut(BaseModel): branch: str | None clone: str | None claimed_by: str | None + preferred_runner_id: str | None = None tokens_used: int | None token_estimate: int | None governor_task_id: str | None = None @@ -103,6 +104,7 @@ class TaskOut(BaseModel): branch: str | None clone: str | None claimed_by: str | None # the runner that owns this task (the spawn gate), or None + preferred_runner_id: str | None = None # the runner restricted to claiming this task (None = local) tokens_used: int | None # cost-weighted input-equivalent tokens used (None until reported) token_estimate: int | None # the agent's forecast of total tokens (set in planning; None until then) governor_task_id: str | None = None # the task that oversees this one, or None for ungoverned tasks @@ -183,6 +185,7 @@ class CreateTaskIn(BaseModel): initial_prompt: str | None = None artifacts: dict[str, str] | None = None depends_on_task_ids: list[str] = [] + preferred_runner_id: str | None = None class DependenciesIn(BaseModel): @@ -471,6 +474,7 @@ async def create_task(body: CreateTaskIn) -> TaskOut: initial_prompt=body.initial_prompt, artifacts=body.artifacts, depends_on_task_ids=body.depends_on_task_ids or None, + preferred_runner_id=body.preferred_runner_id, ) ) diff --git a/src/panopticon/taskservice/service.py b/src/panopticon/taskservice/service.py index fde8546a..6e03cbc9 100644 --- a/src/panopticon/taskservice/service.py +++ b/src/panopticon/taskservice/service.py @@ -255,6 +255,7 @@ async def create_task( initial_prompt: str | None = None, artifacts: dict[str, str] | None = None, depends_on_task_ids: list[str] | None = None, + preferred_runner_id: str | None = None, ) -> Task: repo = await self.get_repo(repo_id) # ensure exists (raises NotFound) if governor_task_id is not None: @@ -267,6 +268,7 @@ async def create_task( now = self._clock() task = wf.start_task(self._id(), repo_id, at=now, memo=memo, initial_prompt=initial_prompt) task.governor_task_id = governor_task_id + task.preferred_runner_id = preferred_runner_id task.updated_at = now # creation time = first mutation await self._store.create_task(task) _log.info("task %s: created (workflow=%s, repo=%s)", task.id, workflow_name, repo_id) @@ -563,10 +565,18 @@ async def claim(self, task_id: str, runner_id: str) -> Task: Compare-and-set: succeeds if the task is unclaimed (idempotent if this runner already holds it); raises :class:`AlreadyClaimed` if a different runner does. The store is the single writer, so the check-and-set is serialized. + + ``preferred_runner_id`` (``None`` treated as ``"local"``) gates which runner may claim: + only the designated runner is allowed, so a remote runner cannot steal a local-only task. """ task = await self.get_task(task_id) if task.claimed_by not in (None, runner_id): raise AlreadyClaimed(f"task {task_id!r} is already claimed by {task.claimed_by!r}") + effective = task.preferred_runner_id or "local" + if effective != runner_id: + raise AlreadyClaimed( + f"task {task_id!r} is reserved for runner {effective!r}; {runner_id!r} may not claim it" + ) task.claimed_by = runner_id self.clear_lifecycle(task_id) # drop any stale phase from a prior owner; this spawn re-reports await self._save_task(task) diff --git a/src/panopticon/taskservice/store_sqlalchemy.py b/src/panopticon/taskservice/store_sqlalchemy.py index d6b8cd9c..64524866 100644 --- a/src/panopticon/taskservice/store_sqlalchemy.py +++ b/src/panopticon/taskservice/store_sqlalchemy.py @@ -124,6 +124,7 @@ class _TaskRow(_Base): branch: Mapped[str | None] = mapped_column(default=None) clone: Mapped[str | None] = mapped_column(default=None) claimed_by: Mapped[str | None] = mapped_column(default=None) + preferred_runner_id: Mapped[str | None] = mapped_column(default=None) tokens_used: Mapped[int | None] = mapped_column(default=None) token_estimate: Mapped[int | None] = mapped_column(default=None) governor_task_id: Mapped[str | None] = mapped_column(ForeignKey("task.id"), default=None) @@ -151,6 +152,7 @@ def to_domain(self) -> Task: branch=self.branch, clone=self.clone, claimed_by=self.claimed_by, + preferred_runner_id=self.preferred_runner_id, tokens_used=self.tokens_used, token_estimate=self.token_estimate, governor_task_id=self.governor_task_id, @@ -175,6 +177,7 @@ def from_domain(cls, task: Task) -> _TaskRow: branch=task.branch, clone=task.clone, claimed_by=task.claimed_by, + preferred_runner_id=task.preferred_runner_id, tokens_used=task.tokens_used, token_estimate=task.token_estimate, governor_task_id=task.governor_task_id, @@ -381,6 +384,7 @@ async def _update_task(self, task: Task, stored: Sequence[HistoryEntry]) -> None row.branch = task.branch row.clone = task.clone row.claimed_by = task.claimed_by + row.preferred_runner_id = task.preferred_runner_id row.tokens_used = task.tokens_used row.token_estimate = task.token_estimate row.governor_task_id = task.governor_task_id diff --git a/src/panopticon/terminal/dashboard.py b/src/panopticon/terminal/dashboard.py index be20668c..2d2e25cb 100644 --- a/src/panopticon/terminal/dashboard.py +++ b/src/panopticon/terminal/dashboard.py @@ -1365,11 +1365,36 @@ def create(result: "tuple[str, bool] | None") -> None: stripped = memo_text.strip() if _apply_memo_filter(stripped): return - if auto_submit and stripped: - self._client.create_task(repo, workflow, stripped, initial_prompt=stripped) + + def _submit(runner_id: str | None) -> None: + if auto_submit and stripped: + self._client.create_task( + repo, workflow, stripped, + initial_prompt=stripped, preferred_runner_id=runner_id, + ) + else: + self._client.create_task( + repo, workflow, stripped or None, preferred_runner_id=runner_id, + ) + self.action_refresh() + + live_runners = self._client.live_runners() + if len(live_runners) >= 2: + # Sort: "local" first, then alphabetically by id. + sorted_runners = sorted( + live_runners, + key=lambda r: ("" if r["id"] == "local" else r["id"]), + ) + runner_ids = [r["id"] for r in sorted_runners] + + def pick_runner(runner_id: str | None) -> None: + _submit(runner_id) + + self.push_screen(ChoiceScreen("runner", runner_ids), pick_runner) else: - self._client.create_task(repo, workflow, stripped or None) - self.action_refresh() + # 0 or 1 runners: no selector — use the single live runner, or None (local). + preferred = live_runners[0]["id"] if len(live_runners) == 1 else None + _submit(preferred) self.push_screen(MemoScreen(auto_submit_default), create) diff --git a/tests/test_api.py b/tests/test_api.py index 73fddaa6..35d5e0f1 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -259,18 +259,19 @@ def test_set_turn_and_blocked(client: TestClient) -> None: def test_claim_release_over_rest(client: TestClient) -> None: - task_id = _new_task(client) + resp = client.post("/tasks", json={"repo_id": "r1", "workflow": "spike", "preferred_runner_id": "host-1"}) + task_id = resp.json()["id"] assert client.get(f"/tasks/{task_id}").json()["claimed_by"] is None claimed = client.put(f"/tasks/{task_id}/claim", json={"runner_id": "host-1"}) assert claimed.status_code == 200 and claimed.json()["claimed_by"] == "host-1" - # a different runner is refused with 409 while it's held + # a different runner is refused with 409 (already claimed by host-1) assert client.put(f"/tasks/{task_id}/claim", json={"runner_id": "host-2"}).status_code == 409 - # release frees it; another runner can then claim + # release frees it; the preferred runner can reclaim assert client.delete(f"/tasks/{task_id}/claim").json()["claimed_by"] is None - assert client.put(f"/tasks/{task_id}/claim", json={"runner_id": "host-2"}).json()["claimed_by"] == "host-2" + assert client.put(f"/tasks/{task_id}/claim", json={"runner_id": "host-1"}).json()["claimed_by"] == "host-1" # -- artifacts ---------------------------------------------------------------------- diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index bdd282e1..3a0e17ca 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -174,6 +174,9 @@ def list_workflows_for_repo(self, repo_id: str) -> list[dict[str, str]]: def list_operations(self, task_id: str) -> dict[str, str]: return self._operations + def live_runners(self) -> list[dict[str, Any]]: + return [] # 0 live runners by default → no runner selector shown + def create_task( self, repo_id: str, @@ -181,6 +184,7 @@ def create_task( memo: str | None = None, *, initial_prompt: str | None = None, + preferred_runner_id: str | None = None, ) -> dict[str, Any]: self.created.append((repo_id, workflow, memo, initial_prompt)) return {"id": "new"} diff --git a/tests/test_host.py b/tests/test_host.py index 8822df39..3cb76e52 100644 --- a/tests/test_host.py +++ b/tests/test_host.py @@ -310,7 +310,7 @@ def test_run_host_spawns_then_provisions_end_to_end(tmp_path: Path) -> None: asyncio.run(service.create_repo(Repo(id="r1", name="acme/widgets", git_url="https://forge/r1.git", default_base="trunk"))) with TestClient(create_app(service)) as http: client = TaskServiceClient(http) - task_id = client.create_task("r1", "spike")["id"] + task_id = client.create_task("r1", "spike", preferred_runner_id="host-1")["id"] runner = _FakeRunner() def one_pass() -> Callable[[], bool]: diff --git a/tests/test_runner_liveness.py b/tests/test_runner_liveness.py index 2e154448..57a5aa0c 100644 --- a/tests/test_runner_liveness.py +++ b/tests/test_runner_liveness.py @@ -110,8 +110,8 @@ def test_reclaim_releases_only_the_dead_runners_non_terminal_claims( ) -> None: service, base = served client = TaskServiceClient(httpx.Client(base_url=base, timeout=10.0)) - mine = client.create_task("r1", "spike")["id"] - other = client.create_task("r1", "spike")["id"] + mine = client.create_task("r1", "spike", preferred_runner_id="host-dead")["id"] + other = client.create_task("r1", "spike", preferred_runner_id="host-live")["id"] client.claim(mine, "host-dead") client.claim(other, "host-live") diff --git a/tests/test_service.py b/tests/test_service.py index d2b29352..e338e66b 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -247,28 +247,47 @@ async def test_blocked_marker_survives_turn_flips(tmp_path: Path) -> None: async def test_claim_is_compare_and_set_and_idempotent_for_the_holder(tmp_path: Path) -> None: svc = await make_service(tmp_path) - task = await svc.create_task("r1", "spike") + task = await svc.create_task("r1", "spike", preferred_runner_id="host-1") assert task.claimed_by is None assert (await svc.claim(task.id, "host-1")).claimed_by == "host-1" assert (await svc.claim(task.id, "host-1")).claimed_by == "host-1" # idempotent for the same runner assert (await svc.get_task(task.id)).claimed_by == "host-1" # persisted -async def test_claim_rejects_a_different_runner(tmp_path: Path) -> None: +async def test_claim_rejects_wrong_preferred_runner(tmp_path: Path) -> None: + # A runner whose id doesn't match preferred_runner_id cannot claim even when the task is unclaimed. svc = await make_service(tmp_path) - task = await svc.create_task("r1", "spike") + task = await svc.create_task("r1", "spike", preferred_runner_id="host-1") + with pytest.raises(AlreadyClaimed): + await svc.claim(task.id, "host-2") + assert (await svc.get_task(task.id)).claimed_by is None # still unclaimed + + +async def test_claim_rejects_a_different_runner_when_already_claimed(tmp_path: Path) -> None: + # A runner already holding the claim: a second runner is rejected (claimed_by mismatch). + svc = await make_service(tmp_path) + task = await svc.create_task("r1", "spike", preferred_runner_id="host-1") await svc.claim(task.id, "host-1") with pytest.raises(AlreadyClaimed): await svc.claim(task.id, "host-2") assert (await svc.get_task(task.id)).claimed_by == "host-1" # unchanged +async def test_claim_none_preferred_means_local_only(tmp_path: Path) -> None: + # preferred_runner_id=None (the default) is equivalent to "local" — remote runners are rejected. + svc = await make_service(tmp_path) + task = await svc.create_task("r1", "spike") # preferred_runner_id=None → effective "local" + with pytest.raises(AlreadyClaimed): + await svc.claim(task.id, "remote-1") + assert (await svc.claim(task.id, "local")).claimed_by == "local" + + async def test_release_returns_the_task_to_unclaimed(tmp_path: Path) -> None: svc = await make_service(tmp_path) - task = await svc.create_task("r1", "spike") - await svc.claim(task.id, "host-1") + task = await svc.create_task("r1", "spike") # effective preferred = "local" + await svc.claim(task.id, "local") assert (await svc.release(task.id)).claimed_by is None - assert (await svc.claim(task.id, "host-2")).claimed_by == "host-2" # now another runner can claim + assert (await svc.claim(task.id, "local")).claimed_by == "local" # preferred runner can reclaim # -- host (runner) liveness + reclaim: connection-scoped, clock-free (mirror of container liveness) -- @@ -307,8 +326,8 @@ async def test_register_runner_reconnect_overlap_keeps_the_runner_live(tmp_path: async def test_reclaim_releases_the_runners_non_terminal_claims(tmp_path: Path) -> None: svc = await make_service(tmp_path) - mine = await svc.create_task("r1", "spike") - other = await svc.create_task("r1", "spike") + mine = await svc.create_task("r1", "spike", preferred_runner_id="host-dead") + other = await svc.create_task("r1", "spike", preferred_runner_id="host-live") await svc.claim(mine.id, "host-dead") await svc.claim(other.id, "host-live") @@ -323,7 +342,7 @@ async def test_reclaim_releases_the_runners_non_terminal_claims(tmp_path: Path) async def test_reclaim_skips_terminal_tasks(tmp_path: Path) -> None: # A terminal task has nothing to respawn, so reclaim leaves its claim alone (no churn). svc = await make_service(tmp_path) - done = await svc.create_task("r1", "spike") + done = await svc.create_task("r1", "spike", preferred_runner_id="host-dead") await svc.claim(done.id, "host-dead") await svc.apply_operation(done.id, "drop") # -> DROPPED (terminal) @@ -336,7 +355,7 @@ async def test_reclaim_skips_terminal_tasks(tmp_path: Path) -> None: async def test_container_status_folds_phase_registration_and_runner_liveness(tmp_path: Path) -> None: svc = await make_service(tmp_path) - task = await svc.create_task("r1", "spike") + task = await svc.create_task("r1", "spike", preferred_runner_id="host-1") async def status() -> str: return svc.container_status(await svc.get_task(task.id)).value @@ -356,7 +375,7 @@ async def status() -> str: async def test_container_status_is_disconnected_when_the_claiming_runner_is_gone(tmp_path: Path) -> None: svc = await make_service(tmp_path) - task = await svc.create_task("r1", "spike") + task = await svc.create_task("r1", "spike", preferred_runner_id="host-1") await svc.claim(task.id, "host-1") await svc.report_lifecycle(task.id, "host-1", LifecyclePhase.AWAITING) # no runner-liveness connection held → claimed by a runner not connected to the task service @@ -367,16 +386,17 @@ async def test_container_status_is_disconnected_when_the_claiming_runner_is_gone async def test_lifecycle_phase_is_cleared_on_release_and_reclaim(tmp_path: Path) -> None: svc = await make_service(tmp_path) - task = await svc.create_task("r1", "spike") + task = await svc.create_task("r1", "spike", preferred_runner_id="host-1") await svc.claim(task.id, "host-1") await svc.report_lifecycle(task.id, "host-1", LifecyclePhase.AWAITING) assert svc.lifecycle(task.id) is not None await svc.release(task.id) assert svc.lifecycle(task.id) is None # a respawn starts clean - await svc.claim(task.id, "host-2") - await svc.report_lifecycle(task.id, "host-2", LifecyclePhase.BUILDING) - await svc.reclaim("host-2") # a dead runner's phase is stale + # Reclaim also clears the phase (for the same preferred runner reclaiming after a crash). + await svc.claim(task.id, "host-1") + await svc.report_lifecycle(task.id, "host-1", LifecyclePhase.BUILDING) + await svc.reclaim("host-1") # a dead runner's phase is stale assert svc.lifecycle(task.id) is None diff --git a/tests/test_skeleton.py b/tests/test_skeleton.py index 57ab2c28..57059a7e 100644 --- a/tests/test_skeleton.py +++ b/tests/test_skeleton.py @@ -96,7 +96,7 @@ def test_record_provisioning_over_rest(client: TaskServiceClient) -> None: def test_container_lifecycle_endpoints_drive_container_status_over_rest( client: TaskServiceClient, ) -> None: - task = client.create_task("r1", "spike") + task = client.create_task("r1", "spike", preferred_runner_id="host-1") task_id = task["id"] assert task["container_status"] == "queued" # unclaimed, non-terminal assert task["lifecycle_detail"] is None @@ -345,7 +345,7 @@ def test_task_out_runner_host_reflects_claiming_runner_host( ) -> None: # TaskOut.runner_host is derived from claimed_by → runner registration host at query time. svc, client = service_and_client - task_id = client.create_task("r1", "spike")["id"] + task_id = client.create_task("r1", "spike", preferred_runner_id="host-1")["id"] assert client.get_task(task_id)["runner_host"] is None # unclaimed reg = asyncio.run(svc.register_runner("host-1", host="myhost.local")) @@ -357,3 +357,43 @@ def test_task_out_runner_host_reflects_claiming_runner_host( asyncio.run(svc.deregister_runner(reg.id)) assert client.get_task(task_id)["runner_host"] is None # runner gone → host unknown + + +def test_preferred_runner_id_round_trips_and_gates_claim(client: TaskServiceClient) -> None: + # preferred_runner_id is exposed on TaskOut and stored correctly. + task = client.create_task("r1", "spike", preferred_runner_id="remote-1") + assert task["preferred_runner_id"] == "remote-1" + task_id = task["id"] + + # A non-matching runner is rejected. + resp = client._http.put(f"/tasks/{task_id}/claim", json={"runner_id": "local"}) + assert resp.status_code == 409 + + # The matching runner succeeds. + client.claim(task_id, "remote-1") + assert client.get_task(task_id)["claimed_by"] == "remote-1" + + +def test_preferred_runner_id_none_is_local_only(client: TaskServiceClient) -> None: + # preferred_runner_id=None means only the "local" runner may claim (not a remote runner). + task_id = client.create_task("r1", "spike")["id"] + assert client.get_task(task_id)["preferred_runner_id"] is None + + # A remote runner is rejected. + resp = client._http.put(f"/tasks/{task_id}/claim", json={"runner_id": "remote-1"}) + assert resp.status_code == 409 + + # "local" is accepted (None and "local" are equivalent). + client.claim(task_id, "local") + assert client.get_task(task_id)["claimed_by"] == "local" + + +def test_preferred_runner_id_local_explicit_behaves_like_none(client: TaskServiceClient) -> None: + # preferred_runner_id="local" is identical to None — only the "local" runner may claim. + task_id = client.create_task("r1", "spike", preferred_runner_id="local")["id"] + + resp = client._http.put(f"/tasks/{task_id}/claim", json={"runner_id": "remote-1"}) + assert resp.status_code == 409 + + client.claim(task_id, "local") + assert client.get_task(task_id)["claimed_by"] == "local" diff --git a/tests/test_spawner.py b/tests/test_spawner.py index fafbaf96..7a1f3587 100644 --- a/tests/test_spawner.py +++ b/tests/test_spawner.py @@ -112,7 +112,7 @@ def _spawner(client: object, runner: object, images: object = None) -> Spawner: def test_spawn_one_claims_then_spawns_a_fresh_task() -> None: client, runner = _FakeClient(repo=_REPO), _FakeRunner() - cid = _spawner(client, runner).spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "ITERATING", "claimed_by": None}) + cid = _spawner(client, runner).spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "host-1"}) assert cid == "panopticon-t1" assert client.claims == [("t1", "host-1")] # claimed for this host first assert runner.spawned[0]["workspace"] == "/tasks/t1" # per-task clone mounted @@ -125,7 +125,8 @@ def test_spawn_one_passes_initial_prompt_as_env_var() -> None: client, runner = _FakeClient(repo=_REPO), _FakeRunner() _spawner(client, runner).spawn_one( {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", - "claimed_by": None, "memo": "build the thing", "initial_prompt": "review your plan"} + "claimed_by": None, "memo": "build the thing", "initial_prompt": "review your plan", + "preferred_runner_id": "host-1"} ) assert runner.spawned[0]["initial_prompt"] == "review your plan" @@ -134,7 +135,7 @@ def test_spawn_one_passes_turn_for_interrupt_prompt_on_respawn() -> None: client, runner = _FakeClient(repo=_REPO), _FakeRunner() _spawner(client, runner).spawn_one( {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "ITERATING", - "claimed_by": None, "turn": "agent"} + "claimed_by": None, "turn": "agent", "preferred_runner_id": "host-1"} ) assert runner.spawned[0]["turn"] == "agent" @@ -143,7 +144,7 @@ def test_spawn_one_passes_the_docker_in_docker_capability() -> None: repo = {**_REPO, "capabilities": {"docker_in_docker": True}} client, runner = _FakeClient(repo=repo), _FakeRunner() _spawner(client, runner).spawn_one( - {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "ITERATING", "claimed_by": None} + {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "host-1"} ) assert runner.spawned[0]["docker_in_docker"] is True # repo opted in → privileged DinD @@ -172,7 +173,7 @@ def test_spawn_one_composes_the_workflow_image_when_it_has_a_layer() -> None: client, runner, runner_id="host-1", cache=cache, tasks_root="/tasks", # type: ignore[arg-type] git=GitClones(run=_no_op_run), images=images, makedirs=lambda _p: None, # type: ignore[arg-type] ) - spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "github-peer-reviewed", "state": "PLANNING", "claimed_by": None}) + spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "github-peer-reviewed", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"}) assert images.built == [("github-peer-reviewed", "r1", ["RUN apt-get install --yes gh"])] # composed base → layer assert runner.spawned[0]["image"] == "panopticon-github-peer-reviewed-r1" # spawned on the composed image @@ -188,7 +189,7 @@ def test_spawn_one_composes_workflow_then_repo_layers() -> None: client, runner, runner_id="host-1", cache=cache, tasks_root="/tasks", # type: ignore[arg-type] git=GitClones(run=_no_op_run), images=images, makedirs=lambda _p: None, # type: ignore[arg-type] ) - spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "github-peer-reviewed", "state": "PLANNING", "claimed_by": None}) + spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "github-peer-reviewed", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"}) assert images.built == [("github-peer-reviewed", "r1", ["RUN apt-get install --yes gh", "RUN pip install uv"])] assert runner.spawned[0]["image"] == "panopticon-github-peer-reviewed-r1" @@ -197,7 +198,7 @@ def test_spawn_one_probes_base_image_during_building_phase() -> None: client, runner = _FakeClient(repo=_REPO), _FakeRunner() images = _FakeImageBuilder() _spawner(client, runner, images=images).spawn_one( - {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None} + {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"} ) assert images.base_checks == 1 # probed exactly once per spawn @@ -205,7 +206,7 @@ def test_spawn_one_probes_base_image_during_building_phase() -> None: def test_spawn_one_reports_the_phase_sequence() -> None: client, runner = _FakeClient(repo=_REPO), _FakeRunner() _spawner(client, runner).spawn_one( - {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None} + {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"} ) # claiming → preparing → building (in spawn_one), then starting → awaiting (from the runner) assert [p for _, p, _ in client.phases] == ["claiming", "preparing", "building", "starting", "awaiting"] @@ -220,7 +221,7 @@ def spawn(self, *args: object, **kwargs: object) -> str: client = _FakeClient(repo=_REPO) with pytest.raises(RuntimeError): _spawner(client, _BoomRunner()).spawn_one( - {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None} + {"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"} ) last_task, last_phase, last_detail = client.phases[-1] assert (last_task, last_phase) == ("t1", "failed") @@ -440,10 +441,20 @@ def test_spawn_one_skips_terminal_and_already_claimed_tasks() -> None: assert client.claims == [] and runner.spawned == [] +def test_spawn_one_skips_tasks_reserved_for_a_different_runner() -> None: + # preferred_runner_id="remote-1" but this spawner is "host-1" → no claim attempt. + client, runner = _FakeClient(repo=_REPO), _FakeRunner() + result = _spawner(client, runner).spawn_one( + {"id": "t1", "repo_id": "r1", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "remote-1"} + ) + assert result is None + assert client.claims == [] # never tried to claim + + def test_spawn_one_skips_when_another_runner_wins_the_claim() -> None: client, runner = _FakeClient(repo=_REPO), _FakeRunner() client.hold("t1", "host-2") # another runner grabbed it between our snapshot and claim - assert _spawner(client, runner).spawn_one({"id": "t1", "repo_id": "r1", "state": "ITERATING", "claimed_by": None}) is None + assert _spawner(client, runner).spawn_one({"id": "t1", "repo_id": "r1", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "host-1"}) is None assert runner.spawned == [] # 409 → no spawn @@ -451,12 +462,25 @@ def test_spawnable_tasks_filters_unclaimed_non_terminal() -> None: class _Lister: def list_tasks(self) -> list[JsonObj]: return [ - {"id": "a", "state": "ITERATING", "claimed_by": None}, # spawnable - {"id": "b", "state": "ITERATING", "claimed_by": "host-1"}, # already claimed + {"id": "a", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": None}, # spawnable (local) + {"id": "b", "state": "ITERATING", "claimed_by": "local"}, # already claimed {"id": "c", "state": "COMPLETE", "claimed_by": None}, # terminal + {"id": "d", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "remote-1"}, # other runner + ] + + assert [t["id"] for t in spawnable_tasks(_Lister(), "local")()] == ["a"] # type: ignore[arg-type] + + +def test_spawnable_tasks_remote_runner_only_sees_its_tasks() -> None: + class _Lister: + def list_tasks(self) -> list[JsonObj]: + return [ + {"id": "a", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": None}, # local-only + {"id": "b", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "remote-1"}, # this runner + {"id": "c", "state": "ITERATING", "claimed_by": None, "preferred_runner_id": "remote-2"}, # other runner ] - assert [t["id"] for t in spawnable_tasks(_Lister())()] == ["a"] # type: ignore[arg-type] + assert [t["id"] for t in spawnable_tasks(_Lister(), "remote-1")()] == ["b"] # type: ignore[arg-type] def test_spawn_runs_repo_hook_with_correct_args() -> None: @@ -473,7 +497,7 @@ def _fake_hook(hook_file: str, task_id: str, repo_name: str, workspace: str) -> git=GitClones(run=_no_op_run), images=_FakeImageBuilder(), # type: ignore[arg-type] run_hook=_fake_hook, makedirs=lambda _p: None, ) - spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None}) + spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"}) assert calls == [("/hooks/acme.sh", "t1", "acme/widgets", "/tasks/t1")] assert runner.spawned # container still spawned after the hook @@ -491,7 +515,7 @@ def _boom(hook_file: str, task_id: str, repo_name: str, workspace: str) -> None: run_hook=_boom, makedirs=lambda _p: None, ) with pytest.raises(RuntimeError, match="hook exited 1"): - spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None}) + spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"}) assert not runner.spawned # docker run was never called assert any(p == "failed" for _, p, _ in client.phases) # reported as FAILED @@ -506,7 +530,7 @@ def test_spawn_skips_hook_when_repo_has_no_hook_file() -> None: git=GitClones(run=_no_op_run), images=_FakeImageBuilder(), # type: ignore[arg-type] run_hook=lambda *a: calls.append(a), makedirs=lambda _p: None, ) - spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None}) + spawner.spawn_one({"id": "t1", "repo_id": "r1", "workflow": "spike", "state": "PLANNING", "claimed_by": None, "preferred_runner_id": "host-1"}) assert not calls # hook never invoked assert runner.spawned # container spawned normally @@ -517,7 +541,8 @@ def test_spawner_against_the_real_service(tmp_path: Path) -> None: asyncio.run(service.create_repo(Repo(id="r1", name="acme/widgets", git_url="https://forge/r1.git"))) with TestClient(create_app(service)) as http: client = TaskServiceClient(http) - task_id = client.create_task("r1", "spike")["id"] + # Create with preferred_runner_id="host-1" so the spawner (runner_id="host-1") can claim it. + task_id = client.create_task("r1", "spike", preferred_runner_id="host-1")["id"] runner = _FakeRunner() spawner = Spawner( client, runner, runner_id="host-1", # type: ignore[arg-type] @@ -525,7 +550,7 @@ def test_spawner_against_the_real_service(tmp_path: Path) -> None: git=GitClones(run=_no_op_run), images=_FakeImageBuilder(), # type: ignore[arg-type] makedirs=lambda _p: None, ) - (task,) = spawnable_tasks(client)() # the fresh task is spawnable + (task,) = spawnable_tasks(client, "host-1")() # the fresh task is spawnable for host-1 assert spawner.spawn_one(task) == f"panopticon-{task_id}" assert client.get_task(task_id)["claimed_by"] == "host-1" # claim recorded on the service - assert spawnable_tasks(client)() == [] # now claimed → no longer spawnable + assert spawnable_tasks(client, "host-1")() == [] # now claimed → no longer spawnable diff --git a/tests/test_store.py b/tests/test_store.py index 9f356656..a199d9ec 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -467,6 +467,7 @@ def _fully_populated_task() -> Task: branch="panopticon/fix-the-widget", clone="/clones/t-full", claimed_by="local", + preferred_runner_id="local", tokens_used=87500, token_estimate=500_000, governor_task_id="t-governor",