Steward hardening: canary truth, cancel-by-id, placement servability

src/skulk/api/main.py

@@ -1567,6 +1567,11 @@ def __init__(
         # this node), so an outstanding failed probe surfaces as `degraded`
         # instead of hiding until the third failure tears the steward down.
         self._steward_canary = StewardCanaryState()
+        # Live steward turns keyed by their advertised outer command id, so
+        # the generic cancel-by-id endpoint can stop a turn whose inner
+        # generation ids the caller never sees. Entries live exactly as long
+        # as their turn's chunk stream.
+        self._steward_turn_harnesses: dict[CommandId, StewardHarness] = {}
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR) if enable_event_log else None
         self._event_log_appends_since_retention_check = 0
         self._system_id = SystemId()
@@ -3371,6 +3376,16 @@ async def _steward_chat_completions(
             chunk_stream = self._extensions.tap_chat_stream(
                 self._extension_context, task_params, chunk_stream
             )
+        # The advertised command id must honor the generic cancel-by-id
+        # contract (POST /v1/cancel/{command_id}). The harness's inner
+        # generation ids are never shown to the caller, so the outer id maps
+        # to the harness for the turn's lifetime and cancel_command routes
+        # it there. Registration happens HERE, before the response exists:
+        # generate_chat_stream advertises the id in its SSE comment before
+        # it pulls the first chunk, so registering lazily inside the stream
+        # would leave a cancel-by-id 404 window during long prefill.
+        self._steward_turn_harnesses[command_id] = harness
+        chunk_stream = self._release_steward_turn_after(command_id, chunk_stream)
         if payload.stream:
             return StreamingResponse(
                 with_sse_keepalive(
@@ -3388,6 +3403,31 @@ async def _steward_chat_completions(
             media_type="application/json",
         )
 
+    async def _release_steward_turn_after(
+        self,
+        command_id: CommandId,
+        chunk_stream: "AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None]",
+    ) -> "AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None]":
+        """Deregister a steward turn from cancel-by-id when its stream ends.
+
+        The caller registers the turn BEFORE constructing the response (the
+        id is advertised to the client before the first chunk is pulled);
+        this wrapper owns only the removal, so normal completion, client
+        disconnect, and cancellation all clean up through one ``finally``.
+
+        Args:
+            command_id: The outer command id the caller registered.
+            chunk_stream: The turn's chunk stream to pass through.
+
+        Yields:
+            The wrapped stream's chunks, unchanged.
+        """
+        try:
+            async for chunk in chunk_stream:
+                yield chunk
+        finally:
+            self._steward_turn_harnesses.pop(command_id, None)
+
     async def _steward_extension_transform(
         self, history: list[StewardChatMessage], *, stream: bool
     ) -> tuple[list[StewardChatMessage], str, TextGenerationTaskParams]:
@@ -3665,6 +3705,16 @@ async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
             or self._audio_transcription_queues.get(command_id)
         )
         if sender is None:
+            # A steward turn advertises one outer command id while its inner
+            # generations run under private ids; route the outer id to the
+            # harness so the generic cancel contract holds there too.
+            steward_turn = self._steward_turn_harnesses.get(command_id)
+            if steward_turn is not None:
+                await steward_turn.cancel_turn()
+                return CancelCommandResponse(
+                    message="Steward turn cancelled.",
+                    command_id=command_id,
+                )
             raise HTTPException(
                 status_code=404,
                 detail="Command not found or already completed",

src/skulk/api/steward.py

@@ -642,6 +642,10 @@ def __init__(self, api: "API") -> None:
         # abandoned stream (client disconnect, cancel button) can stop the
         # runner instead of leaving it generating for nobody.
         self._active_command_id: CommandId | None = None
+        # Latched by cancel_turn: the investigation loop checks it before
+        # every inner dispatch, so cancelling the advertised outer command
+        # stops the whole turn rather than just its current generation.
+        self._turn_cancelled = False
         # Aggregated across the turn's inner generations so the terminal
         # chunk reports real usage instead of null, and the model's actual
         # terminal reason (e.g. length) is preserved.
@@ -713,7 +717,7 @@ async def canary_probe(
         # stream iteration, and the chunk stream's own cancellation handling
         # already sends TaskCancelled and finalizes; a second cancel here
         # would just duplicate it for an already-finalized command.
-        with anyio.move_on_after(CANARY_PROBE_TIMEOUT_SECONDS):
+        with anyio.move_on_after(CANARY_PROBE_TIMEOUT_SECONDS) as deadline_scope:
             async for chunk in chunk_stream:
                 if isinstance(chunk, ErrorChunk):
                     return False
@@ -725,6 +729,12 @@ async def canary_probe(
                     got_text = True
                 if isinstance(chunk, TokenChunk) and chunk.finish_reason is not None:
                     break
+        # A cancelled deadline is a failed probe no matter what arrived
+        # first: partial output followed by a stall is exactly the wedge
+        # this canary exists to catch, and counting it as success would
+        # reset the failure run the teardown threshold depends on.
+        if deadline_scope.cancelled_caught:
+            return False
         return got_text
 
     def steward_instance(self) -> tuple[InstanceId, str] | None:
@@ -952,6 +962,25 @@ async def run_turn_chunks(
                     with anyio.move_on_after(2, shield=True):
                         await self._api.send_task_cancellation(abandoned)
 
+    async def cancel_turn(self) -> None:
+        """Cancel the running turn on behalf of the advertised command id.
+
+        The chat surface advertises one outer command id for the whole turn
+        while the harness dispatches per-step inner generations under ids the
+        client never sees. Cancelling the outer id must therefore stop the
+        inner generation currently in flight AND latch the turn closed so
+        the investigation loop does not simply dispatch its next step.
+
+        Side effects:
+            Sends a task cancellation for the active inner command, if any.
+        """
+        self._turn_cancelled = True
+        active = self._active_command_id
+        if active is None:
+            return
+        self._active_command_id = None
+        await self._api.send_task_cancellation(active)
+
     async def _run_investigation(
         self,
         messages: list[ChatCompletionMessage],
@@ -962,6 +991,8 @@ async def _run_investigation(
         wrap it with abandonment cleanup."""
         reply = ""
         for step_index in range(MAX_STEPS_PER_TURN):
+            if self._turn_cancelled:
+                return
             if step_index == MAX_STEPS_PER_TURN - 1:
                 messages.append(
                     ChatCompletionMessage(
@@ -1128,6 +1159,16 @@ async def _generate_events(
             task_params, target_instance_id=instance_id
         )
         self._active_command_id = command.command_id
+        if self._turn_cancelled:
+            # cancel_turn ran while this dispatch was in flight: the latch
+            # was set before this inner id existed, so cancel_turn had
+            # nothing to stop. Honor it here rather than letting an
+            # accepted cancellation stream one more full generation. The
+            # step ends without a result event, which the loop reads as an
+            # empty final answer.
+            self._active_command_id = None
+            await api.send_task_cancellation(command.command_id)
+            return
         # No extension tap: this is one investigation step, not the turn.
         # The turn's single tap is applied by the caller of
         # run_turn_chunks (API._steward_chat_completions), so observers see

src/skulk/api/tests/test_cancel_command.py

@@ -42,6 +42,7 @@ def _make_api() -> Any:
     api._vision_media_failures = {}  # pyright: ignore[reportPrivateUsage]
     api._vision_media_packet_sender = None  # pyright: ignore[reportPrivateUsage]
     api._cancelled_command_ids = set()  # pyright: ignore[reportPrivateUsage]
+    api._steward_turn_harnesses = {}  # pyright: ignore[reportPrivateUsage]
     api._chunk_reorder = {}  # pyright: ignore[reportPrivateUsage]
     api._data_dedup_cursor = {}  # pyright: ignore[reportPrivateUsage]
     api._data_plane_observer = DataPlaneObserver(  # pyright: ignore[reportPrivateUsage]

src/skulk/api/tests/test_steward_canary.py

@@ -1,9 +1,18 @@
-"""Canary target selection: pure decision logic for the liveness probe."""
+"""Canary behavior: probe target selection and probe outcome semantics."""
 
-from skulk.api.steward import canary_probe_target
+from collections.abc import AsyncGenerator
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, cast
+
+import anyio
+import pytest
+
+import skulk.api.steward as steward_module
+from skulk.api.steward import StewardHarness, canary_probe_target
 from skulk.master.placement import place_instance
 from skulk.master.tests.test_placement import fully_connected_three_nodes
 from skulk.shared.models.model_cards import ModelCard, ModelId, ModelTask
+from skulk.shared.types.chunks import TokenChunk
 from skulk.shared.types.commands import PlaceInstance
 from skulk.shared.types.common import CommandId
 from skulk.shared.types.memory import Memory
@@ -13,6 +22,9 @@
 from skulk.shared.types.worker.runners import RunnerLoading, RunnerReady, RunnerRunning
 from skulk.shared.types.worker.shards import Sharding
 
+if TYPE_CHECKING:
+    from skulk.api.main import API
+
 
 def _steward_placement() -> dict[InstanceId, Instance]:
     topology, node_memory, node_network, _node_ids = fully_connected_three_nodes(
@@ -179,3 +191,63 @@ def test_terminal_lifecycle_tasks_do_not_mute_the_canary() -> None:
     assert (
         canary_probe_target(placed, runners, tasks, host) == instance.instance_id
     )
+
+
+class _WedgedStreamApi:
+    """Stub API whose generation emits text and then stalls forever."""
+
+    def __init__(self, model_id: str) -> None:
+        self._model_id = model_id
+
+    async def running_model_card(self, model_id: ModelId) -> ModelCard:
+        return ModelCard(
+            model_id=ModelId(self._model_id),
+            storage_size=Memory.from_gb(3),
+            n_layers=12,
+            hidden_size=30,
+            supports_tensor=True,
+            tasks=[ModelTask.TextGeneration],
+        )
+
+    async def dispatch_text_generation(
+        self,
+        task_params: TextGenerationTaskParams,
+        target_instance_id: InstanceId | None = None,
+    ) -> object:
+        return SimpleNamespace(command_id=CommandId())
+
+    def text_generation_chunk_stream(
+        self,
+        command: object,
+        task_params: TextGenerationTaskParams,
+        *,
+        extension_tap: bool = True,
+    ) -> "AsyncGenerator[TokenChunk, None]":
+        async def _stream() -> "AsyncGenerator[TokenChunk, None]":
+            yield TokenChunk(
+                model=ModelId(self._model_id),
+                text="O",
+                token_id=-1,
+                usage=None,
+                finish_reason=None,
+            )
+            # The partial-output wedge: text arrived, the terminal never does.
+            await anyio.Event().wait()
+
+        return _stream()
+
+
+async def test_partial_output_then_stall_is_a_failed_probe(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A cancelled deadline must fail the probe even after text arrived.
+
+    Counting the partial output as success would clear the consecutive
+    failure run, so the exact wedge the canary exists to catch could never
+    reach the three-failure teardown threshold (#830).
+    """
+    monkeypatch.setattr(steward_module, "CANARY_PROBE_TIMEOUT_SECONDS", 0.05)
+    api = _WedgedStreamApi("canary-brain")
+    harness = StewardHarness(cast("API", cast(object, api)))
+
+    assert await harness.canary_probe(InstanceId(), "canary-brain") is False

src/skulk/api/tests/test_steward_chunk_stream.py

@@ -206,3 +206,130 @@ async def test_complete_literal_example_survives_in_final_answer() -> None:
     token_chunks = [c for c in chunks if isinstance(c, TokenChunk)]
     content = "".join(c.text for c in token_chunks if not c.is_thinking)
     assert content == answer
+
+
+async def test_cancel_turn_stops_the_generation_and_the_loop() -> None:
+    """Cancelling the advertised outer id ends the whole turn (#830).
+
+    The latch matters as much as the cancellation: without it the
+    investigation loop would treat the cancelled step as complete and
+    dispatch the next one, leaving the turn running behind the caller's
+    back.
+    """
+    cancelled: list[object] = []
+
+    class _Api:
+        async def send_task_cancellation(self, command_id: object) -> None:
+            cancelled.append(command_id)
+
+    # A script that would otherwise tool-loop for the full step budget.
+    harness = _ScriptedHarness(turns=[("", [_call("get_cluster_state")])])
+    harness._api = cast("API", cast(object, _Api()))  # pyright: ignore[reportPrivateUsage]
+    harness._active_command_id = cast(Any, "cmd-inner-1")  # pyright: ignore[reportPrivateUsage]
+
+    stream = harness.run_turn_chunks(
+        [StewardChatMessage(role="user", content="hi")]
+    )
+    first = await stream.__anext__()
+    assert isinstance(first, TokenChunk)
+    generations_before_cancel = len(harness.system_prompts)
+
+    await harness.cancel_turn()
+    assert cancelled == ["cmd-inner-1"]
+
+    async for _chunk in stream:
+        pass
+    # The loop stopped at the latch instead of dispatching further steps.
+    assert len(harness.system_prompts) == generations_before_cancel
+
+
+async def test_cancel_command_routes_a_registered_steward_turn() -> None:
+    """POST /v1/cancel/{outer_id} must reach the turn's harness, not 404."""
+    from skulk.api.main import API
+    from skulk.shared.types.common import CommandId
+
+    cancelled: list[object] = []
+
+    class _Api:
+        async def send_task_cancellation(self, command_id: object) -> None:
+            cancelled.append(command_id)
+
+    harness = _ScriptedHarness(turns=[("irrelevant", [])])
+    harness._api = cast("API", cast(object, _Api()))  # pyright: ignore[reportPrivateUsage]
+    harness._active_command_id = cast(Any, "cmd-inner-9")  # pyright: ignore[reportPrivateUsage]
+
+    outer_id = CommandId()
+    api = API.__new__(API)
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._embedding_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._audio_speech_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._audio_transcription_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._steward_turn_harnesses = {outer_id: harness}  # pyright: ignore[reportPrivateUsage]
+
+    response = await api.cancel_command(outer_id)
+
+    assert response.command_id == outer_id
+    assert cancelled == ["cmd-inner-9"]
+
+
+async def test_cancellation_racing_dispatch_still_cancels_the_step() -> None:
+    """A latch set while dispatch is in flight cancels the fresh command.
+
+    cancel_turn can run between step start and dispatch completion, when no
+    inner id exists yet to cancel. The step must then cancel the command it
+    just obtained and end the turn instead of streaming one more full
+    generation behind an accepted cancellation (#833).
+    """
+    from types import SimpleNamespace
+
+    from skulk.shared.models.model_cards import ModelCard, ModelTask
+    from skulk.shared.types.memory import Memory
+
+    cancelled: list[object] = []
+    stream_requests: list[object] = []
+    harness_holder: list[StewardHarness] = []
+
+    class _RacingApi:
+        async def running_model_card(self, model_id: ModelId) -> ModelCard:
+            return ModelCard(
+                model_id=ModelId("steward-brain"),
+                storage_size=Memory.from_gb(3),
+                n_layers=12,
+                hidden_size=30,
+                supports_tensor=True,
+                tasks=[ModelTask.TextGeneration],
+            )
+
+        async def dispatch_text_generation(
+            self, task_params: object, target_instance_id: object = None
+        ) -> object:
+            # The race: cancellation lands while dispatch is in flight.
+            await harness_holder[0].cancel_turn()
+            return SimpleNamespace(command_id="cmd-fresh-inner")
+
+        def text_generation_chunk_stream(
+            self, command: object, task_params: object, *, extension_tap: bool = True
+        ) -> object:
+            stream_requests.append(command)
+            raise AssertionError("a cancelled step must not open a stream")
+
+        async def send_task_cancellation(self, command_id: object) -> None:
+            cancelled.append(command_id)
+
+    harness = StewardHarness(cast("API", cast(object, _RacingApi())))
+    harness_holder.append(harness)
+    harness.steward_instance = lambda: (InstanceId(), "steward-brain")
+
+    chunks = [
+        chunk
+        async for chunk in harness.run_turn_chunks(
+            [StewardChatMessage(role="user", content="hi")]
+        )
+    ]
+
+    assert cancelled == ["cmd-fresh-inner"]
+    assert stream_requests == []
+    final = chunks[-1]
+    assert isinstance(final, TokenChunk)
+    assert final.finish_reason == "stop"