Skip to content

FIX: Clean up Azure credential/client lifecycle to avoid unclosed sessions#2214

Open
romanlutz wants to merge 2 commits into
microsoft:mainfrom
romanlutz:romanlutz-azure-credential-lifecycle-cleanup
Open

FIX: Clean up Azure credential/client lifecycle to avoid unclosed sessions#2214
romanlutz wants to merge 2 commits into
microsoft:mainfrom
romanlutz:romanlutz-azure-credential-lifecycle-cleanup

Conversation

@romanlutz

Copy link
Copy Markdown
Contributor

Description

PyRIT was leaking aiohttp client sessions and Azure credential connections, surfacing as Unclosed client session / Unclosed connector warnings, because the shared Azure async token provider and the async OpenAI clients were never deterministically closed.

This introduces a managed lifecycle for the Azure async token provider and wires cleanup into the components that consume it:

  • AzureAsyncTokenProvider (pyrit/auth/azure_auth.py): a reference-counted wrapper around the Entra token credential exposing acquire(), release_async(), and close_async(). get_azure_async_token_provider and get_azure_openai_auth now return this managed provider instead of a bare callable. Because a single provider can be shared by multiple targets/scorers, the underlying credential is only closed once the last consumer releases it.
  • OpenAITarget acquires the managed provider on construction and gains a private _cleanup_target_async() that closes its async client and releases its hold on the provider.
  • OpenAIRealtimeTarget.cleanup_target_async() now chains to super()._cleanup_target_async() so the base lifecycle runs.
  • AzureContentFilterScorer gains an idempotent cleanup_scorer_async() that closes the content safety client and releases the managed provider.

Design note for reviewers: cleanup is terminal under Entra/token auth. Once the last consumer releases the shared provider it is closed, so a target or scorer should not be reused after its cleanup has run in that configuration. Under API-key auth there is no managed provider, so cleanup is a no-op and reuse is unaffected.

Tests and Documentation

Unit tests (added/updated, passing):

  • tests/unit/auth/test_azure_auth.py: token provider acquire/release/close lifecycle.
  • tests/unit/prompt_target/target/test_openai_target_auth.py: target cleanup closes the client and releases the provider.
  • tests/unit/score/test_azure_content_filter.py: cleanup_scorer_async is idempotent and releases the provider.
  • tests/integration/score/test_azure_content_filter_integration.py: integration coverage for scorer cleanup.

JupyText: re-executed the affected notebooks with uv run jupytext --execute --to notebook <file>; all produced clean output:

  • doc/code/targets/1_openai_chat_target
  • doc/code/targets/2_openai_responses_target
  • doc/code/scoring/2_float_scale_scorers

Note: doc/code/targets/realtime_target is unchanged by this PR. It reuses a single target after calling cleanup_target_async(), which conflicts with the terminal-cleanup contract described above when running under Entra/token auth (it is fine under API-key auth).

…sions

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 52d2518d-ebe5-4888-8189-6146b6503448
@behnam-o behnam-o self-assigned this Jul 16, 2026
…ntial-lifecycle-cleanup

# Conflicts:
#	doc/code/targets/2_openai_responses_target.ipynb

@behnam-o behnam-o left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The original helper creates an AsyncDefaultAzureCredential but returns only its bearer-token callback. This hides the credential—and therefore its lifecycle—from the caller, so it cannot be closed explicitly. Creating it inside async with does not solve this because returning the callback would immediately exit the context and close the credential before use.

The current PR addresses this with centralized reference counting, but introduces another implicit lifecycle protocol. Every target, scorer, wrapper, and failure path must correctly acquire, release, and propagate cleanup. As currently implemented, cleanup is not reachable through several public abstractions, and the reference-count transitions have race conditions.

A reasonable backward-compatible middle ground is to make AzureAsyncTokenProvider an async context manager:

async with get_azure_async_token_provider(scope) as provider:
    target = OpenAIChatTarget(api_key=provider)
    await run_attack_async(target=target)

Existing callers could continue using the helper without async with, while new code gains a supported way to explicitly scope and close the credential. This creates a migration path toward explicit ownership without an immediate breaking API change.

The context-manager lifecycle should be the preferred usage. Reference counting may remain temporarily for compatibility, but it must be concurrency-safe and should not replace public cleanup contracts for resource-owning targets and scorers.

if self._managed_token_provider:
self._managed_token_provider.acquire()

async def _cleanup_target_async(self) -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cleanup is unreachable for most OpenAI targets. OpenAITarget exposes cleanup only through private _cleanup_target_async(). Except for RealtimeTarget, normal target usage never invokes it, so clients and owned Azure credentials remain open. The unit test calls this method directly and does not validate real lifecycle behavior.

Please add a public, idempotent async cleanup contract to PromptTarget and propagate it through orchestration and composite targets.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I actually explicitly don't want that. I don't want people to have to think about this. We used to have some connection close instructions at the end of some notebooks and it was just annoying. Nobody wants to do that. I thought this might be a middle ground in the sense that it's theoretically possible but we don't expose it with a guarantee that we'll maintain it going forward but maybe this isn't well thought out.


super().__init__(validator=validator or self._DEFAULT_VALIDATOR)

async def cleanup_scorer_async(self) -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Scorer cleanup is hidden by wrappers. Cleanup exists only on AzureContentFilterScorer; Scorer has no cleanup contract, wrappers do not delegate cleanup, and registry reset discards instances without closing them. Registered wrapped Azure scorers therefore cannot be cleaned up through the public API.

Please add async cleanup to Scorer, propagate it through wrappers and composites, and dispose registered instances during async session shutdown.

Comment thread pyrit/auth/azure_auth.py
return
self._consumer_count -= 1
if self._consumer_count == 0:
await self.close_async()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reference counting has race conditions. On final release, the count reaches zero before credential closure completes. Another consumer can acquire during that await and receive a provider that is subsequently marked closed. Concurrent close_async() calls may close twice, while direct closure can invalidate active consumers.

Please synchronize transitions and model explicit OPEN, CLOSING, and CLOSED states. Prefer exception-safe leases, with closure performed only after the final lease is released.

@behnam-o behnam-o left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

assuming we do go with this new pattern for token providers, left some comments that I think should be addressed to make it work they way it is supposed to.

but I have a bigger comment on the core of this problem, and how it is caused by the way the "provider-helper" function (get_azure_async_token_provider) works today ... in it, I also share my thoughts on a good solution

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants