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]:
@@ -3655,8 +3695,27 @@ async def delete_instance(self, instance_id: InstanceId) -> DeleteInstanceRespon
             instance_id=instance_id,
         )
 
-    async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
-        """Cancel an active command by closing its stream and notifying workers."""
+    async def cancel_local_command(self, command_id: CommandId) -> bool:
+        """Cancel one locally registered command stream, if it exists.
+
+        The single cancellation implementation shared by the public
+        cancel-by-id endpoint and the steward harness's turn cancellation:
+        closes the local response queue immediately (so the caller's stream
+        ends now, not when worker-side cancellation lands — a served engine
+        mid-generation may only observe it at completion) and notifies
+        workers through ``TaskCancelled``.
+
+        Args:
+            command_id: The command whose local stream should be cancelled.
+
+        Returns:
+            True when a local queue existed and was cancelled; False when the
+            command is unknown here or already completed.
+
+        Side effects:
+            Sends ``TaskCancelled``, records the id so local stream cleanup
+            suppresses its ``TaskFinished``, and closes the queue.
+        """
         sender = (
             self._text_generation_queues.get(command_id)
             or self._image_generation_queues.get(command_id)
@@ -3665,21 +3724,35 @@ async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
             or self._audio_transcription_queues.get(command_id)
         )
         if sender is None:
-            raise HTTPException(
-                status_code=404,
-                detail="Command not found or already completed",
-            )
-
+            return False
         await self._send(TaskCancelled(cancelled_command_id=command_id))
         # Suppress the final TaskFinished emitted by local stream cleanup so the
         # worker can observe the Cancelled task and deliver runner-local cancel
         # before event-sourced task deletion happens.
         self._cancelled_command_ids.add(command_id)
         sender.close()
+        return True
 
-        return CancelCommandResponse(
-            message="Command cancelled.",
-            command_id=command_id,
+    async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
+        """Cancel an active command by closing its stream and notifying workers."""
+        if await self.cancel_local_command(command_id):
+            return CancelCommandResponse(
+                message="Command cancelled.",
+                command_id=command_id,
+            )
+        # 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",
         )
 
     def _command_task_is_terminal(self, command_id: CommandId) -> bool:
@@ -4646,15 +4719,27 @@ async def _expire_stale_vision_media(self, now: float) -> None:
                     "Vision media timed out waiting for worker verification",
                 )
 
-    async def send_task_cancellation(self, command_id: CommandId) -> None:
+    async def send_task_cancellation(
+        self, command_id: CommandId, *, suppress_local_finish: bool = True
+    ) -> None:
         """Public cancellation seam for internal callers (the steward harness).
 
         Sends the same TaskCancelled command the HTTP cancel endpoint sends,
-        suppressing the final TaskFinished so the worker can observe the
-        cancelled task and stop the runner promptly.
+        by default suppressing the final TaskFinished so the worker can
+        observe the cancelled task and stop the runner promptly.
+
+        Args:
+            command_id: The command to cancel on the workers.
+            suppress_local_finish: Record the id so local stream cleanup
+                skips its TaskFinished. Pass False when NO local stream will
+                ever open for this command (e.g. a cancellation accepted
+                before its chunk stream was created): the marker is only
+                discarded by stream finalization, so retaining it with no
+                stream would leak one set entry per cancellation.
         """
         await self._send(TaskCancelled(cancelled_command_id=command_id))
-        self._cancelled_command_ids.add(command_id)
+        if suppress_local_finish:
+            self._cancelled_command_ids.add(command_id)
 
     async def dispatch_text_generation(
         self,

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,37 @@ 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:
+            Cancels the active inner command through the API's shared local
+            cancellation path (closing its local queue so this turn's stream
+            ends immediately), falling back to a bare worker notification
+            when the queue is already gone.
+        """
+        self._turn_cancelled = True
+        active = self._active_command_id
+        if active is None:
+            return
+        self._active_command_id = None
+        # Close the inner command's LOCAL queue, not just notify workers: a
+        # served engine mid-generation may observe worker-side cancellation
+        # only at completion, and the accepted cancel must end this response
+        # now, not then. When the queue is already gone there is no stream
+        # left to finalize, so the fallback must not retain the local-finish
+        # marker (nothing would ever discard it).
+        if not await self._api.cancel_local_command(active):
+            await self._api.send_task_cancellation(
+                active, suppress_local_finish=False
+            )
+
     async def _run_investigation(
         self,
         messages: list[ChatCompletionMessage],
@@ -962,6 +1003,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 +1171,19 @@ 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. No chunk stream ever opens for this
+            # command, so the local-finish marker must not be retained.
+            self._active_command_id = None
+            await api.send_task_cancellation(
+                command.command_id, suppress_local_finish=False
+            )
+            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]
@@ -188,3 +189,22 @@ async def test_finalize_command_stream_reports_natural_completion() -> None:
     task_finished = api._send.call_args[0][0]
     assert task_finished.finished_command_id == cid
     assert cid not in queue
+
+
+@pytest.mark.asyncio
+async def test_send_task_cancellation_without_stream_retains_no_marker() -> None:
+    """A no-stream cancellation must not leak a local-finish marker.
+
+    The marker is discarded only by stream finalization; a cancellation for
+    a command whose chunk stream never opens (the steward dispatch race)
+    would otherwise grow `_cancelled_command_ids` forever (#833).
+    """
+    api = _make_api()
+    cid = CommandId("no-stream-cmd")
+
+    await api.send_task_cancellation(cid, suppress_local_finish=False)
+    assert cid not in api._cancelled_command_ids
+    api._send.assert_called_once()
+
+    await api.send_task_cancellation(cid)
+    assert cid in api._cancelled_command_ids

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