FIX: Make auth mode an explicit choice instead of an inferred one

[hannahwestra25]

Description

Selecting identity-based authentication did not actually force identity auth.

pyrit/backend/services/target_service.py signalled identity by deleting the api_key:

if request.auth_mode == "identity":
    if "identity" not in target_cls.supported_auth_modes:
        raise ValueError(...)
    params.pop("api_key", None)

But "no api_key passed" is ambiguous — it can mean "the user chose identity" or "read the key from the environment variable". Every auth resolver resolves that ambiguity toward the env var, so the user's explicit choice was silently downgraded to api-key auth whenever a key happened to be present in .env.

Reproduction (on main, prints the env key)

import os
os.environ["OPENAI_CHAT_API_KEY"] = "sk-SECRET-FROM-DOTENV"
from pyrit.auth.openai_auth import resolve_openai_auth
resolved = resolve_openai_auth(
    endpoint="https://my-resource.openai.azure.com/openai/v1",
    api_key=None,  # what target_service passes in identity mode
    api_key_environment_variable="OPENAI_CHAT_API_KEY",
)
print(resolved)  # -> 'sk-SECRET-FROM-DOTENV'  BUG: should be an Entra token provider

On this branch the same three cases now behave correctly:

default      -> str                                     # env key, unchanged behavior
identity     -> function | is env key? False            # Entra token provider
non-azure    -> ValueError: Identity-based authentication requires a recognized Azure ... endpoint

This is not a regression — it's two features that don't compose

pyrit/auth/openai_auth.py was introduced in #2235 purely as a fallback: when no key exists for a recognized Azure endpoint, mint an Entra token rather than fail. That was correct for its purpose. The gap only appeared once the backend let users explicitly choose identity: a fallback mechanism has no way to distinguish "no key available" from "key available but deliberately not wanted."

Affected paths

Path	Resolver
6 OpenAI targets	pyrit/auth/openai_auth.py::resolve_openai_auth — the if api_key_value: return api_key_value env check ran before the is_azure_openai_endpoint Entra fallback
OpenAITextEmbedding	same resolver
AzureMLChatTarget	same chain inlined, is_azure_ml_endpoint
PromptShieldTarget	same chain inlined, is_azure_openai_endpoint

The fix

Make auth intent explicit rather than inferred, in two parts.

1. auth_mode is threaded down to the resolvers, so identity skips the key and env var entirely:

def resolve_openai_auth(*, endpoint, api_key, api_key_environment_variable, auth_mode="api_key"):
    if auth_mode == "identity":
        if not is_azure_openai_endpoint(endpoint):
            raise ValueError(...)
        return get_azure_openai_auth(endpoint)
    ...

auth_mode is plumbed through openai_target.py, azure_ml_chat_target.py, prompt_shield_target.py and openai_text_embedding.py, and target_service.create_target_async passes it explicitly instead of relying solely on popping api_key.

2. ⚠️ The implicit Entra fallback is removed (see Breaking change below). Fixing (1) alone leaves identity and "no key configured" still resolving down the same path, which is the root of the ambiguity.

Import cycle

AuthMode was defined at pyrit/prompt_target/common/prompt_target.py, but pyrit/auth/ may only import from pyrit.auth and pyrit.common (prompt_target depends on auth, not the reverse). The canonical definition moves to the new neutral pyrit/common/auth_mode.py and is re-exported from pyrit.prompt_target.common.prompt_target, so every existing import keeps working. Verified with a clean python -c "import pyrit.prompt_target".

pyrit/cli/_auth.py::AuthMode is a different Literal ("auto" | "azure_cli" | "device_code" | "none") and is untouched.

⚠️ Breaking change: the implicit Entra fallback is removed

Previously, when api_key mode found no key and the endpoint looked like an Azure host, the resolvers silently minted an Entra token. That fallback is deleted.

Why it had to go. It is exactly what makes an explicit auth_mode="identity" indistinguishable from "no key configured" — the two requests converge on the same resolution path, so the resolver cannot honor the explicit one. Keeping it would mean keeping a third, unnamed auth state ("infer from what happens to be in the environment") alongside the two named ones.

New contract:

• auth_mode="api_key" (the default) requires a key: a token-provider callable passed as api_key, an explicit api_key string, or the target's API key environment variable. If none yields a key it raises ValueError rather than guessing.
• auth_mode="identity" ignores keys and the environment variable entirely, and raises ValueError for an endpoint that is not a recognized Azure host — it never mints a token for an unknown host.

Who is affected. Only callers who relied on OpenAIChatTarget(endpoint=<azure>) + az login with no key configured. The migration is one keyword argument, and the error message names it:

No API key available for endpoint 'https://...'. Set the OPENAI_CHAT_API_KEY environment
variable, pass api_key explicitly, or pass auth_mode="identity" to authenticate with
Microsoft Entra ID on a recognized Azure OpenAI / AI Foundry endpoint.

A dedicated test (test_api_key_mode_error_names_the_identity_migration) pins that wording so the migration hint cannot silently regress.

The blast radius is narrower than it looks. Every Entra example in the repo already passes an explicit token provider rather than relying on the fallback — doc/code/targets/1_openai_chat_target.py, 2_openai_responses_target.py, realtime_target.py, prompt_shield_target.py, doc/code/scoring/2_float_scale_scorers.py, doc/code/setup/1_configuration.py — as does TargetInitializer (pyrit/setup/initializers/targets.py). The implicit path was never the documented one. This supersedes the deprecation cycle previously tracked in #3029.

OpenAITextEmbedding is included deliberately. It is a fourth consumer of resolve_openai_auth and would otherwise have lost identity auth with no replacement other than hand-constructing a provider, so it gains auth_mode alongside the targets. This keeps the resolver contract uniform across all of its callers.

Relationship to #2846 (merged)

#2846 has landed on main (799ea49b2) and this branch is rebased onto it.

#2846 fixed the AzureBlobStorageTarget half of this bug — ABS's credential is sas_token, so pop("api_key") was a no-op — by adding a get_auth_mode_parameters classmethod hook that lets a target declare the constructor parameters implied by a chosen auth mode, plus identity_conflicting parameter metadata so the service can strip any flagged credential rather than a hardcoded api_key.

It did not fix the resolver half. pyrit/auth/openai_auth.py is absent from its diff, its base get_auth_mode_parameters returns {}, and only AzureBlobStorageTarget overrode it — so the reproduction above still returned the env key on 799ea49b2.

Rather than leave two competing mechanisms in the tree, this PR adopts #2846's hook as the single channel and deletes the _accepts_auth_mode capability probe this branch previously used.

Tests

tests/unit/auth/test_openai_auth.py plus identity tests in the AzureML, PromptShield, OpenAI-target-auth, embedding and target-service suites:

• Regression (fails without the fix): with the api-key env var set, auth_mode="identity" resolves to a token provider, not the env key — covered for all four resolver paths.
• auth_mode="api_key" and the default preserve the key chain exactly (callable precedence, explicit-key precedence, env var).
• Identity against a non-Azure endpoint raises a clear ValueError.
• api_key mode with no key available raises, and get_azure_openai_auth is asserted not called — the fallback cannot creep back in.
• The migration hint is asserted verbatim.
• target_service: an auth_mode="identity" create request produces a target that does not use the env key (OpenAI and AzureML), and an api_key request still honors the env var.
• A registry contract test (test_identity_targets_accept_an_explicit_auth_mode) asserts that every registered target advertising "identity" overrides get_auth_mode_parameters, so a future identity-capable target cannot silently reintroduce the bug. Verified non-vacuous by temporarily removing an override and confirming the test fails.

Ten existing tests that encoded the old fallback behavior were rewritten to assert the new contract (e.g. test_no_key_recognized_azure_endpoint_auto_mints_entra → ..._raises). These are intentional contract updates, not coverage removal — each asserts the replacement behavior at the same call site.

Review follow-ups

params.auth_mode conflicts (390c1f9). Because auth_mode is now a constructor parameter, the registry also accepted it inside params as a second, competing channel. A request with top-level auth_mode="api_key" and params["auth_mode"]="identity" selected identity, silently ignoring a supplied api_key and bypassing the service's supported_auth_modes gate (that check only ran inside the identity branch). Conflicting values are now rejected, and the request-level auth_mode is forwarded for both modes so it is authoritative in either direction.

This check is deliberately retained on top of #2846: that PR applies the hook with params.update(...), which overwrites a smuggled params["auth_mode"] rather than rejecting it, so the conflicting request would otherwise still be accepted silently.

Docs (c312db1, da72cf5). doc/code/setup/1_configuration documents the explicit auth_mode="identity", the resolution order, and a migration note covering the removed fallback. Both the .py and .ipynb were updated and verified in sync. Stale "Entra ID authentication is used automatically" docstrings on OpenAITarget and OpenAITextEmbedding are corrected.

ty redundant cast (19f2f80). Typing AUTH_MODES as tuple[AuthMode, ...] lets ty narrow through the membership check, making the cast("AuthMode", ...) in _get_supported_auth_modes redundant and failing the ty (type check) hook. The cast and its now-unused typing.cast import are removed; the narrowing is equivalent.

Remaining loose end. #2846 types ABS's parameter as AuthMode | None = None while this PR uses AuthMode = "api_key". Both spell the "infer from credentials" third state differently. Now that the implicit fallback is gone on the resolver side, ABS can be unified to the same spelling — left out of this PR because it would reverse a just-merged design decision without the reviewer's sign-off, and is better done as a focused follow-up.

Validation

Local commands run on the final branch (da72cf550, on top of 799ea49b2):

Command	Result
pytest tests/unit/auth/ tests/unit/backend/test_target_service.py tests/unit/prompt_target/ tests/unit/registry/test_target_registry.py -q	1815 passed, 199 skipped
pytest tests/unit/auth/test_openai_auth.py tests/unit/prompt_target/target/test_openai_chat_target.py tests/unit/prompt_target/target/test_openai_target_auth.py tests/unit/prompt_target/target/test_azure_ml_chat_target.py tests/unit/prompt_target/target/test_prompt_shield_target.py -q	181 passed
pytest tests/unit/embedding -q	13 passed
ruff check pyrit tests/unit	All checks passed
ty check pyrit	no diagnostics in any file changed by this PR
python build_scripts/validate_docs.py	passed
jupytext round-trip of doc/code/setup/1_configuration	.py and .ipynb in sync
python -c "import pyrit.prompt_target"	OK (no import cycle)
Reproduction script above	identity returns a token provider, not the env key

Full-suite validation is from CI, which is green on da72cf550 across every job (all platforms × Python 3.11–3.14, dev and dev_all, plus coverage, docs and frontend). The local full-suite run was not used as evidence: this environment is missing several optional dependencies (mcp, httpx2, setuptools, feedgen), so six modules fail to collect locally on main as well — a pre-existing environment gap, not a property of this branch. CI installs the full --extra all set and exercises those modules.

A previously reported local failure, tests/unit/backend/test_scenario_run_routes.py::TestResumeScenarioRunRoute::test_resume_has_no_get_preflight, was confirmed pre-existing: it touches no auth or target code, passes in isolation, and reproduces identically in a detached worktree at clean main (799ea49b2) with none of this branch's changes present.

[Richard Lundeen (richlundeen)]

Take a look at my comments. None are blockers but recommendations; happy to chat about any further!

[Richard Lundeen (richlundeen)]

I recommend removing the legacy fallback for simplicity. I do not think the compatibility ambiguity is worth it: api_key should require a key or an explicitly supplied provider, and identity should ignore keys and use identity. Please remove automatic identity fallback from api_key mode across the affected targets, rather than add an auto or omitted-mode compatibility path. I would also remove the service's _accepts_auth_mode fallback: every target that advertises identity support should accept the explicit mode, and a target that cannot meet that contract should fail clearly. Please document the compatibility change. (Drafted with GitHub Copilot.)

[hannahwestra25]

Agreed on the end state, but it's a silent breaking change: OpenAIChatTarget(endpoint=<azure>) + az login works and is documented today.

Two notes: no auto/omitted-mode path was added — that fallback is pre-existing and unreordered. And _accepts_auth_mode exists solely because AzureBlobStorageTarget lacks auth_mode; it deletes itself once #2846 lands.

Filed #3029 for the deprecation path. If you'd rather take the break now in this PR, say so — it's a small change.

[Richard Lundeen (richlundeen)]

Yes, please take the compatibility break now in this PR rather than use the deprecation path in #3029. My recommendation remains simplicity: I do not think the compatibility ambiguity is worth it. api_key should require a key or an explicitly supplied provider and fail clearly when neither is available; identity should ignore keys and use identity. Please apply that rule to OpenAI, Azure ML, Prompt Shield, and Blob Storage, without an auto or omitted-mode compatibility path, and document the change and the auth_mode="identity" migration.

I checked the current head (19f2f8031). The conflicting-input fix is addressed, Blob Storage now preserves identity intent and bypasses SAS credentials, and _accepts_auth_mode is gone. The remaining item in this thread is removal of the implicit identity fallback. I understand that the fallback predates this PR; the recommendation is to remove it now, not to treat it as a newly introduced defect. Code changes and merging remain with you.

(Drafted with GitHub Copilot.)