diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/README.md new file mode 100644 index 000000000000..d9eaec407ace --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/README.md @@ -0,0 +1,89 @@ +# Minimal resilient long-running agent — **cancellation** + +A durable, **finite** job that would finish on its own but can be **cancelled +mid-run**. It is the smallest end-to-end illustration of the cancel flow for a +long-running agent (LRA), and depends on **only** the two agentserver packages — +no LLM, no `azure-ai-projects`, no `langgraph`. + +It complements the other minimal samples: + +| Sample | Shape | Cancel | +|--------|-------|--------| +| `resilient_hello_world` | finite, runs to completion | — | +| `resilient_cancellable` | **finite, can be stopped early** | **yes** | +| `resilient_hello_forever` | indefinite, must be stopped | yes | + +## How cancel works + +The job counts through `steps` steps, checkpointing after each. The cancel +endpoint writes a durable **cancel marker** to a separate state-store key; the +job reads that marker **before every step** and, if present, records +`status: "cancelled"` and returns without finishing the remaining steps. + +Using a durable marker rather than the in-process `ctx.cancel` event matters: + +1. **Cross-replica** — the cancel request may land on a different replica than + the one running the job; a durable marker is visible to both. +2. **Crash-durable** — a job recovered after a cancel was requested still sees + the marker and stops. +3. **No ETag race** — the marker lives in its own key, so the cancel write never + collides with the checkpoint's ETag. + +## Run it + +```bash +pip install -r requirements.txt +python app.py # listens on http://localhost:8088 +``` + +Pass the same `?agent_session_id=` on every call so poll/cancel hit the same +session-scoped store: + +```bash +# start a 30-step job (STEP_DELAY defaults to 2s, so ~60s of work) +curl -s -XPOST -H "Content-Type: application/json" \ + -d '{"name": "Ada", "steps": 30}' \ + "http://localhost:8088/invocations?agent_session_id=demo" +# -> {"status": "started", "invocation_id": "", "total_steps": 30} + +# poll — completed_steps climbs while status is "in_progress" +curl -s "http://localhost:8088/invocations/?agent_session_id=demo" +# -> {"status": "in_progress", "completed_steps": 4, "total_steps": 30} + +# cancel mid-run +curl -s -XPOST "http://localhost:8088/invocations//cancel?agent_session_id=demo" +# -> {"status": "cancelling", "invocation_id": ""} + +# poll again — the job stopped early +curl -s "http://localhost:8088/invocations/?agent_session_id=demo" +# -> {"status": "cancelled", "completed_steps": 4, "total_steps": 30} +``` + +`status` is one of `in_progress` / `cancelling` / `cancelled` / `completed` / +`failed`. `cancelling` is the brief window after a cancel is requested but before +the job reaches its next step and finalizes. + +## The key line: enable resilient tasks + +As of `azure-ai-agentserver-core` **2.1.0b1** the durable-task subsystem is +**strictly opt-in**. Before host startup: + +```python +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled +set_resilient_tasks_enabled(True) +``` + +Without it, `cancellable_job.start()` raises `TaskManagerNotInitialized` and +there is no crash recovery. + +## Cancel survives a crash + +1. Start a 30-step job and request cancel while it is a few steps in. +2. Before the job observes the marker, **hard-kill the process** — an *ungraceful* + termination such as `kill -9 ` (SIGKILL). Do **not** use Ctrl-C: that + triggers the host's graceful shutdown, which gives the running task up to ~25s + to finish, during which it observes the marker and exits terminally on its own + — leaving nothing to recover. +3. Restart `python app.py`. The recovery scan re-enters the job with + `ctx.entry_mode == "recovered"`; it reads the still-present cancel marker on + its next step and stops with `status: "cancelled"` — the cancel is not lost. diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/__init__.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/agent.py new file mode 100644 index 000000000000..8ac73f8913d6 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/agent.py @@ -0,0 +1,254 @@ +"""Minimal resilient long-running agent showcasing **cooperative cancellation**. + +Where ``resilient_hello_world`` runs to completion and ``resilient_hello_forever`` +runs until it is stopped, this sample sits in between: a **finite** job that +*would* finish on its own but can be **cancelled mid-run**. It is the smallest +end-to-end illustration of the cancel flow for a durable long-running agent. + +It depends on ONLY ``azure-ai-agentserver-core`` and +``azure-ai-agentserver-invocations`` — no LLM, no cloud — and checkpoints its +progress after every step to a durable state store (``FoundryStateStore``, which +uses a local on-disk backend outside Foundry, so no Azure resources are needed +locally). + +How cancel works (the important bit): + +- The cancel endpoint (``app.py``) writes a durable **cancel marker** to a + separate state-store key. It does NOT rely on an in-process signal. +- Before each step, the task reads that marker. If present, it stops early, + records ``status: "cancelled"`` in its checkpoint, and returns. + +Using a durable marker (rather than the in-process ``ctx.cancel`` event) makes +cancellation correct even when the cancel request lands on a *different* replica +than the one running the task, and it survives a crash/redeploy: a task recovered +after a cancel was requested still sees the marker and stops. Keeping the marker +in its own key means the cancel write never races the checkpoint's ETag. + +Input schema: ``{"name": str, "steps": int?}``. The host also injects the +invocation's ``session_id``/``user_id``/``call_id`` into the durable input so +recovery reopens the same user-isolated store partition and Foundry call identity. + +Environment: + +- ``STEP_DELAY`` — seconds to sleep between steps (default ``2``). Keep it + nonzero so a cancel (or crash) demo has time to land mid-run. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +from typing import Any + +from azure.ai.agentserver.core.storage import FoundryStateStore +from azure.ai.agentserver.core.tasks import TaskContext, task + +logger = logging.getLogger(__name__) + +_STEP_DELAY = float(os.environ.get("STEP_DELAY", "2")) + +# Suffix for the durable "cancel" marker key. The cancel endpoint (app.py) writes +# this key; the task checks it before each step to decide whether to stop early. +# Keeping it in a SEPARATE key means the cancel write never races the +# checkpoint's ETag. +CANCEL_SUFFIX = "/cancel" + + +def checkpoint_store_name(session_id: str) -> str: + """Return the **session-isolated** checkpoint store name. + + ``FoundryStateStore`` is agent-scoped and has no built-in per-session + isolation, so the store is namespaced by the invocation's session id (as the + other resilient samples do). Shared with app.py so the poll/cancel endpoints + read the same scope. + + The session component is **hashed**: a protocol session id can be up to 256 + characters, and the local ``FoundryStateStore`` backend base64-encodes the + whole store name into a single filename, which would blow past the 255-byte + ``NAME_MAX`` and fail this no-cloud sample with ``ENAMETOOLONG``. A fixed-width + SHA-256 digest keeps the name bounded while remaining unique per session. + """ + digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest() + return f"resilient-cancellable/{digest}" + + +def durable_task_id(session_id: str, invocation_id: str, user_id: str) -> str: + """Return the TaskManager task id derived from the user, session and invocation. + + The invocations protocol accepts a *caller-supplied* invocation id, and a + single agent session can serve multiple users, so the invocation id alone is + not a safe identity: two users — or two sessions — reusing an id would collide + on the TaskManager record and let one caller poll or cancel another's job. + Composing the id from ``user_id`` + ``session_id`` + ``invocation_id`` keeps + every start/poll/cancel path isolated. It is also used as the durable + checkpoint item key (and thus the prefix of the cancel-marker key). + + A SHA-256 digest is used (rather than ``f"{user}/{session}/{invocation}"``) + because the provider task-id contract is ``[A-Za-z0-9_-]{1,128}`` (a ``/`` — + and ``.`` or ``:`` — is rejected) and bounded to 128 characters. The hex + digest plus the ``cj-`` prefix uses only ``[a-z0-9-]`` and is a fixed 67 + chars, so it is always valid regardless of how long the protocol ids are. The + ``\\x00`` separators keep the three fields unambiguous. + """ + digest = hashlib.sha256( + f"{user_id}\x00{session_id}\x00{invocation_id}".encode("utf-8") + ).hexdigest() + return f"cj-{digest}" + + +async def open_checkpoint_store(session_id: str, user_id: str) -> FoundryStateStore: + """Open the session-scoped, **user-isolated** checkpoint store. + + A single agent session can serve multiple users, so on top of the + session-scoped store name the store is created with ``user_isolation=True`` + and the explicit ``user_id`` — the platform partitions items per user, so one + user cannot read or cancel another's job even within the same session. Every + start/poll/cancel/recover path opens it the same way (and the task carries + ``user_id`` in its durable input so recovery reopens the same partition). + """ + return await FoundryStateStore.get_or_create( + checkpoint_store_name(session_id), + user_isolation=True, + user_id=user_id or None, + ) + + +@task(name="cancellable_job") +async def cancellable_job(ctx: TaskContext[dict]) -> dict[str, Any]: + """Run ``steps`` steps, checkpointing each, but stop early if cancelled. + + Before every step the task reads the durable cancel marker; if it is present + the task records ``status: "cancelled"`` and returns without finishing the + remaining steps. A crash mid-run resumes from the next step (and still honours + a cancel requested before the crash). + """ + + data = ctx.input or {} + name = str(data.get("name", "world")) + steps = int(data.get("steps", 10)) + # session_id/user_id are carried in the durable input so recovery re-enters + # with the same store scope / user partition as the original run. + session_id = str(data.get("session_id", "")) + user_id = str(data.get("user_id", "")) + cancel_key = f"{ctx.task_id}{CANCEL_SUFFIX}" + + store = await open_checkpoint_store(session_id, user_id) + try: + item = await store.get_item(ctx.task_id) + completed = int((item.value.get("completed_steps", 0) if item else 0) or 0) + status_at_entry = item.value.get("status") if item else None + etag = item.etag if item else None + + if ctx.entry_mode == "recovered": + logger.warning( + "Recovered — resuming '%s' at step %d/%d", name, completed + 1, steps + ) + + # Already finalized (recovered after a terminal write): nothing to do. + if status_at_entry in ("completed", "cancelled"): + return { + "name": name, + "steps": steps, + "completed_steps": completed, + "status": status_at_entry, + } + + try: + done = completed + for i in range(completed, steps): + # ── COOPERATIVE CANCEL CHECK ── + # Durable + cross-replica-safe + crash-durable. Checked BEFORE the + # work so a cancel takes effect within one step interval. + if await store.get_item(cancel_key) is not None: + logger.info( + "cancelled — stopping '%s' at step %d/%d", name, i, steps + ) + await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": i, + "status": "cancelled", + }, + if_match=etag, + ) + return { + "name": name, + "steps": steps, + "completed_steps": i, + "status": "cancelled", + } + + await asyncio.sleep(_STEP_DELAY) # stand-in for long-running work + logger.info("step %d/%d done for %s", i + 1, steps, name) + + # ── CHECKPOINT — the durable crash-recovery boundary ── + ref = await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": i + 1, + "status": "in_progress", + }, + if_match=etag, + ) + etag = ref.etag + done = i + 1 # advance so a later failure records real progress + + logger.info("Finished %d steps for %s", steps, name) + # ── TERMINAL SUCCESS ── + # The one-shot task record is deleted on terminal exit, so the poll + # endpoint has only this durable item to read. Persist an explicit + # terminal status so a finished (or cancelled/failed) run is never + # reported as ``in_progress`` forever. + await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": steps, + "status": "completed", + }, + if_match=etag, + ) + return {"name": name, "steps": steps, "status": "completed"} + except Exception as exc: # noqa: BLE001 — record failure, then re-raise + # Record the *actual* progress (``done``, advanced after each + # successful checkpoint) so failure never rolls ``completed_steps`` + # backward to the entry-time value. + logger.exception("cancellable_job failed for %s", name) + try: + # Keep THIS execution's last-owned ETag. Re-reading the latest + # ETag would defeat the checkpoint CAS: if we lost a race to a + # recovered/new owner, our stale ``done`` must NOT clobber the + # winner's newer progress — the if_match then fails and we skip + # the failure write. + await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": done, + "status": "failed", + "error": str(exc), + }, + if_match=etag, + ) + except Exception: # noqa: BLE001 — never mask the original failure + logger.warning("could not persist failed status for %s", name) + raise + finally: + await store.aclose() + + +__all__ = [ + "cancellable_job", + "checkpoint_store_name", + "durable_task_id", + "open_checkpoint_store", + "CANCEL_SUFFIX", +] diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/app.py new file mode 100644 index 000000000000..3473c5662bd1 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/app.py @@ -0,0 +1,283 @@ +"""Minimal HTTP host for the cancellable resilient job — **start, poll, cancel**. + +- ``POST /invocations`` with body ``{"name": "...", "steps": N}`` — starts the + durable job and returns ``202`` with its ``invocation_id``. +- ``GET /invocations/{invocation_id}`` — JSON snapshot of ``status`` and + ``completed_steps``. Status is one of ``in_progress`` / ``cancelling`` / + ``cancelled`` / ``completed`` / ``failed``. +- ``POST /invocations/{invocation_id}/cancel`` — request cooperative + cancellation. Writes a durable cancel marker; the running job notices it before + its next step and stops with ``status: cancelled``. + +Run it:: + + pip install -r requirements.txt + python app.py + +Then (use the same ``?agent_session_id=`` on every call so poll/cancel hit the +same session-scoped store):: + + # start a 30-step job (STEP_DELAY defaults to 2s, so ~60s of work) + curl -s -XPOST -H "Content-Type: application/json" \\ + -d '{"name": "Ada", "steps": 30}' \\ + "http://localhost:8088/invocations?agent_session_id=demo" + # -> {"status": "started", "invocation_id": "", "total_steps": 30} + + # poll it — completed_steps climbs while status is "in_progress" + curl -s "http://localhost:8088/invocations/?agent_session_id=demo" + # -> {"status": "in_progress", "completed_steps": 4, "total_steps": 30} + + # cancel it mid-run + curl -s -XPOST "http://localhost:8088/invocations//cancel?agent_session_id=demo" + # -> {"status": "cancelling", "invocation_id": ""} + + # poll again — the job stopped early + curl -s "http://localhost:8088/invocations/?agent_session_id=demo" + # -> {"status": "cancelled", "completed_steps": 4, "total_steps": 30} +""" + +from __future__ import annotations + +import json + +from starlette.requests import Request +from starlette.responses import JSONResponse, Response + +from azure.ai.agentserver.core.storage import FoundryStorageConflictError +from azure.ai.agentserver.core.tasks import ( + TaskConflictError, + set_resilient_tasks_enabled, +) +from azure.ai.agentserver.invocations import InvocationAgentServerHost + +try: + from .agent import ( + CANCEL_SUFFIX, + cancellable_job, + durable_task_id, + open_checkpoint_store, + ) +except ImportError: # allows `python app.py` from inside this directory + from agent import ( + CANCEL_SUFFIX, + cancellable_job, + durable_task_id, + open_checkpoint_store, + ) + +# Resilient tasks (durable execution + crash recovery) are strictly opt-in as of +# azure-ai-agentserver-core 2.1.0b1: the host builds the ``TaskManager`` (and runs +# the recovery scan) ONLY when this switch is on, and it must be set before host +# startup (i.e. at module-import time). Without it, ``cancellable_job.start()`` +# raises ``TaskManagerNotInitialized`` and there is no crash recovery. +set_resilient_tasks_enabled(True) + +app = InvocationAgentServerHost() + + +@app.invoke_handler +async def handle_invoke(request: Request) -> Response: + """Start the durable job and return immediately (it runs in the background).""" + try: + data = json.loads(await request.body() or b"{}") + except (json.JSONDecodeError, UnicodeDecodeError): + # UnicodeDecodeError: request bytes are not valid UTF-8. Both are + # malformed-body cases and must be 400, not an internal 500. + return JSONResponse({"error": "invalid JSON body"}, status_code=400) + if not isinstance(data, dict): + return JSONResponse( + {"error": "request body must be a JSON object"}, status_code=400 + ) + name = str(data.get("name", "world")) + # ``steps`` must be a positive integer. Validate the *decoded JSON type* (not + # via ``int(...)``, which would accept ``True`` or truncate ``2.9``). + steps = data.get("steps", 10) + if isinstance(steps, bool) or not isinstance(steps, int): + return JSONResponse({"error": "steps must be an integer"}, status_code=400) + if steps <= 0: + return JSONResponse( + {"error": "steps must be a positive integer"}, status_code=400 + ) + + # Identity: compose the durable task id from user + session + invocation so + # runs cannot collide or be cross-accessed. session_id/user_id are carried in + # the input so recovery reopens the same store scope + user partition; call_id + # is carried so the TaskManager can restore the Foundry call identity for + # hosted store calls after a crash (recovered tasks have no inbound request). + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + call_id: str = request.state.call_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + # The durable record is the authoritative existence gate — NOT start(). A + # one-shot task record is deleted on terminal exit (completed/cancelled/ + # failed), so once a job has ended, ``start()`` would NOT conflict and a reused + # invocation id would spin up a new task against the stale checkpoint/marker. + # So we gate on an atomic ``create_item``: + # * success -> genuinely new job; schedule it. + # * conflict + TERMINAL status -> idempotent; return 409 (never re-run). + # * conflict + NONTERMINAL status -> the record exists but may be an + # ORPHAN (an earlier attempt seeded it then crashed before the task was + # durably scheduled). Fall through and (re-)schedule idempotently — + # ``start()`` recovers the orphan; ``TaskConflictError`` means a job is + # already running. Seeding before ``start()`` also makes the invocation + # visible on every replica immediately (no pre-checkpoint 404 window). + _TERMINAL = ("completed", "cancelled", "failed") + seeded_now = False + store = await open_checkpoint_store(session_id, user_id) + try: + try: + await store.create_item( + task_id, + { + "name": name, + "steps": steps, + "completed_steps": 0, + "status": "in_progress", + }, + ) + seeded_now = True + except FoundryStorageConflictError: + existing = await store.get_item(task_id) + evalue = existing.value if existing else {} + status = evalue.get("status") or "in_progress" + if status in _TERMINAL: + return JSONResponse( + { + "status": status, + "invocation_id": invocation_id, + "detail": "invocation already exists", + }, + status_code=409, + ) + # Nonterminal: fall through to (re-)schedule idempotently. Reuse the + # PERSISTED parameters (name/steps), not this retry's body, so a retry + # reusing the invocation id with a different body cannot mutate or roll + # back the in-flight job. + name = str(evalue.get("name", name)) + steps = int(evalue.get("steps", steps)) + finally: + await store.aclose() + + try: + await cancellable_job.start( + task_id=task_id, + input={ + "name": name, + "steps": steps, + "session_id": session_id, + "user_id": user_id, + "call_id": call_id, + }, + ) + except TaskConflictError: + # A job is already scheduled/running for this id — idempotent. + return JSONResponse( + {"status": "already_started", "invocation_id": invocation_id}, + status_code=409, + ) + except Exception: + # Non-conflict scheduling failure. If we seeded the record this request, + # remove it so a retry can start cleanly rather than being wedged behind a + # permanent ``in_progress`` record with no job to recover. (An ambiguous + # remote accept is still safe: a retry re-seeds and ``start()`` raises + # TaskConflictError, handled as "already running".) + if seeded_now: + try: + cleanup = await open_checkpoint_store(session_id, user_id) + try: + await cleanup.delete_item(task_id) + finally: + await cleanup.aclose() + except Exception: # noqa: BLE001 — best-effort compensation + pass + raise + + return JSONResponse( + {"status": "started", "invocation_id": invocation_id, "total_steps": steps}, + status_code=202, + ) + + +@app.get_invocation_handler +async def handle_get(request: Request) -> Response: + """Return a JSON status snapshot from the durable checkpoint (poll this). + + Status is derived from durable state, so it is correct on any replica and + after a crash. Terminal states (``completed``/``cancelled``/``failed``) are + persisted by the task; ``cancelling`` means a cancel was requested but the + job has not yet reached its next step. A started invocation always has a + durable record (seeded by the invoke handler), so an absent record means a + genuinely unknown invocation. + """ + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + store = await open_checkpoint_store(session_id, user_id) + try: + item = await store.get_item(task_id) + cancel_marker = await store.get_item(f"{task_id}{CANCEL_SUFFIX}") + finally: + await store.aclose() + + if item is None and cancel_marker is None: + return JSONResponse( + {"status": "not_found", "invocation_id": invocation_id}, status_code=404 + ) + + value = item.value if item else {} + done = int(value.get("completed_steps", 0) or 0) + total = int(value.get("steps", 0) or 0) + persisted = value.get("status") + if persisted in ("completed", "cancelled", "failed"): + status = persisted + elif cancel_marker is not None: + # Cancel requested; the job has not yet observed it and finalized. + status = "cancelling" + else: + status = "in_progress" + body = { + "status": status, + "invocation_id": invocation_id, + "completed_steps": done, + "total_steps": total, + } + if status == "failed" and value.get("error"): + body["error"] = value["error"] + return JSONResponse(body) + + +@app.cancel_invocation_handler +async def handle_cancel(request: Request) -> Response: + """Request cooperative cancellation of a running job. + + Writes a durable cancel marker the job checks before each step. Refuses to + persist a marker for an invocation that does not exist (an unknown/arbitrary + id would otherwise poison a later legitimate start with the same session/id). + Existence is decided by the durable record seeded at invoke time (visible on + every replica), not the replica-local ``get_active_run()``. + """ + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + store = await open_checkpoint_store(session_id, user_id) + try: + if await store.get_item(task_id) is None: + return JSONResponse( + {"status": "not_found", "invocation_id": invocation_id}, + status_code=404, + ) + await store.set_item(f"{task_id}{CANCEL_SUFFIX}", {"cancel": True}) + finally: + await store.aclose() + + return JSONResponse({"status": "cancelling", "invocation_id": invocation_id}) + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/requirements.txt b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/requirements.txt new file mode 100644 index 000000000000..c4ace62b69a8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_cancellable/requirements.txt @@ -0,0 +1,9 @@ +# Deliberately minimal: this sample needs ONLY the two agentserver packages — +# no azure-ai-projects, no openai, no langgraph — so it runs with the smallest +# possible dependency and memory footprint. +# +# Constrained to compatible release lines (matching the package's own +# convention): the strictly-opt-in durable-task behaviour and local state store +# this sample relies on require Core 2.1.0b1+. +azure-ai-agentserver-core>=2.1.0b1,<3.0.0 +azure-ai-agentserver-invocations>=1.2.0b1,<2.0.0 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/README.md new file mode 100644 index 000000000000..6010e0feeb16 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/README.md @@ -0,0 +1,90 @@ +# Minimal *indefinite* resilient long-running agent + +A durable background worker that **never finishes on its own** — it ticks +forever until cancelled, surviving crashes and redeploys. It is the companion to +`resilient_hello_world` (which runs a fixed number of steps and completes), and +depends on **only** the two agentserver packages — no LLM, no `azure-ai-projects`, +no `langgraph`. + +## What makes an infinite loop a well-behaved LRA + +The body is a `while True` loop, plus the three things such a loop needs: + +1. **Checkpoint every iteration** — a monotonically increasing `iterations` + cursor is written to a durable state store, so recovery resumes from where it + was. +2. **Graceful shutdown** — on redeploy / SIGTERM the framework sets + `ctx.shutdown`; the loop does `return await ctx.exit_for_recovery()` to release + the lease cleanly so the next instance re-enters and continues. +3. **A stop path** — an explicit cancel writes a durable *stop marker* to the + checkpoint store; the loop re-reads that marker every iteration (independently + of `ctx.cancel`) and returns terminally, so it can be stopped on demand even + when the cancel lands on a different replica than the one running the loop. + +It also sets `timeout=timedelta(days=7)` (the maximum per-turn budget). The +framework's per-turn watchdog is *cooperative*: at the budget it only sets +`ctx.timeout_exceeded` / `ctx.cancel` — it does **not** forcibly end the turn, +and this loop deliberately keeps ticking rather than treating that as a stop. +Re-entry with `ctx.entry_mode == "recovered"` happens on **crash** or +**redeploy** (via `exit_for_recovery` on `ctx.shutdown`), resuming from the +checkpointed iteration; raising the budget to the 7-day maximum just avoids noisy +watchdog signals for a task meant to run indefinitely. + +## Files + +| File | Purpose | +|------|---------| +| `agent.py` | The durable `@task` (`hello_forever`): the infinite, checkpointed loop. | +| `app.py` | `InvocationAgentServerHost`: start / poll / cancel. | +| `requirements.txt` | Just `azure-ai-agentserver-core` + `azure-ai-agentserver-invocations`. | + +## Run it + +```bash +pip install -r requirements.txt +python app.py +``` + +```bash +# start the forever worker (note the invocation_id in the response). Pass the +# same ?agent_session_id= on every call so poll/cancel hit the same store. +curl -s -XPOST -H "Content-Type: application/json" \ + -d '{"name": "Ada"}' \ + "http://localhost:8088/invocations?agent_session_id=demo" +# -> {"status": "started", "invocation_id": ""} + +# poll — iterations keep climbing, status stays "running" +curl -s "http://localhost:8088/invocations/?agent_session_id=demo" +# -> {"status": "running", "iterations": 12} + +# stop it +curl -s -XPOST "http://localhost:8088/invocations//cancel?agent_session_id=demo" +# -> {"status": "cancelling", "invocation_id": ""} +``` + +## The key line: enable resilient tasks + +As of `azure-ai-agentserver-core` **2.1.0b1** the durable-task subsystem is +**strictly opt-in**. Before host startup: + +```python +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled +set_resilient_tasks_enabled(True) +``` + +Without it, `hello_forever.start()` raises `TaskManagerNotInitialized` and there +is no crash recovery. + +## See the recovery + +1. Start the worker and poll until `iterations` is a few in. +2. **Kill the process** (Ctrl-C). +3. Restart `python app.py`. The recovery scan re-enters `hello_forever` with + `ctx.entry_mode == "recovered"` (see the `Recovered '' at iteration N` + log) and the loop continues climbing from `N` — it does **not** reset to 0. + +## Environment + +| Var | Default | Meaning | +|-----|---------|---------| +| `TICK_SECONDS` | `2` | Seconds between iterations. | diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/__init__.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/agent.py new file mode 100644 index 000000000000..ea148747846a --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/agent.py @@ -0,0 +1,223 @@ +"""Minimal *indefinite* resilient long-running agent (invocations protocol). + +Where ``resilient_hello_world`` runs a fixed number of steps and finishes, this +sample **never finishes on its own** — it is a durable background worker that +ticks forever until it is cancelled, surviving crashes and redeploys. + +The differences from the finite sample are exactly the three things an infinite +loop needs to be a well-behaved LRA: + +1. ``while True`` instead of a bounded ``for`` — it always has more work. +2. **Graceful shutdown**: on redeploy / SIGTERM the framework sets + ``ctx.shutdown``; the loop calls ``return await ctx.exit_for_recovery()`` to + release the lease cleanly so the next instance re-enters and continues. +3. **A stop path**: an explicit cancel writes a durable *stop marker* to the + checkpoint store; the loop checks that marker every iteration and returns + terminally, so the worker can actually be stopped on demand — even when the + cancel request lands on a different replica than the one running the loop. + +It also sets ``timeout=timedelta(days=7)`` — the maximum per-turn budget. The +framework's per-turn watchdog is *cooperative*: when the budget is reached it +only sets ``ctx.timeout_exceeded``/``ctx.cancel``; it does not forcibly end the +turn. This loop deliberately does not treat that as a stop, so it simply keeps +ticking. Re-entry with ``ctx.entry_mode == "recovered"`` happens on **crash** or +**redeploy** (via ``exit_for_recovery`` on ``ctx.shutdown``), and the worker +resumes from its checkpointed iteration. Raising the budget to the 7-day maximum +just avoids noisy watchdog signals for a task that is meant to run indefinitely. + +Progress (the ``iterations`` cursor) is checkpointed to a durable, session-scoped +state store (``FoundryStateStore``, which uses a local on-disk backend when +running outside Foundry — so no Azure resources are required to run this locally). + +Input schema: ``{"name": str}``. The host also injects the invocation's +``session_id`` into the durable input so the checkpoint store is isolated per +session (see ``checkpoint_store_name``). + +Environment: + +- ``TICK_SECONDS`` — seconds between iterations (default ``2``). +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +from datetime import timedelta +from typing import Any + +from azure.ai.agentserver.core.storage import FoundryStateStore +from azure.ai.agentserver.core.tasks import TaskContext, task + +logger = logging.getLogger(__name__) + +_TICK = float(os.environ.get("TICK_SECONDS", "2")) + + +def checkpoint_store_name(session_id: str) -> str: + """Return the **session-isolated** checkpoint store name. + + ``FoundryStateStore`` is agent-scoped and has no built-in per-session + isolation, so the store is namespaced by the invocation's session id (as the + other resilient samples do). This keeps one session's worker from being read + or stopped by another — important because POST accepts a caller-supplied + invocation id that a different session could otherwise reuse as a key. Shared + with app.py so the poll/cancel endpoints read the same scope. + + The session component is **hashed**: a protocol session id can be up to 256 + characters, and the local ``FoundryStateStore`` backend base64-encodes the + whole store name into a single filename, which would blow past the 255-byte + ``NAME_MAX`` and fail this no-cloud sample with ``ENAMETOOLONG``. A fixed-width + SHA-256 digest keeps the name bounded while remaining unique per session. + """ + digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest() + return f"resilient-hello-forever/{digest}" + + +def durable_task_id(session_id: str, invocation_id: str, user_id: str) -> str: + """Return the TaskManager task id derived from the user, session and invocation. + + The invocations protocol accepts a *caller-supplied* invocation id, and a + single agent session can serve multiple users, so the invocation id alone (or + even session+invocation) is not a safe identity: two users — or two sessions — + reusing an id would collide on the TaskManager record and let one caller poll + or cancel another's worker. Composing the id from ``user_id`` + ``session_id`` + + ``invocation_id`` keeps every start/poll/cancel path isolated. It is also + used as the durable checkpoint item key (and thus the prefix of the + stop-marker key). + + A SHA-256 digest is used (rather than ``f"{user}/{session}/{invocation}"``) + for two reasons: the provider task-id contract is ``[A-Za-z0-9_-]{1,128}`` + (a ``/`` — and ``.`` or ``:`` — is rejected), and it is bounded to 128 + characters. The hex digest plus the ``hf-`` prefix uses only ``[a-z0-9-]`` + and is a fixed 67 chars, so it is always valid regardless of how long the + protocol ids are. The ``\\x00`` separators keep the three fields unambiguous. + """ + digest = hashlib.sha256( + f"{user_id}\x00{session_id}\x00{invocation_id}".encode("utf-8") + ).hexdigest() + return f"hf-{digest}" + + +async def open_checkpoint_store(session_id: str, user_id: str) -> FoundryStateStore: + """Open the session-scoped, **user-isolated**, non-expiring checkpoint store. + + A single agent session can serve multiple users, so on top of the + session-scoped store name the store is created with ``user_isolation=True`` + and the explicit ``user_id`` — the platform partitions items per user, so one + user cannot read or stop another's worker even within the same session. + ``item_ttl_seconds=-1`` keeps the checkpoint and stop marker from expiring for + an indefinitely-running worker. Every start/poll/cancel/recover path opens it + the same way (and the worker carries ``user_id`` in its durable input so + recovery reopens the same partition). + """ + return await FoundryStateStore.get_or_create( + checkpoint_store_name(session_id), + user_isolation=True, + user_id=user_id or None, + item_ttl_seconds=-1, + ) + + +# Suffix for the durable "stop" marker key. The cancel endpoint (app.py) writes +# this key; the worker checks it to decide whether to stop. Keeping it in a +# SEPARATE key means the cancel write never races the checkpoint's ETag. +STOP_SUFFIX = "/stop" + + +@task(name="hello_forever", timeout=timedelta(days=7)) +async def hello_forever(ctx: TaskContext[dict]) -> dict[str, Any]: + """Tick forever, checkpointing the ``iterations`` cursor after each tick. + + Runs until an explicit cancel. Survives crashes (resumes from the checkpoint) + and redeploys (yields cleanly via ``exit_for_recovery`` and resumes on the + next instance). + """ + + name = str((ctx.input or {}).get("name", "world")) + # ``session_id`` and ``user_id`` are carried in the durable input so recovery + # re-enters with the same store scope / user partition as the original run. + session_id = str((ctx.input or {}).get("session_id", "")) + user_id = str((ctx.input or {}).get("user_id", "")) + stop_key = f"{ctx.task_id}{STOP_SUFFIX}" + + store = await open_checkpoint_store(session_id, user_id) + try: + item = await store.get_item(ctx.task_id) + n = int((item.value.get("iterations", 0) if item else 0) or 0) + etag = item.etag if item else None + + if ctx.entry_mode == "recovered": + logger.warning("Recovered '%s' at iteration %d", name, n) + + try: + while True: + # 1) Graceful redeploy / SIGTERM: release the lease so the + # recovery scan re-enters this worker on the next instance. + if ctx.shutdown.is_set(): + logger.info( + "shutdown — yielding for recovery at iteration %d", n + ) + return await ctx.exit_for_recovery() + + # 2) Explicit stop. The DURABLE stop marker is the single source + # of truth and is checked EVERY iteration. The cancel endpoint + # writes the marker (it does NOT rely on an in-process signal), + # so a stop is observed here regardless of which replica + # received the cancel request. The marker also survives the + # per-turn watchdog and crash/recovery, so it never races the + # checkpoint's ETag. + if await store.get_item(stop_key) is not None: + logger.info("stop requested — stopping at iteration %d", n) + return {"name": name, "iterations": n, "stopped": True} + + # 3) One unit of ongoing work, then CHECKPOINT the durable cursor. + n += 1 + logger.info("iteration %d for %s", n, name) + ref = await store.set_item( + ctx.task_id, + {"name": name, "iterations": n, "status": "running"}, + if_match=etag, + ) + etag = ref.etag + + await asyncio.sleep(_TICK) # heartbeat / interval between ticks; + # the stop marker is re-checked at the top of every iteration, so + # an explicit stop is honoured within one tick regardless of which + # replica received the cancel. + except Exception as exc: # noqa: BLE001 — record failure, then re-raise + # A raised exception is terminal (the one-shot record is deleted), so + # without this the durable checkpoint would keep reporting ``running`` + # forever. + logger.exception("hello_forever failed for %s", name) + try: + # Keep THIS execution's last-owned ETag. Re-reading the latest + # ETag would defeat the checkpoint CAS: if we lost a race to a + # recovered/new owner, our stale ``n`` must NOT clobber the + # winner's newer progress — the if_match then fails and we skip + # the failure write. + await store.set_item( + ctx.task_id, + { + "name": name, + "iterations": n, + "status": "failed", + "error": str(exc), + }, + if_match=etag, + ) + except Exception: # noqa: BLE001 — never mask the original failure + logger.warning("could not persist failed status for %s", name) + raise + finally: + await store.aclose() + + +__all__ = [ + "hello_forever", + "checkpoint_store_name", + "durable_task_id", + "open_checkpoint_store", + "STOP_SUFFIX", +] diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/app.py new file mode 100644 index 000000000000..b8c2041fc6f9 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/app.py @@ -0,0 +1,274 @@ +"""Minimal HTTP host for the *indefinite* "hello forever" resilient agent. + +No streaming — **start, poll, and cancel**: + +- ``POST /invocations`` with body ``{"name": "..."}`` — starts the durable + worker and returns ``202`` with its ``invocation_id``. The worker runs forever + in the background until cancelled. +- ``GET /invocations/{invocation_id}`` — JSON snapshot: whether the worker is + ``running`` and its current ``iterations`` count. +- ``POST /invocations/{invocation_id}/cancel`` — stop the worker. + +Run it:: + + pip install -r requirements.txt + python app.py + +Then (use the invocation_id from the POST response / X-Agent-Invocation-Id, and +the same ``?agent_session_id=`` on every call so poll/cancel hit the same +session-scoped store):: + + # start the forever worker + curl -s -XPOST -H "Content-Type: application/json" \\ + -d '{"name": "Ada"}' \\ + "http://localhost:8088/invocations?agent_session_id=demo" + # -> {"status": "started", "invocation_id": ""} + + # poll it — iterations keep climbing, status stays "running" + curl -s "http://localhost:8088/invocations/?agent_session_id=demo" + # -> {"status": "running", "iterations": 12} + + # stop it + curl -s -XPOST "http://localhost:8088/invocations//cancel?agent_session_id=demo" + # -> {"status": "cancelling", "invocation_id": ""} +""" + +from __future__ import annotations + +import json + +from starlette.requests import Request +from starlette.responses import JSONResponse, Response + +from azure.ai.agentserver.core.storage import FoundryStorageConflictError +from azure.ai.agentserver.core.tasks import ( + TaskConflictError, + set_resilient_tasks_enabled, +) +from azure.ai.agentserver.invocations import InvocationAgentServerHost + +try: + from .agent import ( + STOP_SUFFIX, + durable_task_id, + hello_forever, + open_checkpoint_store, + ) +except ImportError: # allows `python app.py` from inside this directory + from agent import ( + STOP_SUFFIX, + durable_task_id, + hello_forever, + open_checkpoint_store, + ) + +# Resilient tasks (durable execution + crash recovery) are strictly opt-in as of +# azure-ai-agentserver-core 2.1.0b1: ``AgentServerHost`` builds the +# ``TaskManager`` ONLY when this switch is on, and it must be set before host +# startup (i.e. at module-import time). Without it, ``hello_forever.start()`` +# raises ``TaskManagerNotInitialized`` and there is no crash recovery — which +# would defeat the purpose of a long-running agent. +set_resilient_tasks_enabled(True) + +app = InvocationAgentServerHost() + + +@app.invoke_handler +async def handle_invoke(request: Request) -> Response: + """Start the forever worker and return immediately (it runs in the background).""" + try: + data = json.loads(await request.body() or b"{}") + except (json.JSONDecodeError, UnicodeDecodeError): + # UnicodeDecodeError: request bytes are not valid UTF-8. Both are + # malformed-body cases and must be 400, not an internal 500. + return JSONResponse({"error": "invalid JSON body"}, status_code=400) + if not isinstance(data, dict): + return JSONResponse( + {"error": "request body must be a JSON object"}, status_code=400 + ) + name = str(data.get("name", "world")) + + # Identity for this worker. The invocation id is caller-supplied and a session + # may serve multiple users, so compose the durable task id from user + session + # + invocation (see durable_task_id) so workers cannot collide or be + # cross-accessed. session_id/user_id are carried in the input so recovery + # reopens the same store scope + user partition; call_id is carried so the + # TaskManager can restore the Foundry call identity for hosted store calls + # after a crash (recovered tasks have no inbound request). + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + call_id: str = request.state.call_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + # The durable record is the authoritative existence gate — NOT start(). A + # one-shot task record is deleted on terminal exit (stop/failure), so once a + # worker has stopped, ``start()`` would NOT conflict and a reused invocation id + # would spin up a brand-new task that immediately sees the OLD stop marker and + # exits — while POST wrongly reported "started". So we gate on an atomic + # ``create_item``: + # * success -> genuinely new worker; schedule it. + # * conflict + TERMINAL (stop marker present, or status == "failed") -> + # idempotent; return 409 (never re-run). + # * conflict + NONTERMINAL ("running") -> the record exists but may be an + # ORPHAN (an earlier attempt seeded it then crashed before the worker was + # durably scheduled). Fall through and (re-)schedule idempotently — + # ``start()`` recovers the orphan; ``TaskConflictError`` means a worker is + # already running. Seeding before ``start()`` also makes the invocation + # visible on every replica immediately (no pre-checkpoint 404 window). + seeded_now = False + store = await open_checkpoint_store(session_id, user_id) + try: + try: + await store.create_item( + task_id, {"name": name, "iterations": 0, "status": "running"} + ) + seeded_now = True + except FoundryStorageConflictError: + existing = await store.get_item(task_id) + stop_marker = await store.get_item(f"{task_id}{STOP_SUFFIX}") + evalue = existing.value if existing else {} + status = evalue.get("status") or "running" + if stop_marker is not None: + status = "stopped" + if stop_marker is not None or status == "failed": + return JSONResponse( + { + "status": status, + "invocation_id": invocation_id, + "detail": "invocation already exists", + }, + status_code=409, + ) + # Nonterminal ("running"): fall through to (re-)schedule idempotently. + # Reuse the PERSISTED name, not this retry's body, so an orphan retry + # cannot mutate the identity recorded by the original invocation. + name = str(evalue.get("name", name)) + finally: + await store.aclose() + + try: + await hello_forever.start( + task_id=task_id, + input={ + "name": name, + "session_id": session_id, + "user_id": user_id, + "call_id": call_id, + }, + ) + except TaskConflictError: + # A worker is already scheduled/running for this id — idempotent. + return JSONResponse( + {"status": "already_running", "invocation_id": invocation_id}, + status_code=409, + ) + except Exception: + # Non-conflict scheduling failure. If we seeded the record this request, + # remove it so a retry can start cleanly rather than being wedged behind a + # permanent ``running`` record with no worker to recover. (An ambiguous + # remote accept is still safe: a retry re-seeds and ``start()`` raises + # TaskConflictError, handled as "already running".) + if seeded_now: + try: + cleanup = await open_checkpoint_store(session_id, user_id) + try: + await cleanup.delete_item(task_id) + finally: + await cleanup.aclose() + except Exception: # noqa: BLE001 — best-effort compensation + pass + raise + + return JSONResponse( + {"status": "started", "invocation_id": invocation_id}, status_code=202 + ) + + +@app.get_invocation_handler +async def handle_get(request: Request) -> Response: + """Report whether the worker is running, stopped, or failed. + + Status is derived from **durable state**, not process-local run ownership: an + indefinite worker keeps running (possibly on a *different* replica) until it + is explicitly stopped, so ``get_active_run()`` returning ``None`` on the + polled replica means nothing. The invoke handler seeds an initial durable + record, so a started invocation is visible on every replica immediately (no + pre-checkpoint 404 window). Terminal state is durable: the stop marker means + ``stopped``; the worker persists ``status: failed`` on an unhandled exception + (its one-shot record is deleted on terminal exit, so without that a crashed + worker would report ``running`` forever). Absence of both the record and the + marker means the invocation is genuinely unknown. + """ + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + store = await open_checkpoint_store(session_id, user_id) + try: + item = await store.get_item(task_id) + stop_marker = await store.get_item(f"{task_id}{STOP_SUFFIX}") + finally: + await store.aclose() + + if item is None and stop_marker is None: + # No durable record and no stop marker: genuinely unknown invocation. + # (A started invocation always has a record from the invoke handler.) + return JSONResponse( + {"status": "not_found", "invocation_id": invocation_id}, + status_code=404, + ) + + value = item.value if item else {} + iterations = int(value.get("iterations", 0) or 0) + if stop_marker is not None: + status = "stopped" + elif value.get("status") == "failed": + status = "failed" + else: + status = "running" + body = { + "status": status, + "invocation_id": invocation_id, + "iterations": iterations, + } + if status == "failed" and value.get("error"): + body["error"] = value["error"] + return JSONResponse(body) + + +@app.cancel_invocation_handler +async def handle_cancel(request: Request) -> Response: + """Stop the forever worker. + + Refuses to persist a stop marker for an invocation that does not exist: + otherwise cancelling an arbitrary caller-chosen id would make GET report it + ``stopped`` and would poison a *later* legitimate start with the same + session/id (the worker would see the stale marker and exit immediately). + Existence is decided by the **durable record** seeded at invoke time (visible + on every replica), not the replica-local ``get_active_run()``. The worker + re-reads the stop marker at the top of every iteration, so it stops within one + tick without any in-process wake signal. + """ + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + store = await open_checkpoint_store(session_id, user_id) + try: + if await store.get_item(task_id) is None: + return JSONResponse( + {"status": "not_found", "invocation_id": invocation_id}, + status_code=404, + ) + await store.set_item(f"{task_id}{STOP_SUFFIX}", {"stop": True}) + finally: + await store.aclose() + + return JSONResponse({"status": "cancelling", "invocation_id": invocation_id}) + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/requirements.txt b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/requirements.txt new file mode 100644 index 000000000000..c4ace62b69a8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_forever/requirements.txt @@ -0,0 +1,9 @@ +# Deliberately minimal: this sample needs ONLY the two agentserver packages — +# no azure-ai-projects, no openai, no langgraph — so it runs with the smallest +# possible dependency and memory footprint. +# +# Constrained to compatible release lines (matching the package's own +# convention): the strictly-opt-in durable-task behaviour and local state store +# this sample relies on require Core 2.1.0b1+. +azure-ai-agentserver-core>=2.1.0b1,<3.0.0 +azure-ai-agentserver-invocations>=1.2.0b1,<2.0.0 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/README.md b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/README.md new file mode 100644 index 000000000000..75da3f111df6 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/README.md @@ -0,0 +1,94 @@ +# Minimal resilient "hello world" long-running agent + +The smallest possible **long-running agent (LRA)** on the invocations protocol — +**start-and-poll, no streaming**. It depends on **only** the two agentserver +packages (no LLM, no `azure-ai-projects`, no `langgraph`), so it runs with a +minimal dependency and memory footprint while still showing the one thing that +makes an agent an LRA: **it survives a crash and resumes from a checkpoint +instead of starting over.** + +## Files + +| File | Purpose | +|------|---------| +| `agent.py` | The durable `@task` (`hello_world`): counts N steps, checkpointing after each. | +| `app.py` | `InvocationAgentServerHost`: `POST` to start, `GET` to poll status. | +| `requirements.txt` | Just `azure-ai-agentserver-core` + `azure-ai-agentserver-invocations`. | + +## Run it + +```bash +pip install -r requirements.txt +python app.py # listens on http://localhost:8088 +``` + +Start a run (returns immediately — the task keeps running in the background). +Pass `?agent_session_id=` to isolate this run's checkpoints; use the **same** +session id when you poll: + +```bash +curl -s -XPOST -H "Content-Type: application/json" \ + -d '{"name": "Ada", "steps": 10}' \ + "http://localhost:8088/invocations?agent_session_id=demo" +# -> {"status": "started", "invocation_id": "", "total_steps": 10} +``` + +Poll it by that invocation id (repeat every couple of seconds, same session): + +```bash +curl -s "http://localhost:8088/invocations/?agent_session_id=demo" +# -> {"status": "in_progress", "completed_steps": 3, "total_steps": 10} +# ... while running you see completed_steps climb. +# When it finishes: +# -> {"status": "completed", "completed_steps": 10, "total_steps": 10} +``` + +> The one-shot `@task` record is cleaned up on completion, but this sample's +> durable **checkpoint** persists, so the poll keeps reporting `completed`. + +## How an *invocation* becomes *long-running* + +1. `POST /invocations` derives a durable task id — `durable_task_id(session_id, + invocation_id, user_id)`, a SHA-256 digest, **not** the raw invocation id — + and calls `hello_world.start(task_id=)`, which **schedules the + task on the TaskManager and returns immediately**. The work is *not* tied to + the HTTP request — the handler returns `202` while the task runs on. (The + digest keeps the id within the provider's `[A-Za-z0-9_-]{1,128}` contract and + isolates it per user + session; it is also the durable **checkpoint** key.) +2. The task loops, sleeping between steps and writing a durable **checkpoint** + to the state store after each, keyed by that same derived task id. +3. `GET /invocations/{invocation_id}` re-derives the same task id from the + request's identity and reads that durable state, so any client can poll + progress long after the original call returned. When correlating logs or + checkpoints, note they are keyed by the derived task id, not the invocation + id. + +## The key line: enable resilient tasks + +As of `azure-ai-agentserver-core` **2.1.0b1** the durable-task subsystem is +**strictly opt-in**. `AgentServerHost` builds the `TaskManager` (and runs the +crash-recovery scan) **only** when you call, before host startup: + +```python +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled +set_resilient_tasks_enabled(True) +``` + +Without it, `hello_world.start()` raises `TaskManagerNotInitialized` and there is +no crash recovery. Declaring a `@task` alone does **not** turn it on. + +## See the recovery + +1. Start a longer run: `{"name": "Ada", "steps": 30}`. +2. Poll until `completed_steps` is a few in, then **kill the process** (Ctrl-C). +3. Restart `python app.py`. On startup the recovery scan re-enters `hello_world` + with `ctx.entry_mode == "recovered"` (see the `Recovered — resuming …` log), + reads `completed_steps` from the durable checkpoint, and continues from the + next step — it does **not** restart from step 1. Poll again to watch + `completed_steps` climb past where it crashed, ending at `completed`. + +## Environment + +| Var | Default | Meaning | +|-----|---------|---------| +| `STEP_DELAY` | `2` | Seconds between steps. Keep > 0 so a crash lands mid-run. | diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/__init__.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/agent.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/agent.py new file mode 100644 index 000000000000..63a9fb2b44e0 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/agent.py @@ -0,0 +1,215 @@ +"""Minimal "hello world" resilient long-running agent (invocations protocol). + +The smallest possible long-running agent (LRA): it depends on ONLY +``azure-ai-agentserver-core`` and ``azure-ai-agentserver-invocations`` — no LLM, +no streaming, no cloud services — so it runs with a minimal dependency and +memory footprint. + +It counts through ``steps`` steps, sleeping between them to simulate slow work +that outlives a single request, and **checkpoints its progress after every +step** to a durable state store (``FoundryStateStore``, which uses a local +on-disk backend when running outside Foundry — so no Azure resources are +required to run this locally). If the container crashes mid-run, the platform +restarts it and the framework re-enters this task with +``ctx.entry_mode == "recovered"`` — the handler reads ``completed_steps`` from +the checkpoint and resumes at the next step instead of starting over. + +Input schema: ``{"name": str, "steps": int?}``. The host also injects the +invocation's ``session_id`` into the durable input so the checkpoint store is +isolated per session (see ``checkpoint_store_name``). + +Environment: + +- ``STEP_DELAY`` — seconds to sleep between steps (default ``2``). Keep it + nonzero so a crash demo has time to land mid-run. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +import os +from typing import Any + +from azure.ai.agentserver.core.storage import FoundryStateStore +from azure.ai.agentserver.core.tasks import TaskContext, task + +logger = logging.getLogger(__name__) + +_STEP_DELAY = float(os.environ.get("STEP_DELAY", "2")) + + +def checkpoint_store_name(session_id: str) -> str: + """Return the **session-isolated** checkpoint store name. + + ``FoundryStateStore`` is agent-scoped and has no built-in per-session + isolation, so the store is namespaced by the invocation's session id (as the + other resilient samples do). This keeps one session's progress from being + read or overwritten by another — important because POST accepts a + caller-supplied invocation id that a different session could otherwise reuse + as a key. Shared with app.py so the poll endpoint reads the same scope. + + The session component is **hashed**: a protocol session id can be up to 256 + characters, and the local ``FoundryStateStore`` backend base64-encodes the + whole store name into a single filename, which would blow past the 255-byte + ``NAME_MAX`` and fail this no-cloud sample with ``ENAMETOOLONG``. A fixed-width + SHA-256 digest keeps the name bounded while remaining unique per session. + """ + digest = hashlib.sha256(session_id.encode("utf-8")).hexdigest() + return f"resilient-hello-world/{digest}" + + +def durable_task_id(session_id: str, invocation_id: str, user_id: str) -> str: + """Return the TaskManager task id derived from the user, session and invocation. + + The invocations protocol accepts a *caller-supplied* invocation id, and a + single agent session can serve multiple users, so the invocation id alone (or + even session+invocation) is not a safe identity: two users — or two sessions — + reusing an id would collide on the TaskManager record (a ``start()`` 500) and + let one caller poll or cancel another's run. Composing the id from + ``user_id`` + ``session_id`` + ``invocation_id`` keeps every start/poll/cancel + path isolated. It is also used as the durable checkpoint item key. + + A SHA-256 digest is used (rather than ``f"{user}/{session}/{invocation}"``) + for two reasons: the provider task-id contract is ``[A-Za-z0-9_-]{1,128}`` + (a ``/`` — and ``.`` or ``:`` — is rejected), and it is bounded to 128 + characters. The hex digest plus the ``hw-`` prefix uses only ``[a-z0-9-]`` + and is a fixed 67 chars, so it is always valid regardless of how long the + protocol ids are. The ``\\x00`` separators keep the three fields unambiguous. + """ + digest = hashlib.sha256( + f"{user_id}\x00{session_id}\x00{invocation_id}".encode("utf-8") + ).hexdigest() + return f"hw-{digest}" + + +async def open_checkpoint_store(session_id: str, user_id: str) -> FoundryStateStore: + """Open the session-scoped, **user-isolated** checkpoint store. + + A single agent session can serve multiple users, so on top of the + session-scoped store name the store is created with ``user_isolation=True`` + and the explicit ``user_id`` — the platform then partitions items per user, so + one user cannot read or overwrite another's checkpoint even within the same + session. Every start/poll/recover path opens it the same way (and the task + carries ``user_id`` in its durable input so recovery reopens the same + partition). + """ + return await FoundryStateStore.get_or_create( + checkpoint_store_name(session_id), + user_isolation=True, + user_id=user_id or None, + ) + + +@task(name="hello_world") +async def hello_world(ctx: TaskContext[dict]) -> dict[str, Any]: + """Count through ``steps`` steps, checkpointing after each one. + + The checkpoint (``completed_steps``) lives in a durable, session-scoped, + user-isolated state store keyed by ``ctx.task_id``, so a crash mid-run resumes + from the next step. + """ + + data = ctx.input or {} + name = str(data.get("name", "world")) + steps = int(data.get("steps", 10)) + # ``session_id`` and ``user_id`` are carried in the durable input so recovery + # re-enters with the same store scope / user partition as the original run. + session_id = str(data.get("session_id", "")) + user_id = str(data.get("user_id", "")) + + store = await open_checkpoint_store(session_id, user_id) + try: + item = await store.get_item(ctx.task_id) + completed = int((item.value.get("completed_steps", 0) if item else 0) or 0) + status_at_entry = (item.value.get("status") if item else None) + etag = item.etag if item else None + + if ctx.entry_mode == "recovered": + logger.warning( + "Recovered — resuming '%s' at step %d/%d", name, completed + 1, steps + ) + + # Already finalized (e.g. recovered after the terminal write): nothing to + # do, and re-writing would needlessly churn the ETag. + if completed >= steps and status_at_entry == "completed": + return {"name": name, "steps": steps, "status": "complete"} + + try: + done = completed + for i in range(completed, steps): + await asyncio.sleep(_STEP_DELAY) # stand-in for long-running work + logger.info("step %d/%d done for %s", i + 1, steps, name) + + # ── CHECKPOINT — the durable crash-recovery boundary ── + # After this write, a crash resumes at step (i + 2), not step 1. + ref = await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": i + 1, + "status": "in_progress", + }, + if_match=etag, + ) + etag = ref.etag + done = i + 1 # advance so a later failure records real progress + + logger.info("Finished %d steps for %s", steps, name) + # ── TERMINAL SUCCESS ── + # The one-shot task record is deleted on terminal exit, so the poll + # endpoint has only this durable item to read. Persist an explicit + # terminal status so a completed (or failed) run is not indefinitely + # reported as ``in_progress``. + await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": steps, + "status": "completed", + }, + if_match=etag, + ) + return {"name": name, "steps": steps, "status": "complete"} + except Exception as exc: # noqa: BLE001 — record failure, then re-raise + # A raised exception is terminal (the record is deleted), so without + # this the durable item would sit below ``steps`` forever and polling + # would report ``in_progress`` indefinitely. Record the *actual* + # progress (``done``, advanced after each successful checkpoint) so + # failure never rolls ``completed_steps`` backward. Best-effort write; + # if the ETag moved we still re-raise so the framework marks the task + # failed. + logger.exception("hello_world failed for %s", name) + try: + # Keep THIS execution's last-owned ETag. Re-reading the latest + # ETag would defeat the checkpoint CAS: if we lost a race to a + # recovered/new owner, our stale ``done`` must NOT clobber the + # winner's newer progress — the if_match then fails and we simply + # skip the failure write. + await store.set_item( + ctx.task_id, + { + "name": name, + "steps": steps, + "completed_steps": done, + "status": "failed", + "error": str(exc), + }, + if_match=etag, + ) + except Exception: # noqa: BLE001 — never mask the original failure + logger.warning("could not persist failed status for %s", name) + raise + finally: + await store.aclose() + + +__all__ = [ + "hello_world", + "checkpoint_store_name", + "durable_task_id", + "open_checkpoint_store", +] diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/app.py new file mode 100644 index 000000000000..c95e7ee28a40 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/app.py @@ -0,0 +1,243 @@ +"""Minimal HTTP host for the "hello world" resilient long-running agent. + +No streaming — just **start and poll**: + +- ``POST /invocations`` with body ``{"name": "...", "steps": N}`` — starts the + durable task and returns ``202`` with its ``invocation_id``. The task keeps + running in the background after this response returns. +- ``GET /invocations/{invocation_id}`` — returns a JSON snapshot of the run's + ``status`` and ``completed_steps`` (read from the durable checkpoint). Poll it + to watch progress. + +Run it:: + + pip install -r requirements.txt + python app.py + +Then, in another shell (pass ``?agent_session_id=`` to isolate a run's +checkpoints; GET/cancel must use the same session):: + + # start a run — note the invocation_id in the response (also the + # X-Agent-Invocation-Id response header) + curl -s -XPOST -H "Content-Type: application/json" \\ + -d '{"name": "Ada", "steps": 10}' \\ + "http://localhost:8088/invocations?agent_session_id=demo" + # -> {"status": "started", "invocation_id": "", "total_steps": 10} + + # poll it (repeat every couple seconds) — same session id + curl -s "http://localhost:8088/invocations/?agent_session_id=demo" + # -> {"status": "in_progress", "completed_steps": 3, "total_steps": 10} + # ... eventually -> {"status": "completed", "completed_steps": 10, ...} +""" + +from __future__ import annotations + +import json + +from starlette.requests import Request +from starlette.responses import JSONResponse, Response + +from azure.ai.agentserver.core.storage import FoundryStorageConflictError +from azure.ai.agentserver.core.tasks import ( + TaskConflictError, + set_resilient_tasks_enabled, +) +from azure.ai.agentserver.invocations import InvocationAgentServerHost + +try: + from .agent import durable_task_id, hello_world, open_checkpoint_store +except ImportError: # allows `python app.py` from inside this directory + from agent import durable_task_id, hello_world, open_checkpoint_store + +# Resilient tasks (durable execution + crash recovery) are strictly opt-in as of +# azure-ai-agentserver-core 2.1.0b1: ``AgentServerHost`` builds the +# ``TaskManager`` ONLY when this switch is on, and it must be set before host +# startup (i.e. at module-import time). Without it, ``hello_world.start()`` +# raises ``TaskManagerNotInitialized`` and there is no crash recovery — which +# would defeat the purpose of a long-running agent. +set_resilient_tasks_enabled(True) + +app = InvocationAgentServerHost() + + +@app.invoke_handler +async def handle_invoke(request: Request) -> Response: + """Start the durable task and return immediately (it runs in the background).""" + try: + data = json.loads(await request.body() or b"{}") + except (json.JSONDecodeError, UnicodeDecodeError): + # UnicodeDecodeError: request bytes are not valid UTF-8. Both are + # malformed-body cases and must be 400, not an internal 500. + return JSONResponse({"error": "invalid JSON body"}, status_code=400) + if not isinstance(data, dict): + return JSONResponse( + {"error": "request body must be a JSON object"}, status_code=400 + ) + name = str(data.get("name", "world")) + # ``steps`` must be a positive integer. Validate the *decoded JSON type* (not + # via ``int(...)``, which would accept ``True`` or silently truncate ``2.9``). + # A non-positive count writes no checkpoint, so the accepted invocation would + # then poll as 404. + steps = data.get("steps", 10) + if isinstance(steps, bool) or not isinstance(steps, int): + return JSONResponse( + {"error": "steps must be an integer"}, status_code=400 + ) + if steps <= 0: + return JSONResponse( + {"error": "steps must be a positive integer"}, status_code=400 + ) + + # Identity for this run. The invocation id is caller-supplied and a session + # may serve multiple users, so compose the durable task id from user + + # session + invocation (see durable_task_id) so runs cannot collide or be + # cross-accessed. session_id/user_id are carried in the input so recovery + # reopens the same store scope + user partition; call_id is carried so the + # TaskManager can restore the Foundry call identity for hosted store calls + # after a crash (recovered tasks have no inbound request). + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + call_id: str = request.state.call_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + # The durable checkpoint is the authoritative existence record — NOT start(). + # A one-shot task record is deleted on terminal exit, so once a run has + # completed/failed, ``start()`` would NOT conflict and a reused invocation id + # would spin up a brand-new task against the old checkpoint. So we gate on an + # atomic ``create_item``: + # * success -> genuinely new invocation; schedule the task. + # * conflict + TERMINAL status -> idempotent; return 409 (never re-run). + # * conflict + NONTERMINAL status -> the record exists but may be an + # ORPHAN: an earlier attempt created the seed then crashed before the + # TaskManager record became durable, so nothing is left for recovery to + # run. Fall through and (re-)schedule idempotently — ``start()`` recovers + # the orphan, and ``TaskConflictError`` means a task is already running. + # Seeding before ``start()`` also makes the invocation visible on every + # replica immediately (no pre-checkpoint 404 window). + _TERMINAL = ("completed", "failed") + seeded_now = False + store = await open_checkpoint_store(session_id, user_id) + try: + try: + await store.create_item( + task_id, + { + "name": name, + "steps": steps, + "completed_steps": 0, + "status": "in_progress", + }, + ) + seeded_now = True + except FoundryStorageConflictError: + existing = await store.get_item(task_id) + evalue = existing.value if existing else {} + status = evalue.get("status") or "in_progress" + if status in _TERMINAL: + return JSONResponse( + { + "status": status, + "invocation_id": invocation_id, + "detail": "invocation already exists", + }, + status_code=409, + ) + # Nonterminal: fall through to (re-)schedule idempotently. Reuse the + # PERSISTED parameters (name/steps), not this retry's body, so a retry + # that reuses the invocation id with a different body cannot mutate or + # roll back the in-flight job (e.g. a 5/10 checkpoint retried with + # steps=3 must not rewrite progress backward). + name = str(evalue.get("name", name)) + steps = int(evalue.get("steps", steps)) + finally: + await store.aclose() + + # start() schedules the task on the TaskManager and returns right away — the + # work is NOT tied to this HTTP request's lifetime. + try: + await hello_world.start( + task_id=task_id, + input={ + "name": name, + "steps": steps, + "session_id": session_id, + "user_id": user_id, + "call_id": call_id, + }, + ) + except TaskConflictError: + # A task is already scheduled/running for this id — idempotent. + return JSONResponse( + {"status": "already_started", "invocation_id": invocation_id}, + status_code=409, + ) + except Exception: + # Non-conflict scheduling failure. If we created the seed on THIS request, + # remove it so a retry can start cleanly rather than being wedged behind a + # permanent ``in_progress`` record with no task to recover. (If the task + # was ambiguously accepted remotely, a retry re-seeds and ``start()`` + # raises TaskConflictError, which we handle as "already running".) + if seeded_now: + try: + cleanup = await open_checkpoint_store(session_id, user_id) + try: + await cleanup.delete_item(task_id) + finally: + await cleanup.aclose() + except Exception: # noqa: BLE001 — best-effort compensation + pass + raise + + return JSONResponse( + {"status": "started", "invocation_id": invocation_id, "total_steps": steps}, + status_code=202, + ) + + +@app.get_invocation_handler +async def handle_get(request: Request) -> Response: + """Return a JSON status snapshot from the durable checkpoint (poll this).""" + # The invocations protocol addresses the run by the {invocation_id} path + # segment; the durable task id is composed from user + session + invocation so + # a poll only ever resolves this user's own run in this session. + invocation_id: str = request.state.invocation_id + session_id: str = request.state.session_id + user_id: str = request.state.user_id + task_id = durable_task_id(session_id, invocation_id, user_id) + + store = await open_checkpoint_store(session_id, user_id) + try: + item = await store.get_item(task_id) + finally: + await store.aclose() + + if item is None: + # A started invocation always has a durable record (seeded by the invoke + # handler), so an absent record means a genuinely unknown invocation. + return JSONResponse( + {"status": "not_found", "invocation_id": invocation_id}, status_code=404 + ) + + value = item.value or {} + done = int(value.get("completed_steps", 0) or 0) + total = int(value.get("steps", 0) or 0) + # Read the explicit terminal status the task persists. The one-shot task + # record is deleted on terminal exit, so this durable field — not + # ``completed_steps >= total`` — is what distinguishes a still-running task + # from a completed or failed one. A task that raised writes ``failed``; without + # that, a run that died below ``total`` would report ``in_progress`` forever. + status = str(value.get("status") or "in_progress") + body = { + "status": status, + "invocation_id": invocation_id, + "completed_steps": done, + "total_steps": total, + } + if status == "failed" and value.get("error"): + body["error"] = value["error"] + return JSONResponse(body) + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/requirements.txt b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/requirements.txt new file mode 100644 index 000000000000..c4ace62b69a8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_hello_world/requirements.txt @@ -0,0 +1,9 @@ +# Deliberately minimal: this sample needs ONLY the two agentserver packages — +# no azure-ai-projects, no openai, no langgraph — so it runs with the smallest +# possible dependency and memory footprint. +# +# Constrained to compatible release lines (matching the package's own +# convention): the strictly-opt-in durable-task behaviour and local state store +# this sample relies on require Core 2.1.0b1+. +azure-ai-agentserver-core>=2.1.0b1,<3.0.0 +azure-ai-agentserver-invocations>=1.2.0b1,<2.0.0 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py index ba58ee46bdf9..adc58bb9dac4 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_langgraph/app.py @@ -67,6 +67,7 @@ EventStreamNotFoundError, streams, ) +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled from azure.ai.agentserver.invocations import InvocationAgentServerHost try: @@ -81,6 +82,15 @@ # stream id is the per-turn ``invocation_id``. streams.use_in_memory_replay(ttl_seconds=600) +# Resilient tasks (durable execution + crash recovery) are strictly opt-in as of +# azure-ai-agentserver-core 2.1.0b1: ``AgentServerHost`` constructs the +# ``TaskManager`` ONLY when this switch is on, and it must be set before host +# startup (i.e. at module-import time). Without it, ``get_task_manager()`` / +# ``langgraph_session.start()`` raise ``TaskManagerNotInitialized`` and the agent +# runs non-durably with no crash recovery — defeating the purpose of a resilient +# long-running agent. +set_resilient_tasks_enabled(True) + app = InvocationAgentServerHost() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py index b61a88ea70c9..743f3211a8d1 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_multiturn/app.py @@ -37,7 +37,7 @@ from starlette.responses import JSONResponse, Response from azure.ai.agentserver.core.storage import FoundryStateStore -from azure.ai.agentserver.core.tasks import TaskConflictError +from azure.ai.agentserver.core.tasks import TaskConflictError, set_resilient_tasks_enabled from azure.ai.agentserver.invocations import InvocationAgentServerHost try: @@ -45,6 +45,15 @@ except ImportError: # allows `python app.py` from inside this directory from agent import invocation_state_store_name, session_workflow +# Resilient tasks (durable execution + crash recovery) are strictly opt-in as of +# azure-ai-agentserver-core 2.1.0b1: ``AgentServerHost`` constructs the +# ``TaskManager`` ONLY when this switch is on, and it must be set before host +# startup (i.e. at module-import time). Without it, ``get_task_manager()`` / +# ``session_workflow.start()`` raise ``TaskManagerNotInitialized`` and the agent +# runs non-durably with no crash recovery — defeating the purpose of a resilient +# long-running agent. +set_resilient_tasks_enabled(True) + app = InvocationAgentServerHost() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py index 00b64c2813ff..d50121722ae2 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/samples/resilient_research/app.py @@ -77,6 +77,7 @@ EventStreamNotFoundError, streams, ) +from azure.ai.agentserver.core.tasks import set_resilient_tasks_enabled from azure.ai.agentserver.invocations import InvocationAgentServerHost try: @@ -99,6 +100,15 @@ # only thing this sample supplies is the ``cursor_fn``. streams.use_file_backed_replay(cursor_fn=lambda ev: ev["sequence_number"]) +# Resilient tasks (durable execution + crash recovery) are strictly opt-in as of +# azure-ai-agentserver-core 2.1.0b1: ``AgentServerHost`` constructs the +# ``TaskManager`` ONLY when this switch is on, and it must be set before host +# startup (i.e. at module-import time). Without it, ``get_task_manager()`` / +# ``deep_research.start()`` raise ``TaskManagerNotInitialized`` and the agent +# runs non-durably with no crash recovery — defeating the purpose of a resilient +# long-running agent. +set_resilient_tasks_enabled(True) + app = InvocationAgentServerHost() diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_cancellable.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_cancellable.py new file mode 100644 index 000000000000..347b52b82260 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_cancellable.py @@ -0,0 +1,193 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""End-to-end test for the ``resilient_cancellable`` minimal sample. + +Fully self-contained (no LLM, no cloud): the durable checkpoint uses local +file-backed storage rooted at the test's ``tmp_path``. Drives the sample's task +in-process to exercise its core promise — a finite job that runs to completion, +but stops early (``status: cancelled``) when a durable cancel marker is present, +including across a recovery boundary. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +import pytest_asyncio + + +@pytest_asyncio.fixture +async def task_manager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """A real TaskManager backed by the local file provider at tmp_path.""" + + (tmp_path / "tasks").mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("AGENTSERVER_STATE_ROOT", str(tmp_path)) + monkeypatch.delenv("FOUNDRY_HOSTING_ENVIRONMENT", raising=False) + + from azure.ai.agentserver.core.tasks._manager import ( # noqa: WPS433 + TaskManager, + set_task_manager, + ) + + config = type( + "C", + (), + { + "agent_name": "test-cancellable", + "session_id": "test-cancellable-session", + "agent_version": "1.0.0", + "is_hosted": False, + }, + )() + mgr = TaskManager(config=config, shutdown_event=asyncio.Event()) + set_task_manager(mgr) + await mgr.startup() + try: + yield mgr + finally: + await mgr.shutdown() + set_task_manager(None) + + +@pytest.fixture(autouse=True) +def _samples_on_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Prepend the samples dir to ``sys.path`` (auto-restored after each test).""" + samples = Path(__file__).resolve().parent.parent.parent / "samples" + monkeypatch.syspath_prepend(str(samples)) + + +_SESSION = "test-cj-session" +_USER = "test-cj-user" + + +async def _load_item(cj, key: str): + store = await cj.open_checkpoint_store(_SESSION, _USER) + async with store: + return await store.get_item(key) + + +async def _seed_item(cj, key: str, value: dict) -> None: + store = await cj.open_checkpoint_store(_SESSION, _USER) + async with store: + await store.set_item(key, value) + + +async def _wait_for_steps(cj, key: str, minimum: int, timeout: float = 5.0) -> int: + """Poll the durable checkpoint until ``completed_steps`` reaches ``minimum``.""" + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + item = await _load_item(cj, key) + if item is not None: + done = int(item.value.get("completed_steps", 0) or 0) + if done >= minimum: + return done + await asyncio.sleep(0.02) + raise AssertionError( + f"completed_steps did not reach {minimum} within {timeout}s for {key}" + ) + + +@pytest.mark.asyncio +async def test_runs_to_completion_when_not_cancelled( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """Without a cancel marker the job finishes all steps.""" + from resilient_cancellable import agent as cj # noqa: WPS433 + + monkeypatch.setattr(cj, "_STEP_DELAY", 0.0) + + task_id = "cj-complete" + run = await cj.cancellable_job.start( + task_id=task_id, + input={"name": "alice", "steps": 3, "session_id": _SESSION, "user_id": _USER}, + ) + result = await run.result() + + assert result["status"] == "completed" + item = await _load_item(cj, task_id) + assert item.value.get("completed_steps") == 3 + assert item.value.get("status") == "completed" + + +@pytest.mark.asyncio +async def test_cancel_marker_before_start_stops_immediately( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """A cancel marker present before the first step stops the job at step 0.""" + from resilient_cancellable import agent as cj # noqa: WPS433 + + monkeypatch.setattr(cj, "_STEP_DELAY", 0.0) + + task_id = "cj-cancel-early" + await _seed_item(cj, f"{task_id}{cj.CANCEL_SUFFIX}", {"cancel": True}) + + run = await cj.cancellable_job.start( + task_id=task_id, + input={"name": "bob", "steps": 5, "session_id": _SESSION, "user_id": _USER}, + ) + result = await run.result() + + assert result["status"] == "cancelled" + assert result["completed_steps"] == 0 + item = await _load_item(cj, task_id) + assert item.value.get("status") == "cancelled" + + +@pytest.mark.asyncio +async def test_cancel_mid_run_stops_early( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """A cancel marker written mid-run stops the job before it finishes.""" + from resilient_cancellable import agent as cj # noqa: WPS433 + + monkeypatch.setattr(cj, "_STEP_DELAY", 0.02) + + task_id = "cj-cancel-mid" + run = await cj.cancellable_job.start( + task_id=task_id, + input={"name": "carol", "steps": 50, "session_id": _SESSION, "user_id": _USER}, + ) + + # Let it make some progress, then request cancel. + reached = await _wait_for_steps(cj, task_id, minimum=2) + await _seed_item(cj, f"{task_id}{cj.CANCEL_SUFFIX}", {"cancel": True}) + + result = await asyncio.wait_for(run.result(), timeout=5.0) + assert result["status"] == "cancelled" + # Stopped early: at least where we saw it, and short of the full 50. + assert reached <= result["completed_steps"] < 50 + item = await _load_item(cj, task_id) + assert item.value.get("status") == "cancelled" + + +@pytest.mark.asyncio +async def test_recovered_run_honours_pending_cancel( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """A partially-done job with a pending cancel marker cancels on resume.""" + from resilient_cancellable import agent as cj # noqa: WPS433 + + monkeypatch.setattr(cj, "_STEP_DELAY", 0.0) + + task_id = "cj-recover-cancel" + # Simulate a run that had done 2/10 steps and a cancel that landed before a + # crash: both the checkpoint and the cancel marker are already durable. + await _seed_item( + cj, + task_id, + {"name": "dave", "steps": 10, "completed_steps": 2, "status": "in_progress"}, + ) + await _seed_item(cj, f"{task_id}{cj.CANCEL_SUFFIX}", {"cancel": True}) + + run = await cj.cancellable_job.start( + task_id=task_id, + input={"name": "dave", "steps": 10, "session_id": _SESSION, "user_id": _USER}, + ) + result = await run.result() + + assert result["status"] == "cancelled" + # Stopped at the resume point (2), not restarted at 0 and not finished at 10. + assert result["completed_steps"] == 2 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_forever.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_forever.py new file mode 100644 index 000000000000..61977b5a2273 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_forever.py @@ -0,0 +1,214 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""End-to-end test for the ``resilient_hello_forever`` minimal sample. + +The hello-forever sample is an *indefinite* durable worker: it ticks forever, +checkpointing its ``iterations`` cursor after every tick, and only stops on an +explicit cancel that is confirmed by a durable "stop" marker. Like the other +minimal sample it is **fully self-contained** (no LLM, no cloud) — its +checkpoint uses local file-backed storage rooted at the test's ``tmp_path``. + +This is *not* a live test: it imports the sample's task directly and drives it +in-process. It exercises the two behaviours that make an infinite loop a +well-behaved LRA: + +- The worker ticks and checkpoints its ``iterations`` cursor. +- An explicit cancel + durable stop marker stops it terminally (``stopped``), + and it stops from wherever its checkpoint had reached (i.e. it resumes from a + pre-existing checkpoint rather than restarting at 0). +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +import pytest_asyncio + + +@pytest_asyncio.fixture +async def task_manager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """A real TaskManager backed by ``LocalFileTaskProvider`` at tmp_path.""" + + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("AGENTSERVER_STATE_ROOT", str(tmp_path)) + monkeypatch.delenv("FOUNDRY_HOSTING_ENVIRONMENT", raising=False) + + from azure.ai.agentserver.core.tasks._manager import ( # noqa: WPS433 + TaskManager, + set_task_manager, + ) + + config = type( + "C", + (), + { + "agent_name": "test-hello-forever", + "session_id": "test-hello-forever-session", + "agent_version": "1.0.0", + "is_hosted": False, + }, + )() + mgr = TaskManager(config=config, shutdown_event=asyncio.Event()) + set_task_manager(mgr) + await mgr.startup() + try: + yield mgr + finally: + await mgr.shutdown() + set_task_manager(None) + + +@pytest.fixture(autouse=True) +def _samples_on_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Prepend the samples dir to ``sys.path`` (auto-restored after each test).""" + samples = Path(__file__).resolve().parent.parent.parent / "samples" + monkeypatch.syspath_prepend(str(samples)) + + +async def _load_item(hf, key: str): + store = await hf.open_checkpoint_store(_SESSION, _USER) + async with store: + return await store.get_item(key) + + +async def _seed_item(hf, key: str, value: dict) -> None: + store = await hf.open_checkpoint_store(_SESSION, _USER) + async with store: + await store.set_item(key, value) + + +async def _wait_for_iterations(hf, key: str, minimum: int, timeout: float = 5.0) -> int: + """Poll the durable checkpoint until ``iterations`` reaches ``minimum``.""" + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + item = await _load_item(hf, key) + if item is not None: + iters = int(item.value.get("iterations", 0) or 0) + if iters >= minimum: + return iters + await asyncio.sleep(0.02) + raise AssertionError( + f"iterations did not reach {minimum} within {timeout}s for task {key}" + ) + + +# A fixed session + user id so the tests read the same session-scoped, +# user-isolated store the worker writes to (the host injects both from request +# state at runtime). +_SESSION = "test-hf-session" +_USER = "test-hf-user" + + +def test_durable_task_id_is_accepted_by_task_manager() -> None: + """``durable_task_id`` must produce an id the TaskManager accepts. + + Regression guard: a ``/`` separator (or an over-long id) is rejected by + ``start()`` — a bug the task-level tests miss because they pass their own + slash-free ids. Validate the helper's output with the real SDK validator, + including long protocol ids that would blow the 128-char limit if not hashed. + """ + from azure.ai.agentserver.core.tasks._validation import ( # noqa: WPS433 + validate_task_id, + ) + from resilient_hello_forever import agent as hf # noqa: WPS433 + + validate_task_id(hf.durable_task_id("sess-1", "inv-1", "user-1")) + validate_task_id(hf.durable_task_id("s" * 200, "inv_" + "x" * 200, "u" * 200)) + assert hf.durable_task_id("a", "bc", "d") != hf.durable_task_id("ab", "c", "d") + assert hf.durable_task_id("s", "i", "u1") != hf.durable_task_id("s", "i", "u2") + + +@pytest.mark.asyncio +async def test_ticks_then_stops_on_cancel_with_marker( + task_manager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The worker ticks, then an explicit cancel + stop marker stops it.""" + from resilient_hello_forever import agent as hf # noqa: WPS433 + + monkeypatch.setattr(hf, "_TICK", 0.01) + + task_id = "hf-stop" + run = await hf.hello_forever.start( + task_id=task_id, + input={"name": "alice", "session_id": _SESSION, "user_id": _USER}, + ) + + # Let it tick a few times so there is real progress to stop. + reached = await _wait_for_iterations(hf, task_id, minimum=2) + + # Request a stop: write the durable marker, then signal cancel. The worker + # only treats cancel as a real stop when the marker is present. + await _seed_item(hf, f"{task_id}{hf.STOP_SUFFIX}", {"stop": True}) + await run.cancel() + + result = await asyncio.wait_for(run.result(), timeout=5.0) + + assert result["stopped"] is True + assert result["name"] == "alice" + assert result["iterations"] >= reached + + +@pytest.mark.asyncio +async def test_stops_on_marker_without_local_cancel( + task_manager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The durable stop marker alone stops the worker (cross-replica case). + + A cancel routed to a different replica writes the marker but cannot set this + process's ``ctx.cancel``. The worker must still stop, because it re-checks + the marker every iteration independently of ``ctx.cancel``. + """ + from resilient_hello_forever import agent as hf # noqa: WPS433 + + monkeypatch.setattr(hf, "_TICK", 0.01) + + task_id = "hf-marker-only" + run = await hf.hello_forever.start( + task_id=task_id, + input={"name": "carol", "session_id": _SESSION, "user_id": _USER}, + ) + + await _wait_for_iterations(hf, task_id, minimum=2) + + # Write ONLY the durable stop marker — do not call run.cancel(). + await _seed_item(hf, f"{task_id}{hf.STOP_SUFFIX}", {"stop": True}) + + result = await asyncio.wait_for(run.result(), timeout=5.0) + assert result["stopped"] is True + assert result["name"] == "carol" + + +@pytest.mark.asyncio +async def test_resumes_from_existing_checkpoint( + task_manager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A pre-existing checkpoint makes the worker resume, not restart at 0.""" + from resilient_hello_forever import agent as hf # noqa: WPS433 + + monkeypatch.setattr(hf, "_TICK", 0.01) + + task_id = "hf-resume" + await _seed_item(hf, task_id, {"name": "bob", "iterations": 5}) + + run = await hf.hello_forever.start( + task_id=task_id, + input={"name": "bob", "session_id": _SESSION, "user_id": _USER}, + ) + + # It must continue past the seeded cursor rather than counting up from 1. + reached = await _wait_for_iterations(hf, task_id, minimum=6) + assert reached >= 6 + + await _seed_item(hf, f"{task_id}{hf.STOP_SUFFIX}", {"stop": True}) + await run.cancel() + + result = await asyncio.wait_for(run.result(), timeout=5.0) + assert result["stopped"] is True + assert result["iterations"] >= 6 diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_handlers.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_handlers.py new file mode 100644 index 000000000000..659c6fd6de6e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_handlers.py @@ -0,0 +1,366 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""HTTP-handler smoke tests for the minimal resilient samples. + +The other e2e tests drive the ``@task`` functions directly; they deliberately +bypass ``app.py``. This file exercises the actual invoke / poll / cancel +handlers of both minimal samples so handler-level regressions (a ``NameError`` +in the cancel path, a malformed-body 500, a wrong status code, or a +pre-checkpoint 404) are caught. The handler decorators return the original +function unmodified, so we call them directly with a lightweight fake request. + +Fully self-contained: local file-backed task provider + state store rooted at +the test's ``tmp_path``. No LLM, no cloud. +""" + +from __future__ import annotations + +import asyncio +import json +import types +from pathlib import Path + +import pytest +import pytest_asyncio + + +class _Req: + """Minimal stand-in for the Starlette request the handlers consume.""" + + def __init__( + self, + *, + body: bytes = b"", + invocation_id: str = "", + session_id: str = "", + user_id: str = "u", + call_id: str = "c", + ) -> None: + self._body = body + self.state = types.SimpleNamespace( + invocation_id=invocation_id, + session_id=session_id, + user_id=user_id, + call_id=call_id, + ) + + async def body(self) -> bytes: + return self._body + + +@pytest.fixture(autouse=True) +def _samples_on_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Prepend the samples dir to ``sys.path`` (auto-restored after each test).""" + samples = Path(__file__).resolve().parent.parent.parent / "samples" + monkeypatch.syspath_prepend(str(samples)) + + +@pytest_asyncio.fixture +async def task_manager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """A real TaskManager backed by the local file provider at tmp_path.""" + import asyncio # noqa: WPS433 + + (tmp_path / "tasks").mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("AGENTSERVER_STATE_ROOT", str(tmp_path)) + monkeypatch.delenv("FOUNDRY_HOSTING_ENVIRONMENT", raising=False) + + from azure.ai.agentserver.core.tasks import ( # noqa: WPS433 + resilient_tasks_enabled, + set_resilient_tasks_enabled, + ) + from azure.ai.agentserver.core.tasks._manager import ( # noqa: WPS433 + TaskManager, + set_task_manager, + ) + + # Importing the sample app modules flips the process-global opt-in flag to + # True. Save and restore it so the flag does not leak into later tests and + # make their behavior test-order dependent. + prev_enabled = resilient_tasks_enabled() + + config = type( + "C", + (), + { + "agent_name": "test-hello-handlers", + "session_id": "test-hello-handlers-session", + "agent_version": "1.0.0", + "is_hosted": False, + }, + )() + mgr = TaskManager(config=config, shutdown_event=asyncio.Event()) + set_task_manager(mgr) + await mgr.startup() + try: + yield mgr + finally: + await mgr.shutdown() + set_task_manager(None) + set_resilient_tasks_enabled(prev_enabled) + + +def _status(response) -> tuple[int, dict]: + return response.status_code, json.loads(response.body) + + +@pytest.mark.asyncio +async def test_hello_forever_handlers_lifecycle( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """invoke -> running -> cancel -> stopped, plus 404 and 400 paths.""" + from resilient_hello_forever import agent as hf # noqa: WPS433 + import resilient_hello_forever.app as app # noqa: WPS433 + + monkeypatch.setattr(hf, "_TICK", 0.01) + inv, sess = "inv-hf-1", "sess-hf" + + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "ada"}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 202 and body["invocation_id"] == inv + + code, body = _status( + await app.handle_get(_Req(invocation_id=inv, session_id=sess)) + ) + assert code == 200 and body["status"] == "running" + + # The cancel path previously raised NameError after writing the marker. + code, body = _status( + await app.handle_cancel(_Req(invocation_id=inv, session_id=sess)) + ) + assert code == 200 and body["status"] == "cancelling" + + code, body = _status( + await app.handle_get(_Req(invocation_id=inv, session_id=sess)) + ) + assert code == 200 and body["status"] == "stopped" + + # Reusing the same invocation id after it stopped must NOT start a second + # worker against the stale checkpoint/marker — the durable record already + # exists, so invoke returns 409 with the terminal status. + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "ada"}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 409 and body["status"] == "stopped" + + # A failed worker (one-shot record deleted) is surfaced from the durable + # ``status: failed`` the worker persists — not reported as ``running`` forever. + fail_inv = "inv-hf-failed" + fail_tid = hf.durable_task_id(sess, fail_inv, "u") + failed_store = await hf.open_checkpoint_store(sess, "u") + async with failed_store: + await failed_store.set_item( + fail_tid, {"name": "z", "iterations": 4, "status": "failed", "error": "boom"} + ) + code, body = _status( + await app.handle_get(_Req(invocation_id=fail_inv, session_id=sess)) + ) + assert code == 200 and body["status"] == "failed" and body["error"] == "boom" + + # Unknown invocation: poll and cancel both 404, and cancel must NOT persist a + # stop marker for it. + code, _ = _status(await app.handle_get(_Req(invocation_id="nope", session_id=sess))) + assert code == 404 + code, _ = _status( + await app.handle_cancel(_Req(invocation_id="nope2", session_id=sess)) + ) + assert code == 404 + + # Malformed body shapes -> 400 (not 500): non-object JSON and invalid UTF-8. + code, _ = _status( + await app.handle_invoke(_Req(body=b"[]", invocation_id="inv-hf-2", session_id=sess)) + ) + assert code == 400 + code, _ = _status( + await app.handle_invoke( + _Req(body=b"\xff\xfe", invocation_id="inv-hf-3", session_id=sess) + ) + ) + assert code == 400 + + +@pytest.mark.asyncio +async def test_hello_world_recovers_orphaned_seed( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """A nonterminal durable record with no live task is re-scheduled, not 409'd. + + Simulates the orphan window: an earlier attempt seeded ``in_progress`` then + crashed before the TaskManager record became durable. A later POST must + (re-)schedule the task idempotently and run it to completion rather than + wedging behind a permanent ``in_progress`` record with nothing to recover. + """ + from resilient_hello_world import agent as hw # noqa: WPS433 + import resilient_hello_world.app as app # noqa: WPS433 + + monkeypatch.setattr(hw, "_STEP_DELAY", 0.0) + inv, sess = "inv-hw-orphan", "sess-hw" + + # Seed an orphaned nonterminal record directly (no task started) with real + # progress (3/5) so we can prove the retry does not roll it back. + task_id = hw.durable_task_id(sess, inv, "u") + store = await hw.open_checkpoint_store(sess, "u") + async with store: + await store.create_item( + task_id, {"name": "ada", "steps": 5, "completed_steps": 3, "status": "in_progress"} + ) + + # POST with a DIFFERENT body (steps=99): must recover (202) but reuse the + # PERSISTED parameters, not the retry's — so it finishes at the original 5. + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "eve", "steps": 99}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 202, f"expected orphan recovery (202), got {code}: {body}" + + # And it runs to completion at the persisted total (5), preserving progress. + for _ in range(200): + code, body = _status( + await app.handle_get(_Req(invocation_id=inv, session_id=sess)) + ) + if body["status"] == "completed": + break + await asyncio.sleep(0.01) + assert body["status"] == "completed" + assert body["total_steps"] == 5 and body["completed_steps"] == 5 + + +@pytest.mark.asyncio +async def test_hello_world_handlers( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """invoke -> poll (never 404 for a started run), plus 400/404 paths.""" + from resilient_hello_world import agent as hw # noqa: WPS433 + import resilient_hello_world.app as app # noqa: WPS433 + + monkeypatch.setattr(hw, "_STEP_DELAY", 0.0) + inv, sess = "inv-hw-1", "sess-hw" + + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "ada", "steps": 3}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 202 and body["total_steps"] == 3 + + # Drain the (near-instant) background job to a terminal state so subsequent + # assertions are deterministic rather than racing the task's writes. + final = None + for _ in range(200): + code, body = _status( + await app.handle_get(_Req(invocation_id=inv, session_id=sess)) + ) + assert code == 200 and body["status"] in ("in_progress", "completed") + if body["status"] == "completed": + final = body + break + await asyncio.sleep(0.01) + assert final is not None and final["completed_steps"] == 3 + + # Reusing an existing invocation id must not start a second task against the + # existing checkpoint — invoke returns 409 with the current (terminal) status. + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "ada", "steps": 3}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 409 and body["status"] == "completed" + + # A failed run (one-shot record deleted) is surfaced from the durable + # ``status: failed`` the task persists — not stuck at ``in_progress`` forever. + fail_inv = "inv-hw-failed" + fail_tid = hw.durable_task_id(sess, fail_inv, "u") + failed_store = await hw.open_checkpoint_store(sess, "u") + async with failed_store: + await failed_store.set_item( + fail_tid, + {"name": "z", "steps": 5, "completed_steps": 2, "status": "failed", "error": "boom"}, + ) + code, body = _status( + await app.handle_get(_Req(invocation_id=fail_inv, session_id=sess)) + ) + assert code == 200 and body["status"] == "failed" and body["error"] == "boom" + + code, _ = _status(await app.handle_get(_Req(invocation_id="nope", session_id=sess))) + assert code == 404 + + # Malformed / invalid steps -> 400 (non-object body, non-positive, bool, + # float, and invalid UTF-8 bytes). + for raw in (b"[]", b'{"steps": 0}', b'{"steps": true}', b'{"steps": 2.9}', b"\xff\xfe"): + code, _ = _status( + await app.handle_invoke( + _Req(body=raw, invocation_id="inv-hw-x", session_id=sess) + ) + ) + assert code == 400, f"expected 400 for body {raw!r}, got {code}" + + +@pytest.mark.asyncio +async def test_cancellable_handlers_lifecycle( + task_manager, monkeypatch: pytest.MonkeyPatch +) -> None: + """invoke -> cancel -> cancelling -> cancelled, plus reuse 409 / 404 / 400.""" + from resilient_cancellable import agent as cj # noqa: WPS433 + import resilient_cancellable.app as app # noqa: WPS433 + + # Nonzero step delay so the job is still running when we cancel it. + monkeypatch.setattr(cj, "_STEP_DELAY", 0.05) + inv, sess = "inv-cj-1", "sess-cj" + + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "ada", "steps": 50}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 202 and body["total_steps"] == 50 + + # Request cancel — writes the durable marker. + code, body = _status( + await app.handle_cancel(_Req(invocation_id=inv, session_id=sess)) + ) + assert code == 200 and body["status"] == "cancelling" + + # Poll until the job observes the marker and finalizes as ``cancelled`` (it + # may briefly report ``cancelling`` first). + final = None + for _ in range(200): + code, body = _status( + await app.handle_get(_Req(invocation_id=inv, session_id=sess)) + ) + assert code == 200 and body["status"] in ("cancelling", "cancelled") + if body["status"] == "cancelled": + final = body + break + await asyncio.sleep(0.02) + assert final is not None, "job did not reach 'cancelled'" + assert final["completed_steps"] < 50 + + # Reusing the invocation id after cancel must not start a second job. + code, body = _status( + await app.handle_invoke( + _Req(body=b'{"name": "ada", "steps": 50}', invocation_id=inv, session_id=sess) + ) + ) + assert code == 409 and body["status"] == "cancelled" + + # Unknown invocation: poll and cancel both 404 (cancel must not seed a marker). + code, _ = _status(await app.handle_get(_Req(invocation_id="nope", session_id=sess))) + assert code == 404 + code, _ = _status( + await app.handle_cancel(_Req(invocation_id="nope2", session_id=sess)) + ) + assert code == 404 + + # Malformed bodies -> 400 (non-object, non-positive steps, invalid UTF-8). + for raw in (b"[]", b'{"steps": 0}', b"\xff\xfe"): + code, _ = _status( + await app.handle_invoke( + _Req(body=raw, invocation_id="inv-cj-x", session_id=sess) + ) + ) + assert code == 400, f"expected 400 for body {raw!r}, got {code}" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_world.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_world.py new file mode 100644 index 000000000000..ad6e6650f074 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/e2e/test_resilient_hello_world.py @@ -0,0 +1,208 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""End-to-end test for the ``resilient_hello_world`` minimal sample. + +The hello-world sample is **fully self-contained** (no Azure OpenAI, no +Copilot CLI, no Foundry endpoint). Its durable checkpoint uses local +file-backed storage rooted at the test's ``tmp_path``. + +This is *not* a live test: it imports the sample's task directly and drives +it through completion and a resume-from-checkpoint boundary in the same +process. It exercises the durable-checkpoint contract for the sample (the +structure test in ``test_resilient_samples_structure.py`` proves the files +exist and opt in to durable tasks; this file proves the task actually +checkpoints and resumes). + +Coverage: + +- A fresh run counts through every step and finishes ``complete``. +- Each step persists ``completed_steps`` to the durable checkpoint. +- A run whose checkpoint already shows every step done finishes without + redoing any work (resume-from-checkpoint skips completed steps). +- A partially-completed checkpoint resumes at the next step, not step 1. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import pytest_asyncio + + +@pytest_asyncio.fixture +async def task_manager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """A real TaskManager backed by ``LocalFileTaskProvider`` at tmp_path.""" + import asyncio # noqa: WPS433 + + tasks_dir = tmp_path / "tasks" + tasks_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setenv("AGENTSERVER_STATE_ROOT", str(tmp_path)) + monkeypatch.delenv("FOUNDRY_HOSTING_ENVIRONMENT", raising=False) + + from azure.ai.agentserver.core.tasks._manager import ( # noqa: WPS433 + TaskManager, + set_task_manager, + ) + + config = type( + "C", + (), + { + "agent_name": "test-hello-world", + "session_id": "test-hello-world-session", + "agent_version": "1.0.0", + "is_hosted": False, + }, + )() + mgr = TaskManager(config=config, shutdown_event=asyncio.Event()) + set_task_manager(mgr) + await mgr.startup() + try: + yield mgr + finally: + await mgr.shutdown() + set_task_manager(None) + + +@pytest.fixture(autouse=True) +def _samples_on_path(monkeypatch: pytest.MonkeyPatch) -> None: + """Prepend the samples dir to ``sys.path`` (auto-restored after each test).""" + samples = Path(__file__).resolve().parent.parent.parent / "samples" + monkeypatch.syspath_prepend(str(samples)) + + +async def _load_item(hw, key: str): + store = await hw.open_checkpoint_store(_SESSION, _USER) + async with store: + return await store.get_item(key) + + +async def _seed_item(hw, key: str, value: dict): + """Write an initial checkpoint and return its ETag.""" + store = await hw.open_checkpoint_store(_SESSION, _USER) + async with store: + ref = await store.set_item(key, value) + return ref.etag + + +# A fixed session + user id so the tests read the same session-scoped, +# user-isolated store the task writes to (the host injects both from request +# state at runtime). +_SESSION = "test-hw-session" +_USER = "test-hw-user" + + +def test_durable_task_id_is_accepted_by_task_manager() -> None: + """``durable_task_id`` must produce an id the TaskManager accepts. + + Regression guard: a ``/`` separator (or an over-long id) is rejected by the + provider validator — a bug the task-level tests miss because they pass their + own slash-free ids. Validate the helper's output with the real provider + validator, including long protocol ids that would blow the 128-char limit if + not hashed. + """ + from azure.ai.agentserver.core.tasks._validation import ( # noqa: WPS433 + validate_task_id, + ) + from resilient_hello_world import agent as hw # noqa: WPS433 + + validate_task_id(hw.durable_task_id("sess-1", "inv-1", "user-1")) + validate_task_id(hw.durable_task_id("s" * 200, "inv_" + "x" * 200, "u" * 200)) + # Distinct inputs must not collide across any of the three fields. + assert hw.durable_task_id("a", "bc", "d") != hw.durable_task_id("ab", "c", "d") + assert hw.durable_task_id("s", "i", "u1") != hw.durable_task_id("s", "i", "u2") + + +@pytest.mark.asyncio +async def test_runs_to_completion_and_checkpoints_every_step( + task_manager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A fresh run finishes ``complete`` with the final step checkpointed.""" + from resilient_hello_world import agent as hw # noqa: WPS433 + + monkeypatch.setattr(hw, "_STEP_DELAY", 0.0) + + task_id = "hw-complete" + run = await hw.hello_world.start( + task_id=task_id, + input={"name": "alice", "steps": 3, "session_id": _SESSION, "user_id": _USER}, + ) + result = await run.result() + + assert result == {"name": "alice", "steps": 3, "status": "complete"} + + item = await _load_item(hw, task_id) + assert item is not None + assert item.value.get("completed_steps") == 3 + assert item.value.get("steps") == 3 + + +@pytest.mark.asyncio +async def test_completed_checkpoint_skips_all_work( + task_manager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A checkpoint already at ``steps`` finishes without rewriting it. + + An unchanged ETag proves the handler resumed from the finalized checkpoint and + ran zero steps instead of starting over. + """ + from resilient_hello_world import agent as hw # noqa: WPS433 + + monkeypatch.setattr(hw, "_STEP_DELAY", 0.0) + + task_id = "hw-already-done" + seeded_etag = await _seed_item( + hw, + task_id, + {"name": "bob", "steps": 3, "completed_steps": 3, "status": "completed"}, + ) + + run = await hw.hello_world.start( + task_id=task_id, + input={"name": "bob", "steps": 3, "session_id": _SESSION, "user_id": _USER}, + ) + result = await run.result() + + assert result["status"] == "complete" + + item = await _load_item(hw, task_id) + assert item is not None + assert item.value.get("completed_steps") == 3 + assert item.value.get("status") == "completed" + # Already finalized, so no write happened → ETag is unchanged. + assert item.etag == seeded_etag + + +@pytest.mark.asyncio +async def test_partial_checkpoint_resumes_at_next_step( + task_manager, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A partial checkpoint continues to completion instead of restarting.""" + from resilient_hello_world import agent as hw # noqa: WPS433 + + monkeypatch.setattr(hw, "_STEP_DELAY", 0.0) + + task_id = "hw-partial" + seeded_etag = await _seed_item( + hw, + task_id, + {"name": "carol", "steps": 4, "completed_steps": 2}, + ) + + run = await hw.hello_world.start( + task_id=task_id, + input={"name": "carol", "steps": 4, "session_id": _SESSION, "user_id": _USER}, + ) + result = await run.result() + + assert result == {"name": "carol", "steps": 4, "status": "complete"} + + item = await _load_item(hw, task_id) + assert item is not None + assert item.value.get("completed_steps") == 4 + # Work continued from the seed, so the checkpoint was rewritten. + assert item.etag != seeded_etag diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_resilient_samples_structure.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_resilient_samples_structure.py index 20005640345f..b6615affd8e1 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_resilient_samples_structure.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_resilient_samples_structure.py @@ -31,6 +31,7 @@ from __future__ import annotations +import ast from pathlib import Path import pytest @@ -47,6 +48,14 @@ "resilient_research", ) +# Minimal, dependency-light long-running-agent samples (core + invocations only, +# no LLM). They ship a README in addition to the standard three files. +_MINIMAL_RESILIENT_SAMPLES: tuple[str, ...] = ( + "resilient_hello_world", + "resilient_hello_forever", + "resilient_cancellable", +) + _DROPPED_SAMPLES: tuple[str, ...] = ("resilient_claude", "resilient_copilot") _REQUIRED_FILES_PER_SAMPLE: tuple[str, ...] = ( @@ -169,3 +178,140 @@ def test_sample_has_no_retired_name_references(sample_name: str) -> None: # core/invocations packages. The demo has been split into its own branch # and is no longer part of this # package's shipping surface, so the structural guard is no longer relevant. + + +# --------------------------------------------------------------------------- +# 6. Minimal (no-LLM) resilient samples +# --------------------------------------------------------------------------- + +_REQUIRED_FILES_MINIMAL_SAMPLE: tuple[str, ...] = ( + "agent.py", + "app.py", + "requirements.txt", + "README.md", +) + + +@pytest.mark.parametrize("sample_name", _MINIMAL_RESILIENT_SAMPLES) +def test_minimal_resilient_sample_directory_exists(sample_name: str) -> None: + """Each minimal long-running-agent sample directory MUST exist.""" + + p = _sample_path(sample_name) + assert p.is_dir(), ( + f"Minimal resilient invocation sample missing: {p}. " + f"Expected samples: {', '.join(_MINIMAL_RESILIENT_SAMPLES)}." + ) + + +@pytest.mark.parametrize("sample_name", _MINIMAL_RESILIENT_SAMPLES) +@pytest.mark.parametrize("filename", _REQUIRED_FILES_MINIMAL_SAMPLE) +def test_minimal_required_files_per_sample(sample_name: str, filename: str) -> None: + """Each minimal sample ships agent + app + requirements + a README walkthrough.""" + + p = _sample_path(sample_name) / filename + assert p.is_file(), ( + f"Missing required file {filename} for minimal sample {sample_name} " + f"(expected at {p})." + ) + + +@pytest.mark.parametrize("sample_name", _MINIMAL_RESILIENT_SAMPLES) +def test_minimal_sample_has_no_retired_name_references(sample_name: str) -> None: + """Minimal samples MUST NOT reference retired task-framework names.""" + + offenders: list[tuple[str, str]] = [] + for src in _python_sources_under(_sample_path(sample_name)): + text = src.read_text(encoding="utf-8") + for name in _RETIRED_NAMES: + if name in text: + offenders.append((str(src.relative_to(_SAMPLES_DIR)), name)) + assert not offenders, ( + f"Retired task-framework names still referenced in minimal sample " + f"{sample_name}: {offenders}." + ) + + +# --------------------------------------------------------------------------- +# 7. Regression guard: durable tasks are strictly opt-in (2.1.0b1+) +# --------------------------------------------------------------------------- +# +# Since core 2.1.0b1, ``get_task_manager()`` raises ``TaskManagerNotInitialized`` +# unless ``set_resilient_tasks_enabled(True)`` runs before host startup. Every +# resilient invocation sample MUST enable it in ``app.py`` or the durable / +# crash-recovery behaviour the sample advertises silently does not work. + + +def _call_name(node: ast.Call) -> str | None: + """Return the (possibly dotted-tail) function name of a call node.""" + fn = node.func + if isinstance(fn, ast.Attribute): + return fn.attr + if isinstance(fn, ast.Name): + return fn.id + return None + + +def _module_level_call_statements(source: str): + """Yield ``(lineno, call_node)`` for calls that run at import time. + + Only *direct* module-level statements are considered — a bare expression + statement (``set_resilient_tasks_enabled(True)``) or the call on the + right-hand side of a module-level assignment (``app = Host()``). Calls nested + inside a ``def``/``class``/``if``/``with`` are ignored, because they do not + (necessarily) execute on import; ``ast.walk`` would wrongly accept a call + defined inside an uncalled helper. + """ + for stmt in ast.parse(source).body: + value = None + if isinstance(stmt, ast.Expr): + value = stmt.value + elif isinstance(stmt, (ast.Assign, ast.AnnAssign)): + value = stmt.value + if isinstance(value, ast.Call): + yield stmt.lineno, value + + +def _has_module_level_enable_true(source: str) -> bool: + """Whether a real module-level ``set_resilient_tasks_enabled(True)`` runs on import. + + Scans direct module-level statements (a bare expression call, or the call on + the RHS of a module-level assignment) for a ``set_resilient_tasks_enabled`` + call with a literal ``True`` argument — passed either positionally + (``set_resilient_tasks_enabled(True)``) or by keyword + (``set_resilient_tasks_enabled(enabled=True)``). Nested calls (inside a + ``def``/``class``/``if``) are ignored because they do not necessarily execute + on import. + """ + for _lineno, call in _module_level_call_statements(source): + if _call_name(call) != "set_resilient_tasks_enabled": + continue + args_and_kwargs = list(call.args) + [kw.value for kw in call.keywords] + if any(isinstance(a, ast.Constant) and a.value is True for a in args_and_kwargs): + return True + return False + + +@pytest.mark.parametrize( + "sample_name", _REQUIRED_RESILIENT_SAMPLES + _MINIMAL_RESILIENT_SAMPLES +) +def test_resilient_sample_enables_resilient_tasks(sample_name: str) -> None: + """Every resilient sample's ``app.py`` MUST opt in to durable tasks at import. + + Parses the module AST (rather than matching a substring, which a comment or a + call nested in an uncalled helper would satisfy) and asserts a real + ``set_resilient_tasks_enabled(True)`` runs as a **module-level** statement, so + it executes on import — before the host's lifespan starts and reads the switch + (``core/_base.py`` reads it inside ``_lifespan``, not at + ``InvocationAgentServerHost()`` construction, so relative ordering versus the + host object is irrelevant; only import-time execution matters). + """ + + app_py = _sample_path(sample_name) / "app.py" + assert app_py.is_file(), f"Missing app.py for sample {sample_name} ({app_py})." + assert _has_module_level_enable_true(app_py.read_text(encoding="utf-8")), ( + f"Sample {sample_name} does not call set_resilient_tasks_enabled(True) as a " + "module-level statement in app.py. Durable tasks are strictly opt-in since " + "core 2.1.0b1; without an import-time call the host lifespan starts with the " + "switch off, get_task_manager() raises TaskManagerNotInitialized, and " + "long-running recovery is silently disabled." + )