FIX: Clean up Azure credential/client lifecycle to avoid unclosed sessions#2214
FIX: Clean up Azure credential/client lifecycle to avoid unclosed sessions#2214romanlutz wants to merge 2 commits into
Conversation
…sions Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 52d2518d-ebe5-4888-8189-6146b6503448
…ntial-lifecycle-cleanup # Conflicts: # doc/code/targets/2_openai_responses_target.ipynb
There was a problem hiding this comment.
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: |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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: |
There was a problem hiding this comment.
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.
| return | ||
| self._consumer_count -= 1 | ||
| if self._consumer_count == 0: | ||
| await self.close_async() |
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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
Description
PyRIT was leaking
aiohttpclient sessions and Azure credential connections, surfacing asUnclosed client session/Unclosed connectorwarnings, 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 exposingacquire(),release_async(), andclose_async().get_azure_async_token_providerandget_azure_openai_authnow 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.OpenAITargetacquires 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 tosuper()._cleanup_target_async()so the base lifecycle runs.AzureContentFilterScorergains an idempotentcleanup_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_asyncis 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_targetdoc/code/targets/2_openai_responses_targetdoc/code/scoring/2_float_scale_scorersNote:
doc/code/targets/realtime_targetis unchanged by this PR. It reuses a single target after callingcleanup_target_async(), which conflicts with the terminal-cleanup contract described above when running under Entra/token auth (it is fine under API-key auth).