[agentserver-core] Fix steering future leaks and acknowledgment misrouting

.chronus/changes/agentserver-steering-ack-futures.md

@@ -0,0 +1,8 @@
+---
+changeKind: fix
+packages:
+  - sdk/agentserver/azure-ai-agentserver-core
+---
+
+Fixed steering acknowledgment leaks and duplicate turn dispatch when queue or
+turn-state persistence fails.

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_attachments.py

@@ -99,6 +99,29 @@
 _STEERING_QUEUE_CAP = 9
 
 
+def _steering_pending_ack_ids(steering: dict[str, Any], pending_count: int) -> list[str | None]:
+    """Align internal acknowledgment IDs with queued inputs from older records.
+
+    :param steering: Persisted steering state.
+    :type steering: dict[str, Any]
+    :param pending_count: Number of queued inputs.
+    :type pending_count: int
+    :return: IDs aligned with pending inputs, including untagged legacy slots.
+    :rtype: list[str | None]
+    :raises ValueError: If stored IDs cannot align with the pending queue.
+    """
+    if pending_count == 0:
+        return []
+    ids = steering.get("pending_ack_ids")
+    if ids is None:
+        return [None] * pending_count
+    if not isinstance(ids, list) or len(ids) > pending_count or any(
+        value is not None and not isinstance(value, str) for value in ids
+    ):
+        raise ValueError("Invalid steering pending_ack_ids for pending queue")
+    return ids + [None] * (pending_count - len(ids))
+
+
 # --------------------------------------------------------------------------- #
 # Hash helper
 # --------------------------------------------------------------------------- #

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py

@@ -823,6 +823,7 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
         task_id: str,
         input_val: Any,
         existing: Any,
+        ack_id: str,
         input_id: str | None = None,
         if_last_input_id: str | None = None,
     ) -> None:
@@ -837,6 +838,8 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
                 :keyword existing: The previously-fetched task record (used for the
                     first etag attempt; later attempts re-fetch internally).
                 :paramtype existing: Any
+        :keyword ack_id: Internal identifier binding this queue slot to its acknowledgment.
+        :paramtype ack_id: str
         :keyword input_id:  When set, the new input's identity.
                     Used to advance ``payload["last_input_id"]``
                     atomically with the queue append.
@@ -848,13 +851,26 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
         from ._exceptions import (  # pylint: disable=import-outside-toplevel
             SteeringQueueFull,
         )
+        from ._attachments import (  # pylint: disable=import-outside-toplevel
+            _STEERING_INPUT_KEY_PREFIX,
+            _STEERING_THRESHOLD_BYTES,
+            _resolve_input_storage,
+            _steering_pending_ack_ids,
+        )
         from ._models import (  # pylint: disable=import-outside-toplevel
             TaskPatchRequest,
         )
 
         max_retries = 5
         serialized = _serialize_input(input_val)
 
+        def _signal_accepted_append(pending_count: int) -> None:
+            manager._note_lease_refreshed(task_id)  # pylint: disable=protected-access
+            active = manager._active_tasks.get(task_id)  # pylint: disable=protected-access  # noqa: SLF001
+            if active and hasattr(active, "context") and active.context is not None:
+                active._pending_input_count = pending_count  # pylint: disable=protected-access  # noqa: SLF001
+                active.context.cancel.set()
+
         for _attempt in range(max_retries):
             task_info = (
                 existing
@@ -864,9 +880,22 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
             if task_info is None:
                 raise RuntimeError(f"Task {task_id!r} disappeared during steering append")
 
-            #  Re-check the input precondition on each retry to
-            # catch a concurrent steer that may have advanced `last_input_id`
-            # since we last looked.
+            payload = dict(task_info.payload) if task_info.payload else {}
+            steering = dict(payload.get("steering", {}))
+            pending: list[Any] = list(steering.get("pending_inputs", []))
+            pending_ack_ids = _steering_pending_ack_ids(steering, len(pending))
+
+            # The hosted pipeline may retry a PATCH after the service committed
+            # it but the response was lost. Its retry then receives 412 on the
+            # old etag. Re-reading can prove this exact request already landed,
+            # so return success before predecessor/capacity checks can reject it.
+            if ack_id in pending_ack_ids:
+                _signal_accepted_append(len(pending))
+                return
+
+            # Re-check the input precondition on each genuine retry to catch a
+            # concurrent steer that advanced `last_input_id`. A persisted
+            # `ack_id` above proves this request itself already advanced it.
             if _attempt > 0:
                 _check_input_precondition(
                     existing=task_info,
@@ -875,10 +904,6 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
                     if_last_input_id=if_last_input_id,
                 )
 
-            payload = dict(task_info.payload) if task_info.payload else {}
-            steering = dict(payload.get("steering", {}))
-            pending: list[Any] = list(steering.get("pending_inputs", []))
-
             if len(pending) >= _DEFAULT_MAX_PENDING_STEERING:
                 raise SteeringQueueFull(task_id, _DEFAULT_MAX_PENDING_STEERING)
 
@@ -888,12 +913,6 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
             # ``attachments["_steering_input_<seq>"]`` with a ref slot in
             # the queue. The seq counter is monotonic (never reused) so
             # other entries' attachment keys are stable across drains.
-            from ._attachments import (  # pylint: disable=import-outside-toplevel
-                _STEERING_INPUT_KEY_PREFIX,
-                _STEERING_THRESHOLD_BYTES,
-                _resolve_input_storage,
-            )
-
             next_seq = int(steering.get("next_input_seq", 0))
             steering_key = f"{_STEERING_INPUT_KEY_PREFIX}{next_seq}"
             store_mode, queue_entry = _resolve_input_storage(
@@ -908,7 +927,9 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
                 steering["next_input_seq"] = next_seq + 1
 
             pending.append(queue_entry)
+            pending_ack_ids.append(ack_id)
             steering["pending_inputs"] = pending
+            steering["pending_ack_ids"] = pending_ack_ids
             steering["cancel_requested"] = True
             #   SOT: the
             # internal _steering["generation"] payload field is removed
@@ -942,16 +963,12 @@ async def _append_steering_input(  # pylint: disable=protected-access,too-many-l
                         **lease_kwargs,
                     ),
                 )
-                manager._note_lease_refreshed(task_id)  # pylint: disable=protected-access
                 # Signal the running task's cancel event so it can short-circuit.
                 # Spec 031 / FR-001a + SOT §13 ordering invariant: record the
                 # live pending count BEFORE setting cancel, so a handler that
                 # observes ``ctx.cancel.is_set()`` already sees
                 # ``ctx.pending_input_count >= 1``.
-                active = manager._active_tasks.get(task_id)  # pylint: disable=protected-access  # noqa: SLF001
-                if active and hasattr(active, "context") and active.context is not None:
-                    active._pending_input_count = len(pending)  # pylint: disable=protected-access  # noqa: SLF001
-                    active.context.cancel.set()
+                _signal_accepted_append(len(pending))
                 return
             except _HostedConflict as exc:
                 translated = _translate_hosted_conflict(exc, task_id=task_id)
@@ -974,8 +991,8 @@ def _create_steering_ack_run(
         manager: Any,
         task_id: str,
         future: Any,
+        ack_id: str,
         input_id: str | None = None,
-        input_val: Any = None,
     ) -> TaskRun[Output]:
         """Create a TaskRun for a queued steering input.
 
@@ -985,11 +1002,10 @@ def _create_steering_ack_run(
         :type task_id: str
         :param future: Future that will resolve with the next-turn outcome.
         :type future: Any
+        :param ack_id: Internal identifier of the queued steering slot.
+        :type ack_id: str
         :param input_id: The input_id stamped on the queued input (if any).
         :type input_id: str | None
-        :param input_val: The raw queued input value (used to identify the
-            slot when ``cancel()`` is invoked on the returned handle).
-        :type input_val: Any
         :return: A :class:`TaskRun` whose result resolves with the queued turn.
         :rtype: TaskRun[Output]
         """
@@ -998,8 +1014,7 @@ async def _queued_cancel_cb() -> None:
             await manager._cancel_queued_steering_input(  # pylint: disable=protected-access
                 task_id=task_id,
                 future=future,
-                input_id=input_id,
-                input_val=input_val,
+                ack_id=ack_id,
             )
 
         return TaskRun(
@@ -1331,21 +1346,31 @@ async def _lifecycle_start_inner(  # pylint: disable=too-many-locals,too-many-st
             if self._opts.steerable:
                 # Steering path: append input to queue, signal cancel, return ack
                 # pylint: disable=protected-access
-                ack_future = manager._register_steering_future(task_id)
-                await self._append_steering_input(
-                    manager,
-                    task_id=task_id,
-                    input_val=input,
-                    existing=existing,
-                    input_id=input_id,
-                    if_last_input_id=if_last_input_id,
-                )
+                # Keep drain from binding a future before its append is accepted.
+                async with manager._task_write_guard(task_id):
+                    ack_id = _generate_input_id()
+                    ack_future = manager._register_steering_future(task_id, ack_id)
+                    appended = False
+                    try:
+                        await self._append_steering_input(
+                            manager,
+                            task_id=task_id,
+                            input_val=input,
+                            existing=existing,
+                            ack_id=ack_id,
+                            input_id=input_id,
+                            if_last_input_id=if_last_input_id,
+                        )
+                        appended = True
+                    finally:
+                        if not appended:
+                            manager._unregister_steering_future(task_id, ack_id, ack_future)
                 # Set cancel on in-memory context if task runs in this process
                 active = manager._active_tasks.get(task_id)
                 # pylint: enable=protected-access
                 if active:
                     active.context.cancel.set()
-                return self._create_steering_ack_run(manager, task_id, ack_future, input_id=input_id, input_val=input)
+                return self._create_steering_ack_run(manager, task_id, ack_future, ack_id, input_id=input_id)
             raise TaskConflictError(task_id, "in_progress")
 
         # completed (or any other terminal status)
@@ -1771,8 +1796,8 @@ async def delete(self, task_id: str) -> None:
                 exec_task.cancel()
 
         # 2. Resolve all queued steerer futures with TaskCancelled.
-        pending = getattr(mgr, "_pending_steering_futures", {}).pop(task_id, [])
-        for queued_fut in pending:
+        pending = getattr(mgr, "_pending_steering_futures", {}).pop(task_id, {})
+        for queued_fut in pending.values():
             if not queued_fut.done():
                 queued_fut.set_exception(TaskCancelled())