diff --git a/.chronus/changes/agentserver-steering-ack-futures.md b/.chronus/changes/agentserver-steering-ack-futures.md new file mode 100644 index 000000000000..31c789e18d36 --- /dev/null +++ b/.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. diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_attachments.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_attachments.py index 8aef73bc7647..c268721e26a1 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_attachments.py +++ b/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 # --------------------------------------------------------------------------- # diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py index ca16ea56a2d6..79afb27bcd3d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py +++ b/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,6 +851,12 @@ 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, ) @@ -855,6 +864,13 @@ async def _append_steering_input( # pylint: disable=protected-access,too-many-l 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_"]`` 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()) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py index bfb13be53708..dd6ac7bca4af 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py @@ -14,8 +14,9 @@ import logging import os import traceback -from collections.abc import Awaitable, Callable, Iterator, Mapping -from contextlib import contextmanager +import uuid +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping +from contextlib import asynccontextmanager, contextmanager from contextvars import ContextVar from typing import Any, TypeVar @@ -35,6 +36,7 @@ _read_input_value, _ref_key, _resolve_input_storage, + _steering_pending_ack_ids, ) from ._decorator import TaskOptions, _deserialize_input, _resolve_effective_timeout, _serialize_input from ._exceptions import ( @@ -188,7 +190,7 @@ def _parse_turn_started_at(value: Any) -> float | None: def _resolve_queued_steerers_on_terminal( - pending_steering_futures: dict[str, list["asyncio.Future[Any]"]], + pending_steering_futures: dict[str, dict[str, "asyncio.Future[Any]"]], task_id: str, *, current_status: str, @@ -205,9 +207,9 @@ def _resolve_queued_steerers_on_terminal( Pops every queued steerer future for ``task_id`` and resolves each with ``TaskConflictError(current_status=current_status)``. - :param pending_steering_futures: Per-task list of pending steerer + :param pending_steering_futures: Per-task map of pending steerer futures (mutated in-place — emptied for the given ``task_id``). - :type pending_steering_futures: dict[str, list[asyncio.Future[Any]]] + :type pending_steering_futures: dict[str, dict[str, asyncio.Future[Any]]] :param task_id: The task whose queued steerers should be resolved. :type task_id: str :keyword current_status: Status string to carry on @@ -216,8 +218,8 @@ def _resolve_queued_steerers_on_terminal( """ # TaskConflictError is already imported at module top-level (line 24). - queued = pending_steering_futures.pop(task_id, []) - for fut in queued: + queued = pending_steering_futures.pop(task_id, {}) + for fut in queued.values(): if not fut.done(): fut.set_exception(TaskConflictError(task_id, current_status)) @@ -470,19 +472,22 @@ def __init__( self._shutdown_event = shutdown_event or asyncio.Event() self._shutdown_grace_seconds = shutdown_grace_seconds self._active_generation_future: dict[str, asyncio.Future[Any]] = {} - self._pending_steering_futures: dict[str, list[asyncio.Future[Any]]] = {} + self._pending_steering_futures: dict[str, dict[str, asyncio.Future[Any]]] = {} # Layer 2: periodic recovery scan task. Created # at startup() time; cancelled at shutdown(). self._periodic_recovery_task: asyncio.Task[None] | None = None # / C-WQ-1..3 — per-task write-queue # registry. A single asyncio.Lock per task_id serializes all # in-process PATCHes against that task so etag conflicts become - # rare (only cross-process). Lazy-created on first use; dropped - # in ``_active_tasks_pop`` (no leaks). + # rare (only cross-process). Lazy-created on first use; guards + # track holders + waiters so remote-task entries can be reclaimed + # safely when the final user exits. # — also tracks the latest known etag # per task_id outside the _ActiveTask entry, so reclaim/scan # paths (which have no _ActiveTask yet) can still benefit. self._task_write_locks: dict[str, asyncio.Lock] = {} + self._task_write_lock_users: dict[str, int] = {} + self._task_write_cleanup_pending: set[str] = set() self._task_etag_cache: dict[str, str] = {} # SOT §52 — per-turn timeout watchdog registry. Each per-turn # watchdog gets registered here so that the steering-drain @@ -675,7 +680,7 @@ async def list_tasks( raise RuntimeError("Task list did not converge after retryable conflict") from exc raise translated from exc - def _register_steering_future(self, task_id: str) -> asyncio.Future[Any]: + def _register_steering_future(self, task_id: str, ack_id: str) -> asyncio.Future[Any]: """Create and register a future for a queued steering input. Must be called BEFORE ``_append_steering_input()`` to avoid a race @@ -683,23 +688,54 @@ def _register_steering_future(self, task_id: str) -> asyncio.Future[Any]: :param task_id: The task identifier. :type task_id: str + :param ack_id: Unique internal identifier of the queue slot. + :type ack_id: str :return: The registered future. :rtype: asyncio.Future[Any] """ + pending = self._pending_steering_futures.setdefault(task_id, {}) + if ack_id in pending: + raise ValueError(f"Steering acknowledgment ID already registered for task {task_id!r}") loop = asyncio.get_running_loop() future: asyncio.Future[Any] = loop.create_future() - if task_id not in self._pending_steering_futures: - self._pending_steering_futures[task_id] = [] - self._pending_steering_futures[task_id].append(future) + pending[ack_id] = future return future - async def _cancel_queued_steering_input( # pylint: disable=unused-argument + def _remove_steering_future(self, task_id: str, ack_id: str, future: asyncio.Future[Any]) -> None: + """Remove a registered future only if this request owns its ID. + + :param task_id: The task identifier. + :type task_id: str + :param ack_id: Internal identifier of the queued slot. + :type ack_id: str + :param future: The request's registered future. + :type future: asyncio.Future[Any] + """ + pending = self._pending_steering_futures.get(task_id) + if pending is not None and pending.get(ack_id) is future: + del pending[ack_id] + if not pending: + self._pending_steering_futures.pop(task_id, None) + + def _unregister_steering_future(self, task_id: str, ack_id: str, future: asyncio.Future[Any]) -> None: + """Discard only the future owned by a failed steering append. + + :param task_id: The task identifier. + :type task_id: str + :param ack_id: Internal identifier of the attempted queue slot. + :type ack_id: str + :param future: The failed append's registered future. + :type future: asyncio.Future[Any] + """ + self._remove_steering_future(task_id, ack_id, future) + future.cancel() + + async def _cancel_queued_steering_input( self, *, task_id: str, future: asyncio.Future[Any], - input_id: str | None, - input_val: Any, + ack_id: str, ) -> None: """Remove a queued steering input from the chain's pending queue. @@ -712,53 +748,41 @@ async def _cancel_queued_steering_input( # pylint: disable=unused-argument :keyword task_id: The chain task identifier. :keyword future: The queued steerer's result_future. - :keyword input_id: The input_id of the queued slot (used for the - future-list cleanup; the queue entry itself is identified by - ``input_val``). - :keyword input_val: The raw queued value used to identify which - ``pending_inputs`` entry to remove. + :keyword ack_id: Internal ID identifying the exact queued slot, even + when multiple inputs have identical values or public input IDs. """ from ._attachments import _is_ref, _ref_key # pylint: disable=import-outside-toplevel from ._exceptions import TaskCancelled # pylint: disable=import-outside-toplevel - async with self._get_task_write_lock(task_id): + async with self._task_write_guard(task_id): try: task_info = await self._provider_get_tracked(task_id) except Exception: # pylint: disable=broad-exception-caught task_info = None if task_info is None or not task_info.payload: # Chain already gone — just resolve the future. + self._remove_steering_future(task_id, ack_id, future) if not future.done(): future.set_exception(TaskCancelled()) return steering = dict(task_info.payload.get("steering") or {}) pending = list(steering.get("pending_inputs") or []) + pending_ack_ids = _steering_pending_ack_ids(steering, len(pending)) attachments_patch: dict[str, Any] = {} - # Drop the first queue entry whose raw value matches ``input_val``. - removed = False - new_pending: list[Any] = [] - for entry in pending: - if not removed: - raw = entry - if _is_ref(entry): - # For ref-shaped entries, resolve via attachment to - # compare against input_val. If the attachment is - # missing, fall back to ref identity (unlikely). - key = _ref_key(entry) - raw = (task_info.attachments or {}).get(key, entry) - if raw == input_val: - removed = True - if _is_ref(entry): - attachments_patch[_ref_key(entry)] = None - continue - new_pending.append(entry) - if not removed: + if ack_id not in pending_ack_ids: # Queue entry already drained or never landed; just resolve. + self._remove_steering_future(task_id, ack_id, future) if not future.done(): future.set_exception(TaskCancelled()) return - steering["pending_inputs"] = new_pending - steering["cancel_requested"] = len(new_pending) > 0 + index = pending_ack_ids.index(ack_id) + entry = pending.pop(index) + pending_ack_ids.pop(index) + if _is_ref(entry): + attachments_patch[_ref_key(entry)] = None + steering["pending_inputs"] = pending + steering["pending_ack_ids"] = pending_ack_ids + steering["cancel_requested"] = len(pending) > 0 payload_patch: dict[str, Any] = {"steering": steering} try: # Spec 031 / FR-005a+b: the outer lock is already held, so use @@ -781,10 +805,8 @@ async def _cancel_queued_steering_input( # pylint: disable=unused-argument task_id, exc_info=True, ) - # Remove the future from the registered pending list and resolve it. - pending_list = self._pending_steering_futures.get(task_id) or [] - if future in pending_list: - pending_list.remove(future) + # Remove this request's registration and resolve its returned handle. + self._remove_steering_future(task_id, ack_id, future) if not future.done(): future.set_exception(TaskCancelled()) @@ -2055,6 +2077,7 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br # the framework's ``return X`` is the only end-of-turn # signal. Success flow. renewal_cancel.set() + is_multi_turn_success = getattr(opts, "_is_multi_turn", False) try: completed = await self._handle_success( task_id=task_id, @@ -2084,13 +2107,25 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br current_status="failed", ) break + except Exception as persistence_exc: # noqa: BLE001 + # The handler returned successfully, but its durable + # terminal transition did not persist. Do not acknowledge + # the result or dispatch another turn from uncertain state. + if not current_result_future.done(): + current_result_future.set_exception(persistence_exc) + if is_multi_turn_success: + _resolve_queued_steerers_on_terminal( + self._pending_steering_futures, + task_id, + current_status="in_progress", + ) + break # Set the current turn's caller's result_future to the # completion outcome FIRST, then resolve any queued # steerers with TaskConflictError (since the task has now # terminated). The handler's return value is delivered # unchanged to the current caller; the queued steerers # see the "task is busy / terminal" shape per Invariant 1. - is_multi_turn_success = getattr(opts, "_is_multi_turn", False) if not current_result_future.done(): # Both one-shot and multi-turn return the raw Output # unwrapped; multi-turn keeps the chain alive. @@ -2150,11 +2185,10 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br TaskCancelled, ) - if not current_result_future.done(): - current_result_future.set_exception(TaskCancelled()) - is_multi_turn_cancel = getattr(opts, "_is_multi_turn", False) if opts.ephemeral: + if not current_result_future.done(): + current_result_future.set_exception(TaskCancelled()) # One-shot is always ephemeral: delete the persisted # record so the recovery scanner doesn't re-invoke a # cancelled handler. @@ -2172,9 +2206,9 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br # (input + _retry_attempt + _steering.active_input + any # promoted _input attachment cleared atomically — see # SOT §23.8 item #3). The chain stays alive for the next - # turn's ``.start`` / ``.run``. Errors from this write - # surface only via the logger because the caller's - # future is already resolved with TaskCancelled above. + # turn's ``.start`` / ``.run``. The write must succeed + # before TaskCancelled is acknowledged or a queued turn + # is promoted. error_dict: dict[str, Any] = { "type": "cancelled", "message": "Task cancelled", @@ -2186,20 +2220,32 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br opts=opts, error_dict=error_dict, ) - except TaskConflictError: - # 412 RE-READ decided ABANDON; lease is no longer - # ours, another instance / process owns the record. - # Nothing to clean up here — the caller already saw - # TaskCancelled. - pass - except Exception: # pylint: disable=broad-exception-caught + except TaskConflictError as persistence_exc: + if not current_result_future.done(): + current_result_future.set_exception(persistence_exc) + _resolve_queued_steerers_on_terminal( + self._pending_steering_futures, + task_id, + current_status=persistence_exc.current_status, + ) + break + except Exception as persistence_exc: # pylint: disable=broad-exception-caught logger.warning( "Failed to transition multi-turn chain %s " - "to suspended after cancel; chain may need " - "recovery scan to pick up the in_progress record", + "to suspended after cancel; deferring recovery", task_id, exc_info=True, ) + if not current_result_future.done(): + current_result_future.set_exception(persistence_exc) + _resolve_queued_steerers_on_terminal( + self._pending_steering_futures, + task_id, + current_status="in_progress", + ) + break + if not current_result_future.done(): + current_result_future.set_exception(TaskCancelled()) # Promote queued steerers (if any) per the same drain # rule as raise — chain stays alive, queued head takes # over the next turn. @@ -2281,14 +2327,35 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br "traceback": traceback.format_exc(), } - await self._handle_failure( - task_id=task_id, - exc=exc, - opts=opts, - ) + is_multi_turn_failure = getattr(opts, "_is_multi_turn", False) + try: + await self._handle_failure( + task_id=task_id, + exc=exc, + opts=opts, + ) + except TaskConflictError as persistence_exc: + if not current_result_future.done(): + current_result_future.set_exception(persistence_exc) + if is_multi_turn_failure: + _resolve_queued_steerers_on_terminal( + self._pending_steering_futures, + task_id, + current_status=persistence_exc.current_status, + ) + break + except Exception as persistence_exc: # noqa: BLE001 + if not current_result_future.done(): + current_result_future.set_exception(persistence_exc) + if is_multi_turn_failure: + _resolve_queued_steerers_on_terminal( + self._pending_steering_futures, + task_id, + current_status="in_progress", + ) + break # / step 5 — caller's future resolution: # CancelledError → bare TaskCancelled() else TaskFailed. - is_multi_turn_failure = getattr(opts, "_is_multi_turn", False) if not current_result_future.done(): from ._exceptions import TaskCancelled # pylint: disable=import-outside-toplevel @@ -2370,6 +2437,7 @@ async def _try_drain_steering( # pylint: disable=too-many-branches,too-many-sta result_future: asyncio.Future[Any], partial_output: Any | None = None, _conflict_attempt: int = 0, + _expected_active_ack_id: str | None = None, ) -> TaskContext[Any] | None: """Check for pending steering inputs and drain the next one. @@ -2393,6 +2461,9 @@ async def _try_drain_steering( # pylint: disable=too-many-branches,too-many-sta for etag-conflict retries. Bounded so the hosted task store's etag-comparator pre-fix behaviour cannot loop forever. + :keyword _expected_active_ack_id: Internal acknowledgment ID of the + drain transition whose PATCH may have committed before returning + a conflict. Only that recursive retry may reconcile active state. :return: New context for the drained generation, or None. """ # Spec 031 / FR-005 + SOT §25.2 — the read-state + compute-PATCH + @@ -2406,7 +2477,7 @@ async def _try_drain_steering( # pylint: disable=too-many-branches,too-many-sta # cross-process conflicts retry OUTSIDE the lock via the recursion # below (the per-task ``asyncio.Lock`` is non-reentrant). drain_conflict: BaseException | None = None - async with self._get_task_write_lock(task_id): + async with self._task_write_guard(task_id): task_info = await self._provider_get_tracked(task_id) if task_info is None: return None @@ -2414,89 +2485,126 @@ async def _try_drain_steering( # pylint: disable=too-many-branches,too-many-sta payload = dict(task_info.payload) if task_info.payload else {} steering = dict(payload.get("steering", {})) pending = list(steering.get("pending_inputs", [])) - - if not pending: - return None - - # Pop the next input from the queue.: the entry may be - # either a raw inline value (≤ 20 KiB at append) or a ref slot - # pointing into ``task_info.attachments``. Resolve uniformly via - # ``_read_input_value``; if it was a ref, the same drain PATCH - # MUST also delete the attachment (C-9 /). - next_entry = pending.pop(0) - attachments_patch: dict[str, Any] = {} - if _is_ref(next_entry): - attachments_patch[_ref_key(next_entry)] = None - next_input_raw = _read_input_value(next_entry, task_info.attachments) - - # Admission and recovery must see the promoted turn as the current input. - input_mode, input_value = _resolve_input_storage( - next_input_raw, - threshold_bytes=_INPUT_THRESHOLD_BYTES, - key_for_attachment=_FUNCTION_INPUT_KEY, - task_id=task_id, + active_ack_id = steering.get("active_ack_id") + if active_ack_id is not None and not isinstance(active_ack_id, str): + raise ValueError("Invalid steering active_ack_id") + + committed_expected_drain = ( + _expected_active_ack_id is not None + and steering.get("drain_in_progress") is True + and active_ack_id == _expected_active_ack_id ) - payload["input"] = input_value - if input_mode == "attachment": - attachments_patch[_FUNCTION_INPUT_KEY] = next_input_raw - - # Update steering state. (: previous_input is - # no longer mirrored into _steering; only the active input + queue - # state need to survive a crash mid-drain.) - steering["active_input"] = next_input_raw - steering["pending_inputs"] = pending - # SOT: internal - # _steering["generation"] writes removed. The drain transition - # IS the generation advance — no separate counter needed. - steering["cancel_requested"] = len(pending) > 0 - steering["drain_in_progress"] = True - #: the steering drain re-entry is a NEW - # turn-start boundary — write a fresh _turn_started_at so the - # respawned watchdog computes a full per-turn budget. - payload[_TURN_STARTED_AT_KEY] = _utc_now_iso() - payload["steering"] = steering - # SOT §11/§20: the framework does not write payload["output"]; - # no clear is needed at the drain transition. - - try: - etag = getattr(task_info, "etag", None) or None - # Spec 031 (hosted re-test finding) — the multi-turn turn that - # just ended already wrote ``status="suspended"`` (see - # ``_handle_multi_turn_success``). The drain starts a NEW turn, - # so it MUST transition the record back to ``in_progress`` in - # this same PATCH. This is also REQUIRED for the lease-extension - # piggyback to be valid: the hosted task store rejects lease - # *renewal* on a non-in_progress task ("lease renewal is only - # supported for in_progress tasks"), but ACCEPTS lease params as - # part of a suspended→in_progress *claim*. Without the status - # flip the drain PATCH 409s and the steered turn never runs. - await self._provider_update_lock_held( - task_id, - TaskPatchRequest( - status="in_progress", - payload=payload, - attachments=attachments_patch, - if_match=etag, - **self._lease_ext_kwargs(task_id), - ), - ) - # Spec 031 / FR-002 — the drain consumed the head; the steered - # turn's live backlog is the remaining ``pending``. Keep - # ``ctx.pending_input_count`` in sync for the new turn. + if committed_expected_drain: + # This recursive retry is the only caller allowed to reconcile + # an active drain. The Phase-1 PATCH already removed the queue + # head and persisted its raw active input. + next_input_raw = steering.get("active_input") + if next_input_raw is None: + raise RuntimeError(f"Task {task_id!r} committed a drain without an active input") + next_ack_id = _expected_active_ack_id + _steering_pending_ack_ids(steering, len(pending)) active_now = self._active_tasks.get(task_id) if active_now is not None: active_now._pending_input_count = len(pending) - except _HostedConflict as exc: - translated = _translate_hosted_conflict(exc, task_id=task_id) - if translated is not None: - raise translated from exc - drain_conflict = exc - except (EtagConflict, ValueError, TransportClassifiedError) as exc: - if isinstance(exc, TransportClassifiedError) and getattr(exc, "classification", None) != "conflict": - raise - if isinstance(exc, ValueError) and "etag" not in str(exc).lower(): - raise - drain_conflict = exc + else: + if steering.get("drain_in_progress") is True: + # The active turn belongs to an earlier transition, not this + # retry. Dispatching it again would duplicate side effects. + logger.warning( + "Task %s has an active steering drain outside the current retry; " + "leaving it for terminal persistence or recovery", + task_id, + ) + return None + + if not pending: + return None + + pending_ack_ids = _steering_pending_ack_ids(steering, len(pending)) + next_ack_id = pending_ack_ids.pop(0) + if next_ack_id is None: + next_ack_id = f"drain-{uuid.uuid4().hex}" + # Pop the next input from the queue.: the entry may be + # either a raw inline value (≤ 20 KiB at append) or a ref slot + # pointing into ``task_info.attachments``. Resolve uniformly via + # ``_read_input_value``; if it was a ref, the same drain PATCH + # MUST also delete the attachment (C-9 /). + next_entry = pending.pop(0) + attachments_patch: dict[str, Any] = {} + if _is_ref(next_entry): + attachments_patch[_ref_key(next_entry)] = None + next_input_raw = _read_input_value(next_entry, task_info.attachments) + + # Admission and recovery must see the promoted turn as the current input. + input_mode, input_value = _resolve_input_storage( + next_input_raw, + threshold_bytes=_INPUT_THRESHOLD_BYTES, + key_for_attachment=_FUNCTION_INPUT_KEY, + task_id=task_id, + ) + payload["input"] = input_value + if input_mode == "attachment": + attachments_patch[_FUNCTION_INPUT_KEY] = next_input_raw + + # Update steering state. (: previous_input is + # no longer mirrored into _steering; only the active input + queue + # state need to survive a crash mid-drain.) + steering["active_input"] = next_input_raw + steering["active_ack_id"] = next_ack_id + steering["pending_inputs"] = pending + steering["pending_ack_ids"] = pending_ack_ids + # SOT: internal + # _steering["generation"] writes removed. The drain transition + # IS the generation advance — no separate counter needed. + steering["cancel_requested"] = len(pending) > 0 + steering["drain_in_progress"] = True + #: the steering drain re-entry is a NEW + # turn-start boundary — write a fresh _turn_started_at so the + # respawned watchdog computes a full per-turn budget. + payload[_TURN_STARTED_AT_KEY] = _utc_now_iso() + payload["steering"] = steering + # SOT §11/§20: the framework does not write payload["output"]; + # no clear is needed at the drain transition. + + try: + etag = getattr(task_info, "etag", None) or None + # Spec 031 (hosted re-test finding) — the multi-turn turn that + # just ended already wrote ``status="suspended"`` (see + # ``_handle_multi_turn_success``). The drain starts a NEW turn, + # so it MUST transition the record back to ``in_progress`` in + # this same PATCH. This is also REQUIRED for the lease-extension + # piggyback to be valid: the hosted task store rejects lease + # *renewal* on a non-in_progress task ("lease renewal is only + # supported for in_progress tasks"), but ACCEPTS lease params as + # part of a suspended→in_progress *claim*. Without the status + # flip the drain PATCH 409s and the steered turn never runs. + await self._provider_update_lock_held( + task_id, + TaskPatchRequest( + status="in_progress", + payload=payload, + attachments=attachments_patch, + if_match=etag, + **self._lease_ext_kwargs(task_id), + ), + ) + # Spec 031 / FR-002 — the drain consumed the head; the steered + # turn's live backlog is the remaining ``pending``. Keep + # ``ctx.pending_input_count`` in sync for the new turn. + active_now = self._active_tasks.get(task_id) + if active_now is not None: + active_now._pending_input_count = len(pending) + except _HostedConflict as exc: + translated = _translate_hosted_conflict(exc, task_id=task_id) + if translated is not None: + raise translated from exc + drain_conflict = exc + except (EtagConflict, ValueError, TransportClassifiedError) as exc: + if isinstance(exc, TransportClassifiedError) and getattr(exc, "classification", None) != "conflict": + raise + if isinstance(exc, ValueError) and "etag" not in str(exc).lower(): + raise + drain_conflict = exc # Lock released — a genuine (cross-process) conflict retries here, # re-reading the NEW state under a fresh lock acquisition. Bounded so @@ -2518,13 +2626,18 @@ async def _try_drain_steering( # pylint: disable=too-many-branches,too-many-sta result_future=result_future, partial_output=partial_output, _conflict_attempt=_conflict_attempt + 1, + _expected_active_ack_id=next_ack_id, ) - # Pop and bind the next pending steering future (if any) + # Bind only the future registered for this durable queue slot. new_future: asyncio.Future[Any] | None = None - steering_futures = self._pending_steering_futures.get(task_id, []) - if steering_futures: - new_future = steering_futures.pop(0) + steering_futures = self._pending_steering_futures.get(task_id) + if next_ack_id is not None and steering_futures is not None: + new_future = steering_futures.pop(next_ack_id, None) + if not steering_futures: + self._pending_steering_futures.pop(task_id, None) + if next_ack_id is not None and new_future is None: + logger.debug("Steering input %s for task %s has no local acknowledgment future", next_ack_id, task_id) # Resolve the queued steerer's future binding for the new turn. # / (Subscriber): the OLD result_future is NOT @@ -2650,6 +2763,8 @@ async def _handle_multi_turn_success( # pylint: disable=unused-argument if existing_steering: steering_patch = dict(existing_steering) steering_patch["active_input"] = None + steering_patch["active_ack_id"] = None + steering_patch["drain_in_progress"] = False existing_input_slot = task_info.payload.get("input") if _is_ref(existing_input_slot): attachments_patch[_ref_key(existing_input_slot)] = None @@ -2776,6 +2891,27 @@ async def _handle_multi_turn_failure( # pylint: disable=unused-argument :keyword error_dict: The structured error payload for the caller's future. :paramtype error_dict: dict[str, Any] """ + # Structured failure telemetry is emitted regardless of whether the + # terminal transition can be persisted. + active = self._active_tasks.get(task_id) + input_id = None + if active is not None: + input_id = getattr(active.context, "input_id", None) + logger.error( + "resilient_task_handler_failure: task=%s exc_type=%s", + task_id, + type(exc).__name__, + extra={ + "event": "resilient_task_handler_failure", + "event_name": "resilient_task_handler_failure", + "task_id": task_id, + "input_id": input_id, + "error_type": type(exc).__name__, + "error_message": str(exc), + "primitive": "multi_turn_task", + }, + ) + # Steps 2 + 3: PATCH to suspended (NOT completed); clear input + _retry_attempt # + _steering.active_input + promoted _input attachment if any (SOT §23.8 # single-PATCH invariant); NO payload["error"] written; _last_input_id preserved. @@ -2789,6 +2925,8 @@ async def _handle_multi_turn_failure( # pylint: disable=unused-argument if existing_steering: steering_patch = dict(existing_steering) steering_patch["active_input"] = None + steering_patch["active_ack_id"] = None + steering_patch["drain_in_progress"] = False existing_input_slot = task_info.payload.get("input") if _is_ref(existing_input_slot): attachments_patch[_ref_key(existing_input_slot)] = None @@ -2811,10 +2949,6 @@ async def _handle_multi_turn_failure( # pylint: disable=unused-argument attachments=attachments_patch or None, ), ) - except TaskConflictError: - # 412 RE-READ decided ABANDON; propagate so the active caller - # receives the eviction-shape exception. - raise except _HostedConflict as hosted_exc: translated = _translate_hosted_conflict(hosted_exc, task_id=task_id) if translated is None: @@ -2830,33 +2964,6 @@ async def _handle_multi_turn_failure( # pylint: disable=unused-argument ) raise TaskConflictError(task_id, "in_progress") from transport_exc raise - except Exception: # noqa: BLE001 - logger.warning( - "Failed to PATCH multi-turn suspended-on-raise for task %s", - task_id, - exc_info=True, - ) - # — structured failure log/telemetry for every handler - # failure, independent of listener presence. Logged at ERROR per - # (the chain has just lost a turn). - active = self._active_tasks.get(task_id) - input_id = None - if active is not None: - input_id = getattr(active.context, "input_id", None) - logger.error( - "resilient_task_handler_failure: task=%s exc_type=%s", - task_id, - type(exc).__name__, - extra={ - "event": "resilient_task_handler_failure", - "event_name": "resilient_task_handler_failure", - "task_id": task_id, - "input_id": input_id, - "error_type": type(exc).__name__, - "error_message": str(exc), - "primitive": "multi_turn_task", - }, - ) async def _handle_failure( self, @@ -3226,6 +3333,36 @@ def _get_task_write_lock(self, task_id: str) -> asyncio.Lock: self._task_write_locks[task_id] = lock return lock + @asynccontextmanager + async def _task_write_guard(self, task_id: str) -> AsyncIterator[asyncio.Lock]: + """Hold the per-task write lock and reclaim it after its final user. + + Holder/waiter accounting is incremented before lock acquisition, so a + releasing holder cannot delete the registry entry while another caller + is queued on the same lock. + + :param task_id: The task identifier. + :type task_id: str + :return: The acquired per-task write lock. + :rtype: ~collections.abc.AsyncIterator[asyncio.Lock] + """ + lock = self._get_task_write_lock(task_id) + self._task_write_lock_users[task_id] = self._task_write_lock_users.get(task_id, 0) + 1 + try: + async with lock: + yield lock + finally: + remaining = self._task_write_lock_users[task_id] - 1 + if remaining: + self._task_write_lock_users[task_id] = remaining + else: + self._task_write_lock_users.pop(task_id, None) + if task_id not in self._active_tasks and self._task_write_locks.get(task_id) is lock: + self._task_write_locks.pop(task_id, None) + if task_id in self._task_write_cleanup_pending: + self._task_write_cleanup_pending.discard(task_id) + self._task_etag_cache.pop(task_id, None) + def _track_etag(self, task_id: str, etag: str | None) -> None: """— refresh the latest known etag for a task. @@ -3273,7 +3410,11 @@ def _active_tasks_pop(self, task_id: str) -> None: :type task_id: str """ self._active_tasks.pop(task_id, None) + if self._task_write_lock_users.get(task_id): + self._task_write_cleanup_pending.add(task_id) + return self._task_write_locks.pop(task_id, None) + self._task_write_cleanup_pending.discard(task_id) self._task_etag_cache.pop(task_id, None) async def _provider_get_tracked(self, task_id: str) -> Any: @@ -3376,7 +3517,7 @@ async def _provider_update_locked( :return: The provider update response. :rtype: Any """ - async with self._get_task_write_lock(task_id): + async with self._task_write_guard(task_id): return await self._provider_update_lock_held(task_id, patch, force_if_match=force_if_match) async def _terminal_write_locked( # pylint: disable=too-many-statements @@ -3422,7 +3563,7 @@ async def _terminal_write_locked( # pylint: disable=too-many-statements """ prior_lease_owner = patch.lease_owner prior_lease_instance = patch.lease_instance_id - async with self._get_task_write_lock(task_id): + async with self._task_write_guard(task_id): attempts = 0 cached_expiry_count = self._cached_expiry_count(task_id) while True: diff --git a/sdk/agentserver/azure-ai-agentserver-core/docs/task-and-streaming-spec.md b/sdk/agentserver/azure-ai-agentserver-core/docs/task-and-streaming-spec.md index 39d109547120..b9c61327661c 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/docs/task-and-streaming-spec.md +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/task-and-streaming-spec.md @@ -1042,10 +1042,17 @@ chain record itself does not carry the per-turn diagnostic. | Sub-key | Type | Meaning | |---|---|---| | `pending_inputs` | array of input values OR refs (§23) | FIFO of queued steering inputs. | +| `pending_ack_ids` | array of opaque strings or nulls | Internal acknowledgment IDs aligned with `pending_inputs`. Each new append adds a unique ID in the same PATCH; drain and queued cancel remove the corresponding ID. Null marks an untagged slot from a pre-upgrade task record. These are independent of the caller's `input_id`. | | `next_input_seq` | integer | Monotonic counter for promoted-attachment key allocation (NEVER reused). | | `cancel_requested` | boolean | Resilient cancel signal; set on steering append; cleared after drain when pending is empty. | -| `drain_in_progress` | boolean | True between the start of a drain PATCH and the next turn-start; protects against partial drain on crash. | +| `drain_in_progress` | boolean | True after the drain-start PATCH; cleared at turn-start Phase 3 or atomically with active fields at turn end. Protects against partial drain on crash. | | `active_input` | any JSON value OR ref | The single input being drained (mirror copy used by the race-recovery contract). Cleared at suspend / terminal. | +| `active_ack_id` | opaque string or null | Internal acknowledgment ID for `active_input`. Written in the same drain PATCH and cleared with `active_input`; lets a committed drain reconcile after a lost response. | + +Older task records may omit `pending_ack_ids`; those queued slots are treated +as untagged. Workers handling the same chain concurrently must all understand +this field, since older workers do not keep it aligned when draining or +cancelling inputs. Implementers in other languages MUST use these exact key names. A process built in language X must be able to recover a task created @@ -1348,6 +1355,8 @@ These transitions MUST be single PATCHes carrying BOTH `payload` and 3. **Suspend (multi-turn turn-end via `return X`)**: - `payload["input"] = null` - `payload["steering"]["active_input"] = null` + - `payload["steering"]["active_ack_id"] = null` + - `payload["steering"]["drain_in_progress"] = false` - `payload["retry_attempt"] = null` (fresh budget for the next turn) - `attachments["input"] = null` (delete) IF the input was a ref 4. **Steering append (promoted)**: `payload["steering"]["pending_inputs"] @@ -1549,8 +1558,11 @@ Lock lifecycle: - Per-`task_id` `asyncio.Lock` allocated lazily on first write. - Released after the PATCH response is recorded (etag updated). +- A holder/waiter count is registered before acquisition, so releasing + a holder cannot remove the lock while another caller is queued. - Removed from the in-memory lock table when the local active-task - entry is torn down (no leaked locks). + entry is torn down, or when the final holder/waiter exits for a task + that has no local active entry (no leaked remote-task locks). In-process contention now serializes; cross-process contention (another worker reclaimed the lease) still surfaces as 412 because @@ -3024,6 +3036,8 @@ executes this PATCH as a single round-trip: pending.append(input) # raw inline attachments_patch = None 7. steering['pending_inputs'] = pending + steering['pending_ack_ids'] = existing IDs (null-padded for untagged + legacy slots) + [new unique ack ID] steering['cancel_requested'] = True 8. payload_patch = {'steering': steering} if input_id provided: payload_patch['last_input_id'] = input_id @@ -3055,6 +3069,8 @@ Phase 1 — "Drain start" PATCH (atomic across payload + attachments): 4. If pending is empty: return None (no drain happens; caller proceeds to suspend/complete normally). 5. next_entry = pending.pop(0) + next_ack_id = steering['pending_ack_ids'].pop(0) if present, + or null for an untagged legacy entry 6. attachments_patch = {} 7. If next_entry is a ref (§23.3): attachments_patch[ref_key(next_entry)] = None # delete attachment @@ -3062,12 +3078,15 @@ Phase 1 — "Drain start" PATCH (atomic across payload + attachments): else: active_input_value = next_entry 8. steering['active_input'] = active_input_value + steering['active_ack_id'] = next_ack_id 9. steering['pending_inputs'] = pending + steering['pending_ack_ids'] = remaining IDs 10. steering['drain_in_progress'] = True 11. steering['cancel_requested'] = len(pending) > 0 # more pending => keep advisory - 12. payload['steering'] = steering - 13. payload['turn_started_at'] = utc_now_iso() # fresh turn-start boundary - 14. PATCH(task_id, status='in_progress', payload=payload, + 12. payload['input'] = active_input_value or its function-input ref + 13. payload['steering'] = steering + 14. payload['turn_started_at'] = utc_now_iso() # fresh turn-start boundary + 15. PATCH(task_id, status='in_progress', payload=payload, attachments=attachments_patch, lease piggyback, if_match=etag) [NB: status MUST be set to 'in_progress' in this PATCH. The turn-end @@ -3080,36 +3099,49 @@ Phase 1 — "Drain start" PATCH (atomic across payload + attachments): of a suspended->in_progress *claim*. Omitting the status flip makes the Phase-1 PATCH 409 and the steered turn never runs.] - [NB: Phase 1 does NOT set payload['input'] or write a ref/attachment - for active_input. Only the in-memory ctx receives the value (Phase 2). - Recovery from a crash BETWEEN Phase 1 and Phase 3 reads - steering['active_input'] as the source of truth for the input, - via the race-recovery contract. + [NB: Phase 1 publishes the promoted turn as the canonical current + `payload['input']`, using the function-input attachment when required, + in the same PATCH that removes it from the steering queue. Recovery from + a crash BETWEEN Phase 1 and Phase 3 also retains `active_input` as the + steering-specific recovery source. + Only the recursive retry for that exact Phase-1 transition may reconcile + `drain_in_progress=true`, and only when its expected acknowledgment ID + matches `active_ack_id`. A later ordinary drain MUST NOT redispatch the + active turn or consume another pending queue entry. No output co-clear is needed — the framework does not write payload['output'] / output attachments on any transition.] Phase 2 — Handler re-entry (in-memory only): - 15. Construct a fresh TaskContext with: + 16. Construct a fresh TaskContext with: entry_mode='resumed', is_steered_turn=True, input=active_input_value (deserialized via input_type), metadata reused from previous ctx, cancel_event=fresh (re-set if cancel_requested still True), retry_attempt=0. - 16. Update the in-process _ActiveTask.context pointer. - 17. Invoke the handler with the new ctx. + 17. Update the in-process _ActiveTask.context pointer. + 18. Invoke the handler with the new ctx. Phase 3 — "Drain end" PATCH (after handler re-entered): - 18. steering['drain_in_progress'] = False - 19. payload['steering'] = steering - 20. payload['retry_attempt'] = 0 # Drain resets retry budget - 21. PATCH(task_id, payload=payload, lease piggyback) + 19. steering['drain_in_progress'] = False + 20. payload['steering'] = steering + 21. payload['retry_attempt'] = 0 # Drain resets retry budget + 22. PATCH(task_id, payload=payload, lease piggyback) (No attachments touched in Phase 3.) Phase 4 — On the next turn-end: - 22. The handler returns/suspends/raises. The terminal handler clears - active_input as part of its suspend/complete PATCH (§53). + 23. The handler returns/suspends/raises. The terminal handler clears + active_input, active_ack_id, and drain_in_progress together as part + of its suspend/complete PATCH (§53). ``` +**Terminal-persistence barrier.** A turn result, failure, or cancellation is +not acknowledged as final, and no queued successor is dispatched, until the +turn-end PATCH succeeds. If that PATCH fails, the current caller receives the +persistence error, local queued acknowledgment futures receive +`TaskConflictError(current_status="in_progress")`, and the durable task remains +available for recovery. This prevents a stale active turn from being executed +again after its caller already received a definitive handler outcome. + **Race-recovery contract.** If the process crashes: - **Between Phase 1 PATCH and Phase 2 handler entry:** recovery diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_input_cleared_on_suspend.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_input_cleared_on_suspend.py index ec7dbdf8b896..6f73a76d41a0 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_input_cleared_on_suspend.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_input_cleared_on_suspend.py @@ -69,3 +69,60 @@ async def suspending(ctx: TaskContext[dict]) -> None: ) finally: await _teardown_manager(manager, mgr_mod) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turn_outcome", ["return", "raise"]) +async def test_turn_end_atomically_clears_stale_drain_state(tmp_path: Path, turn_outcome: str) -> None: + """A failed Phase-3 clear cannot leave a phantom active drain.""" + from azure.ai.agentserver.core.tasks._models import TaskCreateRequest + + manager, mgr_mod = await _setup_manager(tmp_path) + task_id = f"t-stale-drain-{turn_outcome}" + try: + await manager.provider.create( + TaskCreateRequest( + id=task_id, + agent_name="test-agent", + session_id="test-session", + status="in_progress", + title="stale drain", + payload={ + "input": {"msg": "B"}, + "schema_version": "1", + "steering": { + "active_input": {"msg": "B"}, + "active_ack_id": "ack-b", + "pending_inputs": [{"msg": "C"}], + "pending_ack_ids": ["ack-c"], + "cancel_requested": True, + "drain_in_progress": True, + }, + }, + lease_owner=manager._lease_owner, + lease_instance_id=manager._instance_id, + lease_duration_seconds=60, + ) + ) + + if turn_outcome == "return": + await manager._handle_multi_turn_success(task_id=task_id, opts=None) + else: + await manager._handle_multi_turn_failure( + task_id=task_id, + exc=RuntimeError("planned failure"), + opts=None, + error_dict={"type": "RuntimeError", "message": "planned failure", "traceback": ""}, + ) + + info = await manager.provider.get(task_id) + assert info is not None + assert info.status == "suspended" + steering = info.payload["steering"] + assert steering["active_input"] is None + assert steering["active_ack_id"] is None + assert steering["drain_in_progress"] is False + assert steering["pending_inputs"] == [{"msg": "C"}] + assert steering["pending_ack_ids"] == ["ack-c"] + finally: + await _teardown_manager(manager, mgr_mod) diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering.py index dde0b9884c25..ecfa3b6bfd25 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering.py @@ -15,6 +15,7 @@ task, EntryMode, SteeringQueueFull, + TaskCancelled, TaskConflictError, multi_turn_task, ) @@ -123,40 +124,388 @@ async def regular(ctx: TaskContext[dict]) -> dict: @pytest.mark.asyncio async def test_steering_queue_full(self, tmp_path): - """start() raises SteeringQueueFull when queue is at capacity. - - : the per-task ``max_pending`` knob was - demoted; the framework-wide default - ``_DEFAULT_MAX_PENDING_STEERING`` (10) applies. This test fills the - queue at that default to verify the exception still surfaces. - """ + """Rejected steers do not leak futures or steal accepted acknowledgments.""" from azure.ai.agentserver.core.tasks._decorator import _DEFAULT_MAX_PENDING_STEERING manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() try: - gate = asyncio.Event() - @multi_turn_task(name="chat", steerable=True) async def chat(ctx: TaskContext[dict]) -> dict: await gate.wait() - return {"msg": "done"} + return {"msg": ctx.input["msg"]} run1 = await chat.start(task_id="t1", input={"msg": "A"}) - # Fill the queue to the framework default + accepted = [] for i in range(_DEFAULT_MAX_PENDING_STEERING): - await chat.start(task_id="t1", input={"msg": f"fill-{i}"}) + accepted.append(await chat.start(task_id="t1", input={"msg": f"fill-{i}"})) + registered = dict(manager._pending_steering_futures["t1"]) + assert len(registered) == _DEFAULT_MAX_PENDING_STEERING + + for i in range(3): + with pytest.raises(SteeringQueueFull): + await chat.start(task_id="t1", input={"msg": f"overflow-{i}"}) + assert manager._pending_steering_futures["t1"] == registered + + info = await manager.provider.get("t1") + assert info is not None + assert [entry["msg"] for entry in info.payload["steering"]["pending_inputs"]] == [ + f"fill-{i}" for i in range(_DEFAULT_MAX_PENDING_STEERING) + ] - # Queue is full — should raise - with pytest.raises(SteeringQueueFull): - await chat.start(task_id="t1", input={"msg": "overflow"}) + gate.set() + assert await asyncio.wait_for(run1.result(), timeout=5.0) == {"msg": "A"} + for i, run in enumerate(accepted): + assert await asyncio.wait_for(run.result(), timeout=5.0) == {"msg": f"fill-{i}"} + assert not manager._pending_steering_futures.get("t1") + finally: + gate.set() + await self._teardown_manager(manager, mgr_mod) - #: SteeringQueueFull is bare exception (no max_pending) + @pytest.mark.asyncio + async def test_steering_storage_failure_unregisters_only_failed_future(self, tmp_path, monkeypatch): + manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + await gate.wait() + return {"msg": ctx.input["msg"]} + + first = await chat.start(task_id="t1", input={"msg": "first"}) + accepted = await chat.start(task_id="t1", input={"msg": "accepted"}) + registered = dict(manager._pending_steering_futures["t1"]) + accepted_future = next(iter(registered.values())) + rejected_future: asyncio.Future[Any] | None = None + + async def reject_update(task_id: str, _patch: Any) -> Any: + nonlocal rejected_future + rejected_future = list(manager._pending_steering_futures[task_id].values())[-1] + raise OSError("storage unavailable") + + with monkeypatch.context() as patcher: + patcher.setattr(manager.provider, "update", reject_update) + with pytest.raises(OSError, match="storage unavailable"): + await chat.start(task_id="t1", input={"msg": "rejected"}) + + assert rejected_future is not None + assert rejected_future.cancelled() + assert manager._pending_steering_futures["t1"] == registered + assert not accepted_future.done() + info = await manager.provider.get("t1") + assert info is not None + assert info.payload["steering"]["pending_inputs"] == [{"msg": "accepted"}] gate.set() - await asyncio.wait_for(run1.result(), timeout=5.0) + assert await asyncio.wait_for(first.result(), timeout=5.0) == {"msg": "first"} + assert await asyncio.wait_for(accepted.result(), timeout=5.0) == {"msg": "accepted"} + finally: + gate.set() + await self._teardown_manager(manager, mgr_mod) + + @pytest.mark.asyncio + async def test_steering_cancelled_append_unregisters_future(self, tmp_path, monkeypatch): + manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() + append_started = asyncio.Event() + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + await gate.wait() + return {"msg": ctx.input["msg"]} + + first = await chat.start(task_id="t1", input={"msg": "first"}) + aborted_future: asyncio.Future[Any] | None = None + + async def blocked_update(task_id: str, _patch: Any) -> Any: + nonlocal aborted_future + aborted_future = list(manager._pending_steering_futures[task_id].values())[-1] + append_started.set() + await asyncio.Event().wait() + + with monkeypatch.context() as patcher: + patcher.setattr(manager.provider, "update", blocked_update) + aborted = asyncio.create_task(chat.start(task_id="t1", input={"msg": "aborted"})) + try: + await asyncio.wait_for(append_started.wait(), timeout=5.0) + assert manager._get_task_write_lock("t1").locked() + finally: + aborted.cancel() + with pytest.raises(asyncio.CancelledError): + await aborted + + assert aborted_future is not None + assert aborted_future.cancelled() + assert "t1" not in manager._pending_steering_futures + info = await manager.provider.get("t1") + assert info is not None + assert info.payload.get("steering", {}).get("pending_inputs", []) == [] + + accepted = await chat.start(task_id="t1", input={"msg": "accepted"}) + gate.set() + assert await asyncio.wait_for(first.result(), timeout=5.0) == {"msg": "first"} + assert await asyncio.wait_for(accepted.result(), timeout=5.0) == {"msg": "accepted"} + finally: + gate.set() + await self._teardown_manager(manager, mgr_mod) + + @pytest.mark.asyncio + @pytest.mark.parametrize("cancel_request", [False, True]) + async def test_committed_append_without_response_preserves_later_ack(self, tmp_path, monkeypatch, cancel_request): + manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() + persisted = asyncio.Event() + seen: list[str] = [] + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + await gate.wait() + seen.append(ctx.input["msg"]) + return {"msg": ctx.input["msg"], "turn": len(seen)} + + first = await chat.start(task_id="t1", input={"msg": "first"}) + earlier = await chat.start(task_id="t1", input={"msg": "earlier"}) + original_update = manager.provider.update + + async def commit_without_response(task_id: str, patch: Any) -> Any: + await original_update(task_id, patch) + persisted.set() + if cancel_request: + await asyncio.Event().wait() + raise OSError("response lost after commit") + + with monkeypatch.context() as patcher: + patcher.setattr(manager.provider, "update", commit_without_response) + if cancel_request: + interrupted = asyncio.create_task( + chat.start(task_id="t1", input={"msg": "duplicate"}, input_id="same-id") + ) + try: + await asyncio.wait_for(persisted.wait(), timeout=5.0) + finally: + interrupted.cancel() + with pytest.raises(asyncio.CancelledError): + await interrupted + else: + with pytest.raises(OSError, match="response lost after commit"): + await chat.start(task_id="t1", input={"msg": "duplicate"}, input_id="same-id") + + later = await chat.start(task_id="t1", input={"msg": "duplicate"}, input_id="same-id") + info = await manager.provider.get("t1") + assert info is not None + assert [entry["msg"] for entry in info.payload["steering"]["pending_inputs"]] == [ + "earlier", + "duplicate", + "duplicate", + ] + ack_ids = info.payload["steering"]["pending_ack_ids"] + assert len(ack_ids) == len(set(ack_ids)) == 3 + + gate.set() + assert await asyncio.wait_for(first.result(), timeout=5.0) == {"msg": "first", "turn": 1} + assert await asyncio.wait_for(earlier.result(), timeout=5.0) == {"msg": "earlier", "turn": 2} + assert await asyncio.wait_for(later.result(), timeout=5.0) == {"msg": "duplicate", "turn": 4} + assert seen == ["first", "earlier", "duplicate", "duplicate"] + finally: + gate.set() + await self._teardown_manager(manager, mgr_mod) + + @pytest.mark.asyncio + async def test_committed_append_conflict_reconciles_ack_before_precondition_and_capacity( + self, tmp_path, monkeypatch + ): + from azure.ai.agentserver.core.tasks._client import TransportClassifiedError + + manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + await gate.wait() + return {"msg": ctx.input["msg"]} + + first = await chat.start(task_id="t1", input={"msg": "first"}, input_id="input-0") + accepted = [] + predecessor = "input-0" + for i in range(1, 9): + input_id = f"input-{i}" + accepted.append( + await chat.start( + task_id="t1", + input={"msg": f"accepted-{i}"}, + input_id=input_id, + if_last_input_id=predecessor, + ) + ) + predecessor = input_id + + original_update = manager.provider.update + update_calls = 0 + + async def commit_then_conflict(task_id: str, patch: Any) -> Any: + nonlocal update_calls + update_calls += 1 + result = await original_update(task_id, patch) + raise TransportClassifiedError( + status=412, + classification="conflict", + message="response retry observed the committed etag", + ) + + with monkeypatch.context() as patcher: + patcher.setattr(manager.provider, "update", commit_then_conflict) + final = await chat.start( + task_id="t1", + input={"msg": "final"}, + input_id="input-9", + if_last_input_id=predecessor, + ) + + assert update_calls == 1 + info = await manager.provider.get("t1") + assert info is not None + pending = info.payload["steering"]["pending_inputs"] + ack_ids = info.payload["steering"]["pending_ack_ids"] + assert len(pending) == len(ack_ids) == 9 + assert pending.count({"msg": "final"}) == 1 + assert len(ack_ids) == len(set(ack_ids)) + assert info.payload["last_input_id"] == "input-9" + + await final.cancel() + with pytest.raises(TaskCancelled): + await final.result() + info = await manager.provider.get("t1") + assert info is not None + assert len(info.payload["steering"]["pending_inputs"]) == 8 + assert {"msg": "final"} not in info.payload["steering"]["pending_inputs"] + + gate.set() + assert await asyncio.wait_for(first.result(), timeout=5.0) == {"msg": "first"} + for i, run in enumerate(accepted, start=1): + assert await asyncio.wait_for(run.result(), timeout=5.0) == {"msg": f"accepted-{i}"} + finally: + gate.set() + await self._teardown_manager(manager, mgr_mod) + + @pytest.mark.asyncio + async def test_rejected_remote_steers_do_not_cache_write_locks(self, tmp_path): + from azure.ai.agentserver.core.tasks._models import TaskCreateRequest + + manager, mgr_mod = await self._setup_manager(tmp_path) + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + return {"msg": ctx.input["msg"]} + + for i in range(20): + task_id = f"remote-full-{i}" + await manager.provider.create( + TaskCreateRequest( + id=task_id, + agent_name="test-agent", + session_id="test-session", + status="in_progress", + title="remote full queue", + payload={ + "input": {"msg": "remote-active"}, + "schema_version": "1", + "steering": { + "pending_inputs": [{"msg": f"queued-{j}"} for j in range(9)], + "pending_ack_ids": [None] * 9, + "next_input_seq": 0, + "cancel_requested": True, + }, + }, + lease_owner="foreign-agent:foreign-session", + lease_instance_id=f"foreign-instance-{i}", + lease_duration_seconds=60, + ) + ) + + with pytest.raises(SteeringQueueFull): + await chat.start(task_id=task_id, input={"msg": "rejected"}) + + assert task_id not in manager._task_write_locks + + assert manager._task_write_locks == {} + finally: + await self._teardown_manager(manager, mgr_mod) + + @pytest.mark.asyncio + async def test_legacy_backlog_does_not_claim_new_steering_ack(self, tmp_path): + from azure.ai.agentserver.core.tasks._models import TaskPatchRequest + + manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() + seen: list[str] = [] + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + await gate.wait() + seen.append(ctx.input["msg"]) + return {"msg": ctx.input["msg"], "turn": len(seen)} + + first = await chat.start(task_id="t1", input={"msg": "first"}) + existing = await manager.provider.get("t1") + assert existing is not None + steering = dict(existing.payload.get("steering") or {}) + steering["pending_inputs"] = [{"msg": "remote"}] + steering["cancel_requested"] = True + await manager.provider.update("t1", TaskPatchRequest(payload={"steering": steering}, if_match=existing.etag)) + + local = await chat.start(task_id="t1", input={"msg": "local"}) + info = await manager.provider.get("t1") + assert info is not None + ack_ids = info.payload["steering"]["pending_ack_ids"] + assert ack_ids[0] is None + assert isinstance(ack_ids[1], str) + + gate.set() + assert await asyncio.wait_for(first.result(), timeout=5.0) == {"msg": "first", "turn": 1} + assert await asyncio.wait_for(local.result(), timeout=5.0) == {"msg": "local", "turn": 3} + assert seen == ["first", "remote", "local"] + finally: + gate.set() + await self._teardown_manager(manager, mgr_mod) + + @pytest.mark.asyncio + async def test_queued_cancel_keeps_identical_input_owned_by_other_handle(self, tmp_path): + manager, mgr_mod = await self._setup_manager(tmp_path) + gate = asyncio.Event() + try: + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + await gate.wait() + return {"msg": ctx.input["msg"]} + + first = await chat.start(task_id="t1", input={"msg": "first"}) + earlier = await chat.start(task_id="t1", input={"msg": "same"}, input_id="shared-id") + later = await chat.start(task_id="t1", input={"msg": "same"}, input_id="shared-id") + info = await manager.provider.get("t1") + assert info is not None + earlier_ack_id = info.payload["steering"]["pending_ack_ids"][0] + + await later.cancel() + with pytest.raises(TaskCancelled): + await later.result() + info = await manager.provider.get("t1") + assert info is not None + assert info.payload["steering"]["pending_inputs"] == [{"msg": "same"}] + assert info.payload["steering"]["pending_ack_ids"] == [earlier_ack_id] + + gate.set() + assert await asyncio.wait_for(first.result(), timeout=5.0) == {"msg": "first"} + assert await asyncio.wait_for(earlier.result(), timeout=5.0) == {"msg": "same"} finally: + gate.set() await self._teardown_manager(manager, mgr_mod) @pytest.mark.asyncio @@ -780,6 +1129,237 @@ async def _setup_manager(self, provider): await manager.startup() return manager, mgr_mod + @staticmethod + def _commit_then_conflict_provider(delegate): + from azure.ai.agentserver.core.tasks._client import TransportClassifiedError + + class DrainCommitThenConflictProvider: + def __init__(self): + self._delegate = delegate + self._armed = True + self.drain_conflicts = 0 + + async def create(self, request): + return await self._delegate.create(request) + + async def get(self, task_id): + return await self._delegate.get(task_id) + + async def list(self, **kwargs): + return await self._delegate.list(**kwargs) + + async def delete(self, task_id, *, force=False, cascade=False): + await self._delegate.delete(task_id, force=force, cascade=cascade) + + async def update(self, task_id, patch): + payload = getattr(patch, "payload", None) or {} + steering = payload.get("steering") or {} + is_drain = steering.get("drain_in_progress") is True + if self._armed and is_drain: + self._armed = False + self.drain_conflicts += 1 + await self._delegate.update(task_id, patch) + raise TransportClassifiedError( + status=412, + classification="conflict", + message="response retry observed the committed drain etag", + ) + return await self._delegate.update(task_id, patch) + + return DrainCommitThenConflictProvider() + + @staticmethod + def _failed_turn_cleanup_provider(delegate): + from azure.core.exceptions import ServiceRequestError + + class FailedTurnCleanupProvider: + def __init__(self): + self._delegate = delegate + self.phase3_failures = 0 + self.terminal_failures = 0 + + async def create(self, request): + return await self._delegate.create(request) + + async def get(self, task_id): + return await self._delegate.get(task_id) + + async def list(self, **kwargs): + return await self._delegate.list(**kwargs) + + async def delete(self, task_id, *, force=False, cascade=False): + await self._delegate.delete(task_id, force=force, cascade=cascade) + + async def update(self, task_id, patch): + payload = getattr(patch, "payload", None) or {} + steering = payload.get("steering") or {} + is_phase3_clear = ( + getattr(patch, "status", None) is None + and steering.get("drain_in_progress") is False + and steering.get("active_input") is not None + ) + is_terminal_suspend = ( + getattr(patch, "status", None) == "suspended" + and steering.get("active_input") is None + and self.phase3_failures == 1 + ) + if is_phase3_clear and self.phase3_failures == 0: + self.phase3_failures += 1 + raise ServiceRequestError("injected Phase-3 cleanup failure") + if is_terminal_suspend and self.terminal_failures == 0: + self.terminal_failures += 1 + raise ServiceRequestError("injected terminal persistence failure") + return await self._delegate.update(task_id, patch) + + return FailedTurnCleanupProvider() + + @pytest.mark.asyncio + @pytest.mark.parametrize("queue_next", [False, True]) + @pytest.mark.parametrize("turn_outcome", ["return", "raise", "cancel"]) + async def test_failed_phase3_and_terminal_write_do_not_redispatch_active_turn( + self, tmp_path, queue_next, turn_outcome + ): + from azure.core.exceptions import ServiceRequestError + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + + provider = self._failed_turn_cleanup_provider(LocalFileTaskProvider(Path(str(tmp_path)))) + manager, mgr_mod = await self._setup_manager(provider) + first_release = asyncio.Event() + first_started = asyncio.Event() + second_release = asyncio.Event() + second_started = asyncio.Event() + seen: list[str] = [] + task_id = f"t-double-cleanup-{turn_outcome}-{queue_next}" + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + message = ctx.input["msg"] + seen.append(message) + if message == "A": + first_started.set() + await first_release.wait() + return {"msg": message} + if message == "B": + second_started.set() + await second_release.wait() + if turn_outcome == "return": + return {"msg": message} + if turn_outcome == "cancel": + raise asyncio.CancelledError() + raise RuntimeError("planned B failure") + return {"msg": message} + + run_a = await chat.start(task_id=task_id, input={"msg": "A"}) + await asyncio.wait_for(first_started.wait(), timeout=5.0) + run_b = await chat.start(task_id=task_id, input={"msg": "B"}) + first_release.set() + await asyncio.wait_for(second_started.wait(), timeout=5.0) + run_c = await chat.start(task_id=task_id, input={"msg": "C"}) if queue_next else None + second_release.set() + + assert await asyncio.wait_for(run_a.result(), timeout=5.0) == {"msg": "A"} + with pytest.raises(ServiceRequestError, match="terminal persistence failure"): + await asyncio.wait_for(run_b.result(), timeout=5.0) + if run_c is not None: + with pytest.raises(TaskConflictError) as exc_info: + await asyncio.wait_for(run_c.result(), timeout=5.0) + assert exc_info.value.current_status == "in_progress" + + for _ in range(100): + if task_id not in manager._active_tasks: + break + await asyncio.sleep(0.01) + + assert provider.phase3_failures == 1 + assert provider.terminal_failures == 1 + assert seen == ["A", "B"] + info = await manager.provider.get(task_id) + assert info is not None + assert info.status == "in_progress" + steering = info.payload["steering"] + assert steering["active_input"] == {"msg": "B"} + assert steering["drain_in_progress"] is True + expected_pending = [{"msg": "C"}] if queue_next else [] + assert steering["pending_inputs"] == expected_pending + finally: + first_release.set() + second_release.set() + await manager.shutdown() + mgr_mod._manager = None + + @pytest.mark.asyncio + async def test_committed_single_slot_drain_conflict_binds_accepted_future(self, tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + + provider = self._commit_then_conflict_provider(LocalFileTaskProvider(Path(str(tmp_path)))) + manager, mgr_mod = await self._setup_manager(provider) + gate = asyncio.Event() + started = asyncio.Event() + seen: list[str] = [] + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + message = ctx.input["msg"] + seen.append(message) + if message == "A": + started.set() + await gate.wait() + return {"msg": message} + + run_a = await chat.start(task_id="t1", input={"msg": "A"}) + await asyncio.wait_for(started.wait(), timeout=5.0) + run_b = await chat.start(task_id="t1", input={"msg": "B"}) + gate.set() + + assert await asyncio.wait_for(run_a.result(), timeout=5.0) == {"msg": "A"} + assert await asyncio.wait_for(run_b.result(), timeout=5.0) == {"msg": "B"} + assert provider.drain_conflicts == 1 + assert seen == ["A", "B"] + assert not manager._pending_steering_futures.get("t1") + finally: + gate.set() + await manager.shutdown() + mgr_mod._manager = None + + @pytest.mark.asyncio + async def test_committed_multi_slot_drain_conflict_does_not_skip_queue_head(self, tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + + provider = self._commit_then_conflict_provider(LocalFileTaskProvider(Path(str(tmp_path)))) + manager, mgr_mod = await self._setup_manager(provider) + gate = asyncio.Event() + started = asyncio.Event() + seen: list[str] = [] + try: + + @multi_turn_task(name="chat", steerable=True) + async def chat(ctx: TaskContext[dict]) -> dict: + message = ctx.input["msg"] + seen.append(message) + if message == "A": + started.set() + await gate.wait() + return {"msg": message} + + run_a = await chat.start(task_id="t1", input={"msg": "A"}) + await asyncio.wait_for(started.wait(), timeout=5.0) + run_b = await chat.start(task_id="t1", input={"msg": "B"}) + run_c = await chat.start(task_id="t1", input={"msg": "C"}) + gate.set() + + assert await asyncio.wait_for(run_a.result(), timeout=5.0) == {"msg": "A"} + assert await asyncio.wait_for(run_b.result(), timeout=5.0) == {"msg": "B"} + assert await asyncio.wait_for(run_c.result(), timeout=5.0) == {"msg": "C"} + assert provider.drain_conflicts == 1 + assert seen == ["A", "B", "C"] + assert not manager._pending_steering_futures.get("t1") + finally: + gate.set() + await manager.shutdown() + mgr_mod._manager = None + @pytest.mark.asyncio async def test_drain_recovers_from_cross_process_conflict(self, tmp_path): from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering_attachment_queue.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering_attachment_queue.py index 3d66fbd61994..b0f72aeb612c 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering_attachment_queue.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_steering_attachment_queue.py @@ -256,6 +256,8 @@ async def runner(ctx: TaskContext[dict]) -> dict: assert len(pending_pre) == 2 assert _ref_key(pending_pre[0]) == "steering_input_0" assert _ref_key(pending_pre[1]) == "steering_input_1" + ack_ids_pre = info_pre.payload["steering"]["pending_ack_ids"] + assert len(ack_ids_pre) == len(set(ack_ids_pre)) == 2 assert info_pre.attachments["steering_input_0"] == a_value assert info_pre.attachments["steering_input_1"] == b_value @@ -269,6 +271,7 @@ async def runner(ctx: TaskContext[dict]) -> dict: pending_mid = info_mid.payload["steering"]["pending_inputs"] # Only B left in the queue. assert len(pending_mid) == 1 + assert info_mid.payload["steering"]["pending_ack_ids"] == [ack_ids_pre[1]] # B's attachment key MUST still be steering_input_1 (not renamed to _0). assert _ref_key(pending_mid[0]) == "steering_input_1" # A's attachment is gone; B's is unchanged. diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_write_queue.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_write_queue.py index ea837f6444dd..ca0fa18e1fa8 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_write_queue.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_write_queue.py @@ -110,7 +110,7 @@ async def test_reads_do_not_acquire_lock(local) -> None: @multi_turn_task(name="reads_no_lock") async def my_task(ctx: TaskContext[str]) -> str: del ctx - async with manager._get_task_write_lock("t-reads"): # pylint: disable=protected-access + async with manager._task_write_guard("t-reads"): # pylint: disable=protected-access in_write_barrier.set() await release_write.wait() return "done" @@ -137,6 +137,83 @@ async def my_task(ctx: TaskContext[str]) -> str: mgr_mod._manager = None +@pytest.mark.asyncio +async def test_remote_write_lock_cleanup_waits_for_queued_users(local) -> None: + """Remote-task lock cleanup cannot split queued callers across two locks.""" + manager = TaskManager(config=_config_stub(), provider=local) + first_acquired = asyncio.Event() + release_first = asyncio.Event() + second_acquired = asyncio.Event() + release_second = asyncio.Event() + third_acquired = asyncio.Event() + + async def _first() -> None: + async with manager._task_write_guard("t-remote"): # pylint: disable=protected-access + first_acquired.set() + await release_first.wait() + + async def _second() -> None: + async with manager._task_write_guard("t-remote"): # pylint: disable=protected-access + second_acquired.set() + await release_second.wait() + + async def _third() -> None: + async with manager._task_write_guard("t-remote"): # pylint: disable=protected-access + third_acquired.set() + + first = asyncio.create_task(_first()) + await first_acquired.wait() + shared_lock = manager._task_write_locks["t-remote"] # pylint: disable=protected-access + second = asyncio.create_task(_second()) + await asyncio.sleep(0) + assert manager._task_write_lock_users["t-remote"] == 2 # pylint: disable=protected-access + + release_first.set() + await second_acquired.wait() + assert manager._task_write_locks["t-remote"] is shared_lock # pylint: disable=protected-access + + third = asyncio.create_task(_third()) + await asyncio.sleep(0) + assert not third_acquired.is_set() + assert manager._task_write_lock_users["t-remote"] == 2 # pylint: disable=protected-access + + release_second.set() + await asyncio.gather(first, second, third) + assert third_acquired.is_set() + assert "t-remote" not in manager._task_write_locks # pylint: disable=protected-access + assert "t-remote" not in manager._task_write_lock_users # pylint: disable=protected-access + + +@pytest.mark.asyncio +async def test_active_teardown_defers_registry_cleanup_until_waiter_exits(local) -> None: + """Active teardown preserves one shared lock until its queued user exits.""" + manager = TaskManager(config=_config_stub(), provider=local) + acquired = asyncio.Event() + release = asyncio.Event() + + async def _holder() -> None: + async with manager._task_write_guard("t-active"): # pylint: disable=protected-access + acquired.set() + await release.wait() + + manager._active_tasks["t-active"] = object() # type: ignore[assignment] # pylint: disable=protected-access + manager._task_etag_cache["t-active"] = "etag-1" # pylint: disable=protected-access + holder = asyncio.create_task(_holder()) + await acquired.wait() + + manager._active_tasks_pop("t-active") # pylint: disable=protected-access + assert "t-active" in manager._task_write_locks # pylint: disable=protected-access + assert "t-active" in manager._task_write_cleanup_pending # pylint: disable=protected-access + assert manager._task_etag_cache["t-active"] == "etag-1" # pylint: disable=protected-access + + release.set() + await holder + assert "t-active" not in manager._task_write_locks # pylint: disable=protected-access + assert "t-active" not in manager._task_write_lock_users # pylint: disable=protected-access + assert "t-active" not in manager._task_write_cleanup_pending # pylint: disable=protected-access + assert "t-active" not in manager._task_etag_cache # pylint: disable=protected-access + + @pytest.mark.asyncio async def test_lock_removed_when_active_entry_torn_down(local) -> None: """/ C-WQ-1 — when the task's active-entry is torn