AgentServer: integrate resilient tasks with memory snapshot session activation

[Nathandrake229]

Summary

Separate process-side task registration from session-bound TaskManager
construction/startup for golden memory snapshots.

Replaces #49485. That draft was stacked on
feature/agentserver-core-lifecycle-hooks. GitHub closed it when #49217 merged
and its base branch was deleted, and a closed PR cannot be retargeted or
reopened. This PR carries the same change on main (based on 037e8b7, which
contains #49217), so nothing is stacked any more.

Differences from #49485

The change is otherwise identical to #49485's final head (4540307). Rebasing
onto the merged lifecycle code and running the full Core and Responses
suites (#49485 was validated with targeted selectors only) needed these
adjustments:

• After-restore rollback (Core). The merged lifecycle code snapshots the
  whole AgentConfig for rollback and no longer captures the previous session
  ID/GUID that the worker-quiesce check compares with the task runtime binding.
  The two captures are restored next to the new snapshot. Without them both
  rollback paths raised NameError; eight targeted tests caught it.
• TaskManager.shutdown(wait_for_workers=False). AgentServer: integrate resilient tasks with memory snapshot session activation #49485 always awaited the
  cancelled workers. That made a grace-expired shutdown persist suspended
  instead of leaving the record in_progress, which breaks the existing
  test_grace_expires_treated_like_crash. Waiting is now requested only by the
  pre-rollback quiesce; ordinary shutdown keeps its existing contract.
  Rollback additionally preserves unfinished records/inputs: it joins workers
  under their current binding before expiring leases, allowing the original
  task to recover on retry rather than becoming terminal cancellation.
• Responses fallback. With FOUNDRY_MEMORY_SNAPSHOT_ENABLED absent or false, a missing
  task runtime keeps the existing in-process fallback (the orchestrator already
  documents it for in-process clients whose lifespan never runs). The
  enabled-but-uninitialized error applies only while startup is deferred (true).
  Absence uses the cold-start path directly, including local/in-process clients.
  Deferred-start failures remove only the exact caller-scoped execution before
  propagating the platform error, rather than retaining work that never started.
• Test isolation. test_snapshot_task_activation.py restores os.environ
  after each test (a restored session ID leaked into test_tracing.py) and pins
  the logger level in its one log assertion (test_logger.py leaves
  azure.ai.agentserver at ERROR).
  Responses tests use a per-test cold-start env fixture instead of global
  setdefault, so inherited capture-mode configuration cannot change ordinary
  tests. Snapshot and absent-mode tests explicitly override that default.
• Lint and format. Pragmas or extraction keep the three touched source files
  pylint-clean (they are clean on main); _manager.py is reflowed to the
  repository's Black width.
• Housekeeping. CHANGELOG entries are merged into the existing Core
  2.3.0b1 and Responses 2.4.0b1 sections. Responses _version.py and
  api.metadata.yml no longer change because main already carries 2.4.0b1.

Design specification

Resilient tasks with memory snapshots

The spec covers lifecycle ordering, trigger/enablement separation, phase retry,
owned runtime cleanup, compatibility, persistence boundaries, validation, and
the unresolved startup-task alternatives.

Implementation

• Keep task definitions and Responses task registration in process memory.
• FOUNDRY_MEMORY_SNAPSHOT_ENABLED=true defers enabled tasks in capture;
  absent/false uses normal cold/local session initialization.
• set_resilient_tasks_enabled(False) disables the subsystem at both startup
  and after-restore regardless of the environment signal.
• After-restore activates requested tasks after session hydration and existing
  dependency callbacks, before successful restore commit.
• Retain restore-ID initialization phases, serialize startup, and ensure one
  periodic recovery scanner.
• Preserve immediate best-effort recovery and existing lease/filter/fencing
  behavior.
• Own late-created runtimes through shutdown and wait for workers to stop before
  incompatible configuration rollback.
• While startup is deferred, propagate enabled-but-uninitialized Responses
  errors instead of silently falling back; intentional resilience opt-out and
  cold startup remain non-durable/in-process as before.

Compatibility and scope

Updated launcher contract: every temporary capture sandbox must set the
variable to true. Ordinary cold/local launchers may omit it or set false.
After-restore activates enabled tasks while the captured true value remains;
no flag override is required.

Coordinate rollout with service PR
https://msdata.visualstudio.com/Vienna/_git/vienna/pullrequest/2357422.
An unmarked capture would initialize enabled tasks eagerly. The earlier SDK
prototype treated absence as deferred, so service and SDK changes must be
reviewed together.

Platform opt-in is separately definition.enable_memory_snapshot, as proposed
by Azure/azure-rest-api-specs#46759. The SDK variable alone does not enable
platform snapshots.

No new HTTP 503 mapping, startup submission queue/replay, public session-start
callback, worker/lease policy, storage backend, or .NET change is included.
Q3 remains open: the spec offers an explicit session-start callback and a
startup launch-intent registry as unselected alternatives.

Controlled wheel labels/strict local-version pins are not published in this PR.
Responses requires the lifecycle-enabled Core version (>=2.3.0b2); no package
release is requested.

Validation

Environment rollback retry fix (October 20, 9f833dca90b4)

Retain the original environment/configuration rollback snapshots when worker
cleanup fails. After pending cleanup succeeds, restore those snapshots before
computing incoming overrides, including when cleanup completed during a
cancelled retry or host shutdown. Clear them only after rollback or successful
restore, so failed values cannot leak into omitted overrides or become the
captured baseline when an override is repeated.

• Expanded the cleanup fault matrix to 28 cases, checking omitted/repeated
  endpoint overrides and a later restore returning to the original baseline.
  Both omission/repetition regressions failed before this fix.
• 1,276 passed, 64 skipped across the full Core suite and Responses
  startup-failure contracts; all 80 activation cases pass.
• Black and Python 3.10 mypy pass. Pylint reports only the three unchanged
  logging diagnostics in _base.py; no new findings.
• This follow-up is locally qualified; the hosted campaign remains pinned to
  its earlier source.

Cleanup failure retry fix (October 20, c9e8e8139f5b)

Rollback retains reconciliation handles across transient provider read/claim
failures. A later restore finishes shielded cleanup under the original binding
before rehydrating configuration or invoking callbacks, then disposes the
stopped manager before starting a replacement. Ordinary shutdown preserves
pending cleanup rather than discarding queued work.

• 14 new fault-injection cases pass: read/claim failures once and twice,
  a claim committed before its error, same/new restore IDs and GUIDs,
  cancellation during cleanup retry, and ordinary shutdown.
• The promoted input executes once; the following queued input remains intact.
  Repeated successful hooks still reuse the activated runtime.
• 787 focused tests passed, 51 skipped, covering snapshot lifecycle, all
  Core task tests and Responses startup-failure contracts.
• After extracting shutdown-handle retention, 98 final focused tests passed,
  including all 66 snapshot activation cases.
• Black, Python 3.10 mypy and whitespace checks pass. Pylint has no new findings:
  four logging diagnostics reproduce unchanged on the previous head.

This follow-up has local fault-injection qualification; the hosted campaign
below remains pinned to its earlier source, not this cleanup-retry fix.

Previous full-suite baseline (October 19, 72d7822002d8)

Validated after merging current main, using Core, Responses and Invocations
source from this worktree (no foreign editable-package imports):

• Full Core: 1,230 passed, 64 skipped, 0 failed.
• Recovery/rollback records the restore ID before startup, owns workers until
  they actually finish even after their active task entry disappears, and joins
  them before rebinding configuration. A completed steering turn can durably
  stage the next queued input as recoverable without executing it during cleanup.
  Responses requires Core 2.3.0b2.
• Full Responses, with capture env true inherited: 2,309 passed, 83 skipped,
  0 failed. The previously reported sample-8 import failures disappeared when
  the sibling Invocations source was loaded from the correct worktree.
• 52 snapshot activation cases pass, including matching-binding retries,
  cooperative exits, completed handlers with blocked renewals, queued-input
  recovery after normal return, and original task/input recovery.
  39 existing cancellation/timeout cases pass.
• Sequential lifespans reset startup admission. Rollback reconciles a queued
  transition interrupted after turn suspension, claims and expires its lease,
  and resumes queued input on retry without executing it during cleanup.
• 18 Responses contract cases pass with the external snapshot variable
  unset and with it set to true. Four new scoped-runtime cleanup cases and
  per-test environment isolation failed before their fixes.
• Black, pylint and mypy (Python 3.10 target) pass for the changed code.

The startup truth table, restore callback ordering, ordinary shutdown and
explicit task cancellation are unchanged.

West US 2 retest on 72d7822002d8

The fresh r4 controlled-image campaign passed the full matrix on exact source
72d7822002d8a36d14c05f6eacf64bc7385c2e31, including the repeated-lifespan and
interrupted-drain fixes. Core 2.3.0b2, Responses 2.4.0b1 and Invocations
1.2.0 were packaged into fresh immutable images; staged-source, wheel and
runtime hashes matched:

• A: 7/7 passed: six resilience/env combinations plus an independent
  restored session; actual capture/restore state and loaded source hashes
  match; enabled jobs renew leases and disabled cases create no task records.
• B passed: stored sync/background, foreground/background streaming,
  non-durable store=false streaming, GET/input-items/DELETE, and cancellation
  with persisted cancelled state and no reappearance within 30 seconds.
• C passed: hard crash after checkpoints 1 and 2; the original task was
  autonomously reclaimed with generation 0->1 and a new fenced worker before
  GET/new ingress. Recovery read the response-store checkpoint, retained output
  IDs/content and executed the remaining phases.
• D passed: 360-second quiet interval after completed work, then one
  same-session invocation; new restore ID/worker, preserved heap and renewed lease.
• Cleanup: 18/18 owned sessions verified idle/stopped; zero failures.

The initial 3-second crash observer missed the short generation-1 reclaim window
before ephemeral task deletion. That evidence was retained; both crash cases
passed on fresh sessions with 250-ms polling and unchanged acceptance criteria.

This does not qualify active jobs crossing idle/freeze, pre-suspend cleanup,
automatic regional service env injection, or a startup launch registry. The
earlier r3 campaign on 16a16b1bae0a also passed, with 16/16 sessions cleaned;
r4 supplies hosted qualification for the latest follow-up fixes.

Earlier contract qualification

Updated capture-only contract (Python 3.14, source imports from this PR):

• Core: 1,021 passed, 59 skipped. An initial OTLP gRPC subprocess timeout
  passed two isolated reruns and the full rerun.
• Responses: 1,865 passed, 83 skipped, with the same two sample-8 import
  failures described below.
• Targeted suites: 31 Core and 13 Responses cases passed, including all six
  enabled/disabled x absent/false/true startup/restore combinations, retained
  capture flag with no restore override, and rejection of cross-session
  rebinding. Both streaming and non-streaming absent-mode fallbacks pass.
• Black with eng/black-pyproject.toml and pylint on changed production code
  pass.
• Coordinated service tests: 15 focused signal cases, 10,555 HostedAgents and
  env-constant tests, and 37 header guards pass.

Earlier West US 2 hosted SDK qualification: 7/7 passed, using immutable images
built from b213b1e71747 and verified loaded Core/Responses source-byte hashes.
All six enabled/disabled x absent/false/true combinations passed, plus a second
true-enabled restored session. Actual before-snapshot/after-restore evidence
showed no capture runtime and retained true; enabled runtime activation followed
restore callbacks. Independent sessions shared the golden heap marker but had
distinct restore/session IDs. Enabled jobs renewed backend leases under the
correct task/owner/native-session scope; disabled cases created no task records.
All seven owned sessions were stopped and verified idle.

The regional task API observations had empty heartbeat fields; renewal was
proved directly by increasing lease expiry for the same in-progress task.
The initial harness false-negative is preserved separately from the fresh
passing run. This qualifies the SDK with controlled-image signals, not
automatic injection by the service PR, which is not regionally deployed.

The table below records the earlier retarget validation, before the capture-only
contract update. Local, Python 3.14, clean environment (no FOUNDRY_MEMORY_SNAPSHOT_ENABLED),
compared against main (037e8b7) in a separate worktree:

Suite	main	This PR
Core tests	990 passed, 59 skipped	1016 passed, 59 skipped
Responses tests	1858 passed, 83 skipped, 2 failed	1862 passed, 83 skipped, 2 failed
Invocations tests	647 passed, 6 skipped	647 passed, 6 skipped
Optimization tests	197 passed	197 passed
Activity tests	95 passed, 11 failed, 1 skipped	95 passed, 11 failed, 1 skipped

The 2 Responses failures (test_sample8_mixin_composition_both_protocols and
test_sample_module_imports_and_registers[sample_08_mixin_composition.py]) and
the Activity failures are present on main in this environment and were not
investigated; azure-ai-agentserver-ghcopilot has the same 2 collection errors
on both.

• mypy (Python 3.10 target) on Core _base.py and tasks/_manager.py and
  Responses hosting/_orchestrator.py: no issues.
• pylint (repository pylintrc) on the same three files: no messages, as on
  main.
• git diff --check: clean.

The broader crash/idle/concurrency evidence below comes from an earlier
controlled build and was not re-run for this contract update. The separately
built controlled packages were exercised with the
existing resilient streaming sample in a hosted test:

• Cold and initial snapshot activation with real lease acquisition/renewal.
• Autonomous hard crash after checkpoints 1 and 2, without a new invocation
  during the orphan interval.
• Retained golden heap marker, new restore ID/worker, incremented fencing
  generation, checkpoint re-entry, and original response completion.
• Independent sessions, two concurrent tasks sharing one runtime, and
  cancellation without resurrection during the observation window.
• Natural idle/resume with a 120-second policy: after six quiet minutes,
  same logical session/capture with a new restore ID and worker, renewed
  leases, and completed outputs. Management still reported active; runtime
  evidence established restoration. Exact suspension timing was not measured.

Those packages predate the final lifecycle prerequisite changes and the revised
capture-only signal contract, so they support the design but do not verify this
exact commit. Startup task submissions remain untested. This is not release
qualification.

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).
10 pipeline(s) were filtered out due to trigger conditions.
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Local resilient samples remain incompatible with the new default, and the Responses test configuration is environment-dependent.

2 open findings

• Local resilient launchers omit the required snapshot environment variable · New
• Use unconditional environment setup to preserve the cold-start test contract · New

What changed in this PR

Integrates resilient task startup with memory-snapshot session activation while preserving cold-start behavior.

Changes:

• Defers task manager initialization until after restore in snapshot mode.
• Adds retry, recovery, shutdown, and rollback lifecycle handling.
• Updates Responses integration, tests, dependencies, and documentation.

File	Description
core/​_base.py	Integrates task activation into host lifecycle.
core/​tasks/​_manager.py	Adds phased startup and worker-aware shutdown.
core/​tests/​test_snapshot_task_activation.py	Tests snapshot activation and recovery.
core/​tests/​tasks/​test_task_manager_optin.py	Pins cold-start behavior in tests.
core/​README.md	Documents snapshot initialization.
core/​docs/​memory-snapshot-task-integration.md	Adds the design specification.
core/​CHANGELOG.md	Records lifecycle integration.
responses/​hosting/​_orchestrator.py	Distinguishes deferred startup from opt-out.
responses/​tests/​contract/​test_resilient_start_failure.py	Tests fallback and activation behavior.
responses/​tests/​conftest.py	Configures cold-start test defaults.
responses/​README.md	Documents compatibility requirements.
responses/​pyproject.toml	Raises the minimum Core dependency.
responses/​CHANGELOG.md	Records Responses behavior changes.

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Deferred-start failures can leave phantom in-progress responses, and the absent-variable behavior conflicts with the stated launcher contract.

3 open findings

• Clean up in-progress runtime before re-raising deferred activation errors · New
• Align absent-variable handling with the documented launcher contract · New
• Use unconditional environment setup to preserve the cold-start test contract

1 resolved since last review

• Local resilient launchers omit the required snapshot environment variable

🧠 Review effort: Balanced

[Copilot]

🟡 Changes recommended

Rollback can lose recoverable work, and deferred Responses failures leave stale records that block retries.

6 open issues

• Preserve recoverable task records during restore rollback · New
• Clean up in-progress records before propagating startup errors · New
• Clean up in-progress runtime before re-raising deferred activation errors
• Restore the session environment variable during test teardown · New
• Align absent-variable handling with the documented launcher contract
• Use unconditional environment setup to preserve the cold-start test contract

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Unresolved rollback, lease-identity, and minimum-dependency issues can compromise safe restoration and recovery.

4 open issues

• Record restore ID before recovery startup · New
• Preserve worker and lease handles during shutdown · New
• Prevent steering promotion during rollback cleanup · New
• Raise Core dependency floor for task deferral · New

6 resolved since last review

• Clean up in-progress records before propagating startup errors
• Preserve recoverable task records during restore rollback
• Clean up in-progress runtime before re-raising deferred activation errors
• Restore the session environment variable during test teardown
• Align absent-variable handling with the documented launcher contract
• Use unconditional environment setup to preserve the cold-start test contract

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Rollback can overlook running renewal workers and leave queued steering inputs excluded from recovery.

2 open issues

• Track and await renewal workers during rollback · New
• Preserve queued inputs when rollback shutdown completes normally · New

4 resolved since last review

• Raise Core dependency floor for task deferral
• Prevent steering promotion during rollback cleanup
• Preserve worker and lease handles during shutdown
• Record restore ID before recovery startup

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Unresolved lifecycle and rollback-recovery defects require fixes, and session-binding concurrency still warrants human review.

2 open issues

• Reset task manager shutdown gate on lifespan startup · New
• Reconcile queued input when rollback cancels steering drain · New

2 resolved since last review

• Preserve queued inputs when rollback shutdown completes normally
• Track and await renewal workers during rollback

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

A transient rollback reconciliation failure leaves a stopped manager installed, blocking subsequent restore retries.

1 open issue

• Restore retries reuse stopped TaskManager after transient failure · New

2 resolved since last review

• Reconcile queued input when rollback cancels steering drain
• Reset task manager shutdown gate on lifespan startup

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Intertwined restore and cancellation paths require human review, and cleanup retries still lose the original environment/configuration rollback snapshots.

1 open issue

• Preserve failed restore environment snapshots across cleanup retries · New

1 resolved since last review

• Restore retries reuse stopped TaskManager after transient failure

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Runtime reuse can retain stale provider configuration or acknowledge restoration without an initialized task manager.

1 open issue

• Rebuild inactive providers when restored configuration changes · New

2 previously missed issues

In code that hasn't changed since last review

Move task deferral feature entry from 2.3.0b1 to 2.3.0b2

sdk/​agentserver/​azure-ai-agentserver-core/​CHANGELOG.md:53

Move this feature entry to the unreleased 2.3.0b2 Features Added section. It currently advertises capture-time task deferral under 2.3.0b1, although this implementation ships in b2 and Responses explicitly requires b2. This could lead consumers to use b1 expecting the flag to prevent eager task initialization.

Reinitialize disposed runtime before acknowledging cached restores

sdk/​agentserver/​azure-ai-agentserver-core/​azure/​ai/​agentserver/​core/​_base.py:567

Runtime disposal leaves completed_restore_ids intact. Reuse a snapshot-enabled host in a second lifespan and send the same successful restore payload: startup defers tasks, then the completed-ID guard returns 200 without recreating the manager or recovery scanner. Resilient requests subsequently fail with TaskManagerNotInitialized. The sequential-lifespan test misses this because each cycle uses a different restore ID. Track runtime activation separately from callback completion, or re-establish the missing runtime before acknowledging cached restore success, without replaying completed callbacks. Add a same-payload lifespan regression.

1 resolved since last review

• Preserve failed restore environment snapshots across cleanup retries

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Cleanup failures bypass the standardized lifecycle error response.

3 open issues

• Move snapshot deferral entry to the 2.3.0b2 release · New
• Preserve rollback error response when cleanup fails · New
• Rebuild inactive providers when restored configuration changes

_{Review effort: Balanced}