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,10 +3376,27 @@ 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.
+        # Deregistration wraps the OUTERMOST response iterator: an inner
+        # layer (the keepalive wrapper emits bytes before pulling its
+        # source) can be abandoned before ever starting, and an unstarted
+        # generator's finally never runs — only the iterator Starlette
+        # itself drives is guaranteed to start and therefore to clean up.
+        self._steward_turn_harnesses[command_id] = harness
         if payload.stream:
             return StreamingResponse(
-                with_sse_keepalive(
-                    generate_chat_stream(command_id, chunk_stream),
+                self._release_steward_turn_after(
+                    command_id,
+                    with_sse_keepalive(
+                        generate_chat_stream(command_id, chunk_stream),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -3384,10 +3406,42 @@ async def _steward_chat_completions(
                 },
             )
         return StreamingResponse(
-            collect_chat_response(command_id, chunk_stream),
+            self._release_steward_turn_after(
+                command_id,
+                collect_chat_response(command_id, chunk_stream),
+            ),
             media_type="application/json",
         )
 
+    async def _release_steward_turn_after(
+        self,
+        command_id: CommandId,
+        response_stream: "AsyncIterator[str]",
+    ) -> "AsyncGenerator[str, None]":
+        """Deregister a steward turn from cancel-by-id when its response 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. It must wrap the OUTERMOST
+        response iterator — the one Starlette drives — because that is the
+        only generator guaranteed to be started (and thus finalized) even
+        when the client disconnects after the first keepalive byte; an
+        abandoned inner generator that never started never runs its
+        ``finally``.
+
+        Args:
+            command_id: The outer command id the caller registered.
+            response_stream: The fully assembled response body iterator.
+
+        Yields:
+            The wrapped response's items, unchanged.
+        """
+        try:
+            async for item in response_stream:
+                yield item
+        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 +3709,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 +3738,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 +4733,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