From 7772d77167366d50886593b7ae8add02c3bd2f53 Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Tue, 19 May 2026 14:35:28 -0700
Subject: [PATCH 1/8] adding files to add entra auth option to targets

---
 .../Config/CreateTargetDialog.test.tsx        | 148 +++++++++++++++
 .../components/Config/CreateTargetDialog.tsx  |  72 +++++++-
 frontend/src/types/index.ts                   |   1 +
 pyrit/backend/models/targets.py               |  10 +-
 pyrit/backend/services/target_service.py      | 109 ++++++++++-
 pyrit/prompt_target/azure_ml_chat_target.py   |  85 +++++++--
 tests/unit/backend/test_target_service.py     | 171 ++++++++++++++++++
 .../target/test_azure_ml_chat_target.py       | 126 +++++++++++++
 8 files changed, 693 insertions(+), 29 deletions(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx
index b8a9dcfc0c..70966a9a9b 100644
--- a/frontend/src/components/Config/CreateTargetDialog.test.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx
@@ -497,4 +497,152 @@ describe("CreateTargetDialog", () => {
 
     expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
   });
+
+  it("should hide the API Key field and omit api_key/include auth_mode when Entra is selected", async () => {
+    const onCreated = jest.fn();
+    const user = userEvent.setup();
+    mockedTargetsApi.createTarget.mockResolvedValue({
+      target_registry_name: "openai_chat_entra",
+      target_type: "OpenAIChatTarget",
+    });
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} onCreated={onCreated} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "OpenAIChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-resource.openai.azure.com/"
+    );
+    fireEvent.change(endpointInput, {
+      target: { value: "https://my-resource.openai.azure.com/" },
+    });
+
+    // API Key field is visible by default.
+    expect(
+      screen.getByPlaceholderText("API key (stored in memory only)")
+    ).toBeInTheDocument();
+
+    // Select Entra radio.
+    await user.click(
+      screen.getByRole("radio", {
+        name: /Microsoft Entra ID/,
+      })
+    );
+
+    // API Key field is hidden in Entra mode.
+    expect(
+      screen.queryByPlaceholderText("API key (stored in memory only)")
+    ).not.toBeInTheDocument();
+
+    await user.click(screen.getByText("Create Target"));
+
+    await waitFor(() => {
+      expect(mockedTargetsApi.createTarget).toHaveBeenCalledWith({
+        type: "OpenAIChatTarget",
+        params: {
+          endpoint: "https://my-resource.openai.azure.com/",
+        },
+        auth_mode: "entra",
+      });
+      expect(onCreated).toHaveBeenCalled();
+    });
+  });
+
+  it("should clear a previously-typed API key when switching to Entra", async () => {
+    const onCreated = jest.fn();
+    const user = userEvent.setup();
+    mockedTargetsApi.createTarget.mockResolvedValue({
+      target_registry_name: "openai_chat_entra",
+      target_type: "OpenAIChatTarget",
+    });
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} onCreated={onCreated} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "OpenAIChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-resource.openai.azure.com/"
+    );
+    fireEvent.change(endpointInput, {
+      target: { value: "https://my-resource.openai.azure.com/" },
+    });
+
+    // Type a key, then switch to Entra.
+    fireEvent.change(
+      screen.getByPlaceholderText("API key (stored in memory only)"),
+      { target: { value: "sk-typed-before-switch" } }
+    );
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra ID/ })
+    );
+
+    await user.click(screen.getByText("Create Target"));
+
+    await waitFor(() => {
+      const call = mockedTargetsApi.createTarget.mock.calls[0][0];
+      expect(call.auth_mode).toBe("entra");
+      expect(call.params).not.toHaveProperty("api_key");
+    });
+  });
+
+  it("should warn the user when Entra is selected for a non-Azure OpenAI endpoint", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "OpenAIChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-resource.openai.azure.com/"
+    );
+    fireEvent.change(endpointInput, { target: { value: "https://api.openai.com/" } });
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra ID/ })
+    );
+
+    expect(
+      screen.getByText(/Microsoft Entra ID authentication is only supported for Azure endpoints/)
+    ).toBeInTheDocument();
+  });
+
+  it("should NOT warn when Entra is selected for a recognized Azure endpoint", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "OpenAIChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-resource.openai.azure.com/"
+    );
+    fireEvent.change(endpointInput, {
+      target: { value: "https://my-resource.openai.azure.com/" },
+    });
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra ID/ })
+    );
+
+    expect(
+      screen.queryByText(/only supported for Azure endpoints/)
+    ).not.toBeInTheDocument();
+  });
 });
diff --git a/frontend/src/components/Config/CreateTargetDialog.tsx b/frontend/src/components/Config/CreateTargetDialog.tsx
index 4bce08f563..c8b2870fa0 100644
--- a/frontend/src/components/Config/CreateTargetDialog.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.tsx
@@ -9,6 +9,8 @@ import {
   Button,
   Input,
   Label,
+  Radio,
+  RadioGroup,
   Select,
   Switch,
   Text,
@@ -32,6 +34,26 @@ const TARGET_TYPE_CONFIG: Record<string, 'openai' | 'azureml'> = {
 
 const SUPPORTED_TARGET_TYPES = Object.keys(TARGET_TYPE_CONFIG)
 
+type AuthMode = 'api_key' | 'entra'
+
+// Mirrors the backend's strict hostname-suffix check (target_service.py).
+// Used to warn the user before they submit; the backend remains the source of truth.
+const AZURE_OPENAI_HOSTNAME_SUFFIXES = [
+  '.openai.azure.com',
+  '.ai.azure.com',
+  '.services.ai.azure.com',
+  '.cognitiveservices.azure.com',
+]
+
+function isAzureOpenAiEndpoint(endpoint: string): boolean {
+  try {
+    const host = new URL(endpoint).hostname.toLowerCase()
+    return AZURE_OPENAI_HOSTNAME_SUFFIXES.some((s) => host.endsWith(s))
+  } catch {
+    return false
+  }
+}
+
 interface CreateTargetDialogProps {
   open: boolean
   onClose: () => void
@@ -45,6 +67,7 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
   const [modelName, setModelName] = useState('')
   const [hasDifferentUnderlying, setHasDifferentUnderlying] = useState(false)
   const [underlyingModel, setUnderlyingModel] = useState('')
+  const [authMode, setAuthMode] = useState<AuthMode>('api_key')
   const [apiKey, setApiKey] = useState('')
   const [maxNewTokens, setMaxNewTokens] = useState('400')
   const [temperature, setTemperature] = useState('1.0')
@@ -54,7 +77,11 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
   const [error, setError] = useState<string | null>(null)
   const [fieldErrors, setFieldErrors] = useState<{ targetType?: string; endpoint?: string }>({})
 
-  const isAzureML = TARGET_TYPE_CONFIG[targetType] === 'azureml'
+  const targetKind = TARGET_TYPE_CONFIG[targetType]
+  const isAzureML = targetKind === 'azureml'
+  const isOpenAi = targetKind === 'openai'
+  const isEntra = authMode === 'entra'
+  const showNonAzureEntraWarning = isEntra && isOpenAi && endpoint !== '' && !isAzureOpenAiEndpoint(endpoint)
 
   const resetForm = () => {
     setTargetType('')
@@ -62,6 +89,7 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
     setModelName('')
     setHasDifferentUnderlying(false)
     setUnderlyingModel('')
+    setAuthMode('api_key')
     setApiKey('')
     setMaxNewTokens('400')
     setTemperature('1.0')
@@ -94,7 +122,7 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
         endpoint,
       }
       if (modelName) params.model_name = modelName
-      if (apiKey) params.api_key = apiKey
+      if (!isEntra && apiKey) params.api_key = apiKey
 
       if (hasDifferentUnderlying && underlyingModel) params.underlying_model = underlyingModel
 
@@ -112,6 +140,7 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
       await targetsApi.createTarget({
         type: targetType,
         params,
+        ...(isEntra ? { auth_mode: 'entra' as const } : {}),
       })
       resetForm()
       onCreated()
@@ -243,15 +272,40 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
                 </>
               )}
 
-              <Field label="API Key">
-                <Input
-                  type="password"
-                  placeholder="API key (stored in memory only)"
-                  value={apiKey}
-                  onChange={(_, data) => setApiKey(data.value)}
-                />
+              <Field label="Authentication">
+                <RadioGroup
+                  value={authMode}
+                  onChange={(_, data) => {
+                    const next = data.value as AuthMode
+                    setAuthMode(next)
+                    if (next === 'entra') setApiKey('')
+                  }}
+                >
+                  <Radio value="api_key" label="API Key" />
+                  <Radio value="entra" label="Microsoft Entra ID (DefaultAzureCredential)" />
+                </RadioGroup>
               </Field>
 
+              {showNonAzureEntraWarning && (
+                <MessageBar intent="warning">
+                  <MessageBarBody>
+                    Microsoft Entra ID authentication is only supported for Azure endpoints
+                    (*.openai.azure.com or *.ai.azure.com). This request will be rejected by the server.
+                  </MessageBarBody>
+                </MessageBar>
+              )}
+
+              {!isEntra && (
+                <Field label="API Key">
+                  <Input
+                    type="password"
+                    placeholder="API key (stored in memory only)"
+                    value={apiKey}
+                    onChange={(_, data) => setApiKey(data.value)}
+                  />
+                </Field>
+              )}
+
               <Label size="small" style={{ color: tokens.colorNeutralForeground3 }}>
                 Targets can also be auto-populated by adding an initializer (e.g. <code>airt</code>) to your{' '}
                 <code>~/.pyrit/.pyrit_conf</code> file, which reads endpoints from your <code>.env</code> and{' '}
diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts
index 2f72d902ce..6881f2e7da 100644
--- a/frontend/src/types/index.ts
+++ b/frontend/src/types/index.ts
@@ -85,6 +85,7 @@ export interface TargetListResponse {
 export interface CreateTargetRequest {
   type: string
   params: Record<string, unknown>
+  auth_mode?: 'api_key' | 'entra'
 }
 
 // --- Converters ---
diff --git a/pyrit/backend/models/targets.py b/pyrit/backend/models/targets.py
index e2eca8d911..1bd319c117 100644
--- a/pyrit/backend/models/targets.py
+++ b/pyrit/backend/models/targets.py
@@ -11,7 +11,7 @@
 This module defines the Instance models for runtime target management.
 """
 
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -80,3 +80,11 @@ class CreateTargetRequest(BaseModel):
 
     type: str = Field(..., description="Target type (e.g., 'OpenAIChatTarget')")
     params: dict[str, Any] = Field(default_factory=dict, description="Target constructor parameters")
+    auth_mode: Literal["api_key", "entra"] = Field(
+        "api_key",
+        description=(
+            "Authentication mode. 'api_key' uses the api_key in params (default, backwards-compatible). "
+            "'entra' uses Microsoft Entra ID (DefaultAzureCredential); requires an Azure endpoint and is "
+            "supported by OpenAI-family targets and AzureMLChatTarget."
+        ),
+    )
diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py
index af058dc2d9..c51ec9a611 100644
--- a/pyrit/backend/services/target_service.py
+++ b/pyrit/backend/services/target_service.py
@@ -12,10 +12,13 @@
 - Retrieved from registry (pre-registered at startup or created earlier)
 """
 
+import logging
 from functools import lru_cache
 from typing import Any
+from urllib.parse import urlparse
 
 from pyrit import prompt_target
+from pyrit.auth import get_azure_async_token_provider, get_azure_openai_auth
 from pyrit.backend.mappers.target_mappers import target_object_to_instance
 from pyrit.backend.models.common import PaginationInfo
 from pyrit.backend.models.targets import (
@@ -24,8 +27,44 @@
     TargetListResponse,
 )
 from pyrit.prompt_target import PromptTarget
+from pyrit.prompt_target.azure_ml_chat_target import AzureMLChatTarget
+from pyrit.prompt_target.openai.openai_target import OpenAITarget
 from pyrit.registry.object_registries import TargetRegistry
 
+logger = logging.getLogger(__name__)
+
+# Scope for Azure Machine Learning managed online endpoints.
+_AZURE_ML_SCOPE = "https://ml.azure.com/.default"
+
+# Recognised Azure OpenAI / AI Foundry hostname suffixes. Used for strict
+# endpoint validation when Entra ID auth is requested, so a bearer token is
+# only ever issued for a known Microsoft-operated endpoint.
+_AZURE_OPENAI_HOSTNAME_SUFFIXES = (
+    ".openai.azure.com",
+    ".ai.azure.com",
+    ".services.ai.azure.com",
+    ".cognitiveservices.azure.com",
+)
+
+
+def _is_azure_openai_endpoint(endpoint: str) -> bool:
+    """
+    Return True if ``endpoint`` resolves to a known Azure OpenAI / AI Foundry host.
+
+    Strict hostname-suffix check (not a substring search) so a bearer token is
+    never issued for an attacker-controlled endpoint whose URL merely contains
+    the word "azure".
+
+    Args:
+        endpoint (str): The endpoint URL to validate.
+
+    Returns:
+        bool: True if the endpoint's hostname ends with a recognised Azure OpenAI /
+        AI Foundry suffix; False otherwise (including for malformed URLs).
+    """
+    hostname = (urlparse(endpoint).hostname or "").lower()
+    return any(hostname.endswith(suffix) for suffix in _AZURE_OPENAI_HOSTNAME_SUFFIXES)
+
 
 def _build_target_class_registry() -> dict[str, type]:
     """
@@ -165,24 +204,84 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn
         Instantiates the target with the given type and params,
         then registers it in the registry under its registry name.
 
+        When ``request.auth_mode == "entra"``, an Azure token-provider callable is
+        injected as ``api_key`` before instantiation. Any ``api_key`` value in the
+        original ``params`` is discarded in that case. Entra ID is supported for
+        OpenAI-family targets (against Azure endpoints) and ``AzureMLChatTarget``.
+
         Args:
-            request: The create target request with type and params.
+            request: The create target request with type, params, and auth_mode.
 
         Returns:
             TargetInstance with the new target's details.
 
         Raises:
-            ValueError: If the target type is not found.
+            ValueError: If the target type is not found, or if Entra ID is requested
+                for an unsupported target type, or if Entra ID is requested for an
+                OpenAI target against a non-Azure endpoint.
         """
-        # Instantiate from request params and register (uses unique_name as key by default)
         target_class = self._get_target_class(target_type=request.type)
-        target_obj = target_class(**request.params)
+
+        # Build a fresh params dict so we never mutate the incoming Pydantic model.
+        params: dict[str, Any] = dict(request.params)
+
+        if request.auth_mode == "entra":
+            params = self._apply_entra_auth(target_class=target_class, target_type=request.type, params=params)
+
+        target_obj = target_class(**params)
         self._registry.register_instance(target_obj)
 
-        # Build response from the registered instance
         target_registry_name = target_obj.get_identifier().unique_name
         return self._build_instance_from_object(target_registry_name=target_registry_name, target_obj=target_obj)
 
+    @staticmethod
+    def _apply_entra_auth(*, target_class: type, target_type: str, params: dict[str, Any]) -> dict[str, Any]:
+        """
+        Replace any ``api_key`` in ``params`` with an Entra ID token provider for
+        the given target class.
+
+        Args:
+            target_class (type): The target class being instantiated. Used to select the
+                correct token-provider scope.
+            target_type (str): The user-facing target type name (used only in error messages).
+            params (dict[str, Any]): The target constructor parameters from the request.
+                This dict is not mutated.
+
+        Returns:
+            dict[str, Any]: A new params dict with ``api_key`` replaced by an async
+            token-provider callable suitable for the target class.
+
+        Raises:
+            ValueError: If the target type does not support Entra ID, or if an
+                OpenAI target is given a non-Azure endpoint.
+        """
+        new_params = dict(params)
+        if "api_key" in new_params:
+            # User error: api_key isn't used in Entra mode. Don't log the value.
+            logger.debug("Discarding 'api_key' from params because auth_mode='entra'.")
+            new_params.pop("api_key", None)
+
+        if issubclass(target_class, OpenAITarget):
+            endpoint = new_params.get("endpoint")
+            if not isinstance(endpoint, str) or not endpoint:
+                raise ValueError("Entra ID authentication requires an 'endpoint' in params.")
+            if not _is_azure_openai_endpoint(endpoint):
+                raise ValueError(
+                    "Entra ID authentication requires an Azure endpoint "
+                    f"(*.openai.azure.com or *.ai.azure.com). Got: {endpoint}"
+                )
+            new_params["api_key"] = get_azure_openai_auth(endpoint)
+            return new_params
+
+        if issubclass(target_class, AzureMLChatTarget):
+            new_params["api_key"] = get_azure_async_token_provider(_AZURE_ML_SCOPE)
+            return new_params
+
+        raise ValueError(
+            f"Target type '{target_type}' does not support Entra ID authentication. "
+            "Supported types are OpenAI-family targets and AzureMLChatTarget."
+        )
+
 
 @lru_cache(maxsize=1)
 def get_target_service() -> TargetService:
diff --git a/pyrit/prompt_target/azure_ml_chat_target.py b/pyrit/prompt_target/azure_ml_chat_target.py
index be4e463779..a3a63d8f65 100644
--- a/pyrit/prompt_target/azure_ml_chat_target.py
+++ b/pyrit/prompt_target/azure_ml_chat_target.py
@@ -3,10 +3,12 @@
 
 import logging
 import warnings
-from typing import Any
+from collections.abc import Awaitable, Callable
+from typing import Any, cast
 
 from httpx import HTTPStatusError
 
+from pyrit.auth import ensure_async_token_provider
 from pyrit.common import default_values, net_utility
 from pyrit.exceptions import (
     EmptyResponseException,
@@ -61,7 +63,7 @@ def __init__(
         self,
         *,
         endpoint: str | None = None,
-        api_key: str | None = None,
+        api_key: str | Callable[[], str | Awaitable[str]] | None = None,
         model_name: str = "",
         message_normalizer: MessageListNormalizer[Any] | None = None,
         max_new_tokens: int = 400,
@@ -79,8 +81,12 @@ def __init__(
         Args:
             endpoint (str | None): The endpoint URL for the deployed Azure ML model.
                 Defaults to the value of the AZURE_ML_MANAGED_ENDPOINT environment variable.
-            api_key (str | None): The API key for accessing the Azure ML endpoint.
-                Defaults to the value of the `AZURE_ML_KEY` environment variable.
+            api_key (str | Callable[[], str | Awaitable[str]] | None): The API key for accessing
+                the Azure ML endpoint, or a callable that returns a bearer token (sync or async).
+                Pass a token provider (e.g. ``get_azure_async_token_provider("https://ml.azure.com/.default")``)
+                to authenticate with Microsoft Entra ID against an AML managed online endpoint.
+                Synchronous providers are automatically wrapped via ``ensure_async_token_provider``.
+                Defaults to the value of the ``AZURE_ML_KEY`` environment variable.
             model_name (str): The name of the model being used (e.g., "Llama-3.2-3B-Instruct").
                 Used for identification purposes. Defaults to empty string.
             message_normalizer (MessageListNormalizer[Any] | None): **Deprecated.** Use
@@ -185,22 +191,45 @@ def _build_identifier(self) -> ComponentIdentifier:
             },
         )
 
-    def _initialize_vars(self, endpoint: str | None = None, api_key: str | None = None) -> None:
+    def _initialize_vars(
+        self,
+        endpoint: str | None = None,
+        api_key: str | Callable[[], str | Awaitable[str]] | None = None,
+    ) -> None:
         """
         Set the endpoint and key for accessing the Azure ML model. Use this function to manually
         pass in your own endpoint uri and api key. Defaults to the values in the .env file for the variables
         stored in self.endpoint_uri_environment_variable and self.api_key_environment_variable (which default to
-        "AZURE_ML_MANAGED_ENDPOINT" and "AZURE_ML_KEY" respectively). It is recommended to set these variables
-        in the .env file and call _set_env_configuration_vars rather than passing the uri and key directly to
-        this function or the target constructor.
+        "AZURE_ML_MANAGED_ENDPOINT" and "AZURE_ML_KEY" respectively).
+
+        If ``api_key`` is a callable (synchronous or asynchronous), it is treated as an Entra ID
+        token provider. The callable is stored on ``self._api_key_provider`` and resolved per-request
+        inside ``_get_headers_async``. Synchronous providers are wrapped via
+        ``ensure_async_token_provider``. In callable mode the ``AZURE_ML_KEY`` environment variable
+        is intentionally NOT consulted, so the caller-provided provider always wins.
 
         Args:
-            endpoint (str, optional): The endpoint uri for the deployed Azure ML model.
-            api_key (str, optional): The API key for accessing the Azure ML endpoint.
+            endpoint (str | None): The endpoint uri for the deployed Azure ML model.
+            api_key (str | Callable[[], str | Awaitable[str]] | None): A static API key string,
+                an async/sync callable returning a bearer token, or None to fall back to the
+                ``AZURE_ML_KEY`` environment variable.
         """
         self._endpoint = default_values.get_required_value(
             env_var_name=self.endpoint_uri_environment_variable, passed_value=endpoint
         )
+
+        if callable(api_key):
+            normalized = ensure_async_token_provider(api_key)
+            # ``ensure_async_token_provider`` returns the input unchanged when it is None
+            # or a string and returns a callable when the input is callable. We already
+            # verified the input is callable above, so narrow the return type for ty.
+            assert callable(normalized), "ensure_async_token_provider must return a callable for callable input"
+            provider = cast("Callable[[], Awaitable[str]]", normalized)
+            self._api_key_provider: Callable[[], Awaitable[str]] | None = provider
+            self._api_key = ""
+            return
+
+        self._api_key_provider = None
         self._api_key = default_values.get_required_value(
             env_var_name=self.api_key_environment_variable, passed_value=api_key
         )
@@ -269,7 +298,7 @@ async def _complete_chat_async(
         Returns:
             str: The generated response message.
         """
-        headers = self._get_headers()
+        headers = await self._get_headers_async()
         payload = await self._construct_http_body_async(messages)
 
         response = await net_utility.make_request_and_raise_if_error_async(
@@ -320,17 +349,45 @@ async def _construct_http_body_async(
 
     def _get_headers(self) -> dict[str, str]:
         """
-        Headers for accessing inference endpoint deployed in AML.
+        Headers for accessing inference endpoint deployed in AML using a static API key.
+
+        Only valid when this target was configured with a string ``api_key``. When an Entra ID
+        token provider was supplied, callers must use ``_get_headers_async`` instead, because
+        the token must be fetched asynchronously and refreshed on each request.
 
         Returns:
             headers(dict): contains bearer token as AML key and content-type: JSON
+
+        Raises:
+            ValueError: If this target was configured with a token-provider callable.
         """
-        headers: dict[str, str] = {
+        if self._api_key_provider is not None:
+            raise ValueError(
+                "AzureMLChatTarget was configured with an Entra ID token provider. "
+                "Use _get_headers_async() instead of _get_headers()."
+            )
+        return {
             "Content-Type": "application/json",
             "Authorization": ("Bearer " + self._api_key),
         }
 
-        return headers
+    async def _get_headers_async(self) -> dict[str, str]:
+        """
+        Headers for accessing the AML inference endpoint, resolving the bearer token from the
+        configured Entra ID token provider when one is set.
+
+        Returns:
+            headers(dict): contains bearer token (static key or freshly-acquired Entra token) and
+            content-type: JSON.
+        """
+        if self._api_key_provider is None:
+            return self._get_headers()
+
+        token = await self._api_key_provider()
+        return {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer " + token,
+        }
 
     def _validate_request(self, *, normalized_conversation: list[Message]) -> None:
         pass
diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py
index dcbe8b8b4d..0a4f8b34eb 100644
--- a/tests/unit/backend/test_target_service.py
+++ b/tests/unit/backend/test_target_service.py
@@ -306,6 +306,177 @@ async def test_create_target_with_different_underlying_model(self, sqlite_instan
         assert result.underlying_model_name == "gpt-4o"
 
 
+class TestCreateTargetEntraAuth:
+    """Tests for TargetService.create_target_async with auth_mode='entra'."""
+
+    async def test_create_openai_target_with_entra_injects_token_provider(self, sqlite_instance) -> None:
+        """OpenAI Entra path: api_key is replaced with the async token provider callable."""
+
+        async def sentinel_provider() -> str:
+            return "test-token"
+
+        with patch(
+            "pyrit.backend.services.target_service.get_azure_openai_auth",
+            return_value=sentinel_provider,
+        ) as mock_get_auth:
+            service = TargetService()
+
+            request = CreateTargetRequest(
+                type="OpenAIChatTarget",
+                params={
+                    "endpoint": "https://test.openai.azure.com/",
+                    "model_name": "gpt-4o",
+                },
+                auth_mode="entra",
+            )
+
+            result = await service.create_target_async(request=request)
+
+            mock_get_auth.assert_called_once_with("https://test.openai.azure.com/")
+            target_obj = service.get_target_object(target_registry_name=result.target_registry_name)
+            assert target_obj is not None
+            # OpenAI target preserves async callables verbatim through ensure_async_token_provider.
+            assert target_obj._api_key is sentinel_provider  # type: ignore[attr-defined]
+
+    async def test_create_openai_target_with_entra_drops_user_api_key(self, sqlite_instance) -> None:
+        """Any api_key supplied alongside auth_mode='entra' must be discarded."""
+
+        async def sentinel_provider() -> str:
+            return "test-token"
+
+        with patch(
+            "pyrit.backend.services.target_service.get_azure_openai_auth",
+            return_value=sentinel_provider,
+        ):
+            service = TargetService()
+
+            request = CreateTargetRequest(
+                type="OpenAIChatTarget",
+                params={
+                    "endpoint": "https://test.openai.azure.com/",
+                    "model_name": "gpt-4o",
+                    "api_key": "should-be-ignored",
+                },
+                auth_mode="entra",
+            )
+
+            result = await service.create_target_async(request=request)
+
+            target_obj = service.get_target_object(target_registry_name=result.target_registry_name)
+            assert target_obj is not None
+            assert target_obj._api_key is sentinel_provider  # type: ignore[attr-defined]
+            # The literal "should-be-ignored" string must never appear.
+            assert target_obj._api_key != "should-be-ignored"  # type: ignore[attr-defined]
+
+    async def test_create_openai_target_with_entra_does_not_mutate_request_params(self, sqlite_instance) -> None:
+        """The CreateTargetRequest.params object must remain unchanged after creation."""
+
+        async def sentinel_provider() -> str:
+            return "test-token"
+
+        with patch(
+            "pyrit.backend.services.target_service.get_azure_openai_auth",
+            return_value=sentinel_provider,
+        ):
+            service = TargetService()
+
+            original_params = {
+                "endpoint": "https://test.openai.azure.com/",
+                "model_name": "gpt-4o",
+                "api_key": "original-key",
+            }
+            request = CreateTargetRequest(
+                type="OpenAIChatTarget",
+                params=dict(original_params),
+                auth_mode="entra",
+            )
+
+            await service.create_target_async(request=request)
+
+            # The caller's request.params must be unchanged after the call.
+            assert request.params == original_params
+
+    async def test_create_openai_target_with_entra_non_azure_endpoint_raises(self, sqlite_instance) -> None:
+        """Entra ID requires a known Azure OpenAI / AI Foundry hostname suffix."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="OpenAIChatTarget",
+            params={"endpoint": "https://api.openai.com/"},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="Azure endpoint"):
+            await service.create_target_async(request=request)
+
+    async def test_create_openai_target_with_entra_substring_lookalike_endpoint_raises(self, sqlite_instance) -> None:
+        """Substring 'azure' in the hostname must not be enough to pass Entra validation."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="OpenAIChatTarget",
+            # Hostname contains 'azure' but does NOT end with an approved suffix.
+            params={"endpoint": "https://evil-azure.example.com/"},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="Azure endpoint"):
+            await service.create_target_async(request=request)
+
+    async def test_create_openai_target_with_entra_missing_endpoint_raises(self, sqlite_instance) -> None:
+        """Entra ID for OpenAI must reject a missing endpoint with a clear error."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="OpenAIChatTarget",
+            params={},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="endpoint"):
+            await service.create_target_async(request=request)
+
+    async def test_create_azureml_target_with_entra_injects_token_provider(self, sqlite_instance) -> None:
+        """AzureML Entra path: api_key is replaced with the ML scope token provider."""
+
+        async def sentinel_provider() -> str:
+            return "aml-token"
+
+        with patch(
+            "pyrit.backend.services.target_service.get_azure_async_token_provider",
+            return_value=sentinel_provider,
+        ) as mock_get_provider:
+            service = TargetService()
+
+            request = CreateTargetRequest(
+                type="AzureMLChatTarget",
+                params={"endpoint": "https://my-aml.region.inference.ml.azure.com/score"},
+                auth_mode="entra",
+            )
+
+            result = await service.create_target_async(request=request)
+
+            mock_get_provider.assert_called_once_with("https://ml.azure.com/.default")
+            target_obj = service.get_target_object(target_registry_name=result.target_registry_name)
+            assert target_obj is not None
+            # AzureMLChatTarget stores the provider on _api_key_provider; static _api_key is cleared.
+            assert target_obj._api_key_provider is sentinel_provider  # type: ignore[attr-defined]
+            assert target_obj._api_key == ""  # type: ignore[attr-defined]
+
+    async def test_create_target_entra_unsupported_type_raises(self, sqlite_instance) -> None:
+        """Entra ID is only supported for OpenAI-family and AzureMLChatTarget."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="TextTarget",
+            params={},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="does not support Entra"):
+            await service.create_target_async(request=request)
+
+
 class TestTargetServiceSingleton:
     """Tests for get_target_service singleton function."""
 
diff --git a/tests/unit/prompt_target/target/test_azure_ml_chat_target.py b/tests/unit/prompt_target/target/test_azure_ml_chat_target.py
index e9517d3ec8..8ae23b274f 100644
--- a/tests/unit/prompt_target/target/test_azure_ml_chat_target.py
+++ b/tests/unit/prompt_target/target/test_azure_ml_chat_target.py
@@ -236,3 +236,129 @@ def test_valid_temperature_and_top_p(patch_central_database):
     )
     assert target._temperature == 1.5
     assert target._top_p == 0.9
+
+
+def test_initialization_with_async_token_provider_stores_provider(patch_central_database):
+    """A callable api_key is stored as the token provider and _api_key is cleared."""
+
+    async def async_provider() -> str:
+        return "entra-token"
+
+    target = AzureMLChatTarget(
+        endpoint="http://aml-test-endpoint.com",
+        api_key=async_provider,
+    )
+
+    assert target._api_key_provider is async_provider
+    assert target._api_key == ""
+
+
+def test_initialization_with_sync_token_provider_is_wrapped(patch_central_database):
+    """A sync token provider is wrapped so it can be awaited in the async send path."""
+
+    def sync_provider() -> str:
+        return "entra-token"
+
+    target = AzureMLChatTarget(
+        endpoint="http://aml-test-endpoint.com",
+        api_key=sync_provider,
+    )
+
+    # ensure_async_token_provider wraps it, so it's not the same callable but is still callable.
+    assert target._api_key_provider is not None
+    assert target._api_key_provider is not sync_provider
+    assert callable(target._api_key_provider)
+    assert target._api_key == ""
+
+
+def test_initialization_with_callable_ignores_env_var(patch_central_database):
+    """When a callable api_key is provided, AZURE_ML_KEY must not leak through."""
+
+    async def async_provider() -> str:
+        return "entra-token"
+
+    with patch.dict(os.environ, {AzureMLChatTarget.api_key_environment_variable: "env-key"}):
+        target = AzureMLChatTarget(
+            endpoint="http://aml-test-endpoint.com",
+            api_key=async_provider,
+        )
+
+    assert target._api_key_provider is async_provider
+    assert target._api_key == ""
+
+
+def test_get_headers_raises_when_token_provider_configured(patch_central_database):
+    """Sync _get_headers must refuse to operate when an Entra provider is configured."""
+
+    async def async_provider() -> str:
+        return "entra-token"
+
+    target = AzureMLChatTarget(
+        endpoint="http://aml-test-endpoint.com",
+        api_key=async_provider,
+    )
+
+    with pytest.raises(ValueError, match="token provider"):
+        target._get_headers()
+
+
+async def test_get_headers_async_resolves_token_from_provider(patch_central_database):
+    """_get_headers_async awaits the configured provider and builds a Bearer header."""
+
+    async def async_provider() -> str:
+        return "fresh-entra-token"
+
+    target = AzureMLChatTarget(
+        endpoint="http://aml-test-endpoint.com",
+        api_key=async_provider,
+    )
+
+    headers = await target._get_headers_async()
+
+    assert headers == {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer fresh-entra-token",
+    }
+
+
+async def test_get_headers_async_falls_back_to_static_key(patch_central_database):
+    """With a static api_key, _get_headers_async returns the same headers as _get_headers."""
+    target = AzureMLChatTarget(
+        endpoint="http://aml-test-endpoint.com",
+        api_key="static-key",
+    )
+
+    headers = await target._get_headers_async()
+
+    assert headers == {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer static-key",
+    }
+
+
+async def test_complete_chat_async_uses_token_provider_per_request(patch_central_database):
+    """Each send re-acquires the token from the provider (auto-refresh)."""
+    token_provider = AsyncMock(side_effect=["token-1", "token-2"])
+
+    target = AzureMLChatTarget(
+        endpoint="http://aml-test-endpoint.com",
+        api_key=token_provider,
+    )
+
+    messages = [
+        Message(message_pieces=[MessagePiece(role="user", conversation_id="abc", original_value="hi")]),
+    ]
+
+    with patch("pyrit.common.net_utility.make_request_and_raise_if_error_async", new_callable=AsyncMock) as mock:
+        mock_response = MagicMock()
+        mock_response.json.return_value = {"output": "ok"}
+        mock.return_value = mock_response
+
+        await target._complete_chat_async(messages)
+        await target._complete_chat_async(messages)
+
+        assert token_provider.await_count == 2
+        headers_first = mock.call_args_list[0].kwargs["headers"]
+        headers_second = mock.call_args_list[1].kwargs["headers"]
+        assert headers_first["Authorization"] == "Bearer token-1"
+        assert headers_second["Authorization"] == "Bearer token-2"

From 8d0c9198c6f77c086b6ec604355f703ba62c4703 Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Tue, 19 May 2026 15:00:03 -0700
Subject: [PATCH 2/8] formatting

---
 frontend/src/components/Config/CreateTargetDialog.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.tsx b/frontend/src/components/Config/CreateTargetDialog.tsx
index c8b2870fa0..4324f353f9 100644
--- a/frontend/src/components/Config/CreateTargetDialog.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.tsx
@@ -282,7 +282,7 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
                   }}
                 >
                   <Radio value="api_key" label="API Key" />
-                  <Radio value="entra" label="Microsoft Entra ID (DefaultAzureCredential)" />
+                  <Radio value="entra" label="Microsoft Entra Authentication" />
                 </RadioGroup>
               </Field>
 

From 51ec04260b18927e4d30d091cdaafd3ba2e6064c Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Wed, 20 May 2026 09:13:16 -0700
Subject: [PATCH 3/8] fixing UI for error dialog box, minor formatting

---
 .../Config/CreateTargetDialog.styles.ts       |  8 ++++++++
 .../components/Config/CreateTargetDialog.tsx  | 12 +++++------
 pyrit/backend/services/target_service.py      | 20 ++++++-------------
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.styles.ts b/frontend/src/components/Config/CreateTargetDialog.styles.ts
index 34f023786a..67fdb2aa2f 100644
--- a/frontend/src/components/Config/CreateTargetDialog.styles.ts
+++ b/frontend/src/components/Config/CreateTargetDialog.styles.ts
@@ -6,4 +6,12 @@ export const useCreateTargetDialogStyles = makeStyles({
     flexDirection: 'column',
     gap: tokens.spacingVerticalL,
   },
+  warningMessage: {
+    width: '100%',
+  },
+  warningMessageBody: {
+    whiteSpace: 'normal',
+    overflowWrap: 'anywhere',
+    wordBreak: 'break-word',
+  },
 })
diff --git a/frontend/src/components/Config/CreateTargetDialog.tsx b/frontend/src/components/Config/CreateTargetDialog.tsx
index 4324f353f9..c949c87e1a 100644
--- a/frontend/src/components/Config/CreateTargetDialog.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.tsx
@@ -36,8 +36,8 @@ const SUPPORTED_TARGET_TYPES = Object.keys(TARGET_TYPE_CONFIG)
 
 type AuthMode = 'api_key' | 'entra'
 
-// Mirrors the backend's strict hostname-suffix check (target_service.py).
-// Used to warn the user before they submit; the backend remains the source of truth.
+// Mirrors backend's hostname-suffix check (list in target_service.py).
+// The backend still does the check and will reject unsupported endpoints, but this allows us to show a warning in the UI if the user selects Microsoft Entra authentication with a non-Azure OpenAI endpoint.
 const AZURE_OPENAI_HOSTNAME_SUFFIXES = [
   '.openai.azure.com',
   '.ai.azure.com',
@@ -287,10 +287,10 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
               </Field>
 
               {showNonAzureEntraWarning && (
-                <MessageBar intent="warning">
-                  <MessageBarBody>
-                    Microsoft Entra ID authentication is only supported for Azure endpoints
-                    (*.openai.azure.com or *.ai.azure.com). This request will be rejected by the server.
+                <MessageBar intent="warning" className={styles.warningMessage}>
+                  <MessageBarBody className={styles.warningMessageBody}>
+                    Error: Entra auth only works with Azure OpenAI / AI Foundry endpoints (for example,
+                    *.openai.azure.com or *.ai.azure.com).
                   </MessageBarBody>
                 </MessageBar>
               )}
diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py
index c51ec9a611..ab645e06c4 100644
--- a/pyrit/backend/services/target_service.py
+++ b/pyrit/backend/services/target_service.py
@@ -204,11 +204,6 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn
         Instantiates the target with the given type and params,
         then registers it in the registry under its registry name.
 
-        When ``request.auth_mode == "entra"``, an Azure token-provider callable is
-        injected as ``api_key`` before instantiation. Any ``api_key`` value in the
-        original ``params`` is discarded in that case. Entra ID is supported for
-        OpenAI-family targets (against Azure endpoints) and ``AzureMLChatTarget``.
-
         Args:
             request: The create target request with type, params, and auth_mode.
 
@@ -216,13 +211,13 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn
             TargetInstance with the new target's details.
 
         Raises:
-            ValueError: If the target type is not found, or if Entra ID is requested
+            ValueError: If the target type is not found, if Entra ID is requested
                 for an unsupported target type, or if Entra ID is requested for an
                 OpenAI target against a non-Azure endpoint.
         """
         target_class = self._get_target_class(target_type=request.type)
 
-        # Build a fresh params dict so we never mutate the incoming Pydantic model.
+        # Copy params so we can modify values (eg api_key) without changing request.params.
         params: dict[str, Any] = dict(request.params)
 
         if request.auth_mode == "entra":
@@ -237,15 +232,13 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn
     @staticmethod
     def _apply_entra_auth(*, target_class: type, target_type: str, params: dict[str, Any]) -> dict[str, Any]:
         """
-        Replace any ``api_key`` in ``params`` with an Entra ID token provider for
+        Replace ``api_key`` in ``params`` with an Entra ID token provider for
         the given target class.
 
         Args:
-            target_class (type): The target class being instantiated. Used to select the
-                correct token-provider scope.
-            target_type (str): The user-facing target type name (used only in error messages).
-            params (dict[str, Any]): The target constructor parameters from the request.
-                This dict is not mutated.
+            target_class (type): The target class being instantiated
+            target_type (str): The user-facing target type name
+            params (dict[str, Any]): The target constructor parameters from the request
 
         Returns:
             dict[str, Any]: A new params dict with ``api_key`` replaced by an async
@@ -257,7 +250,6 @@ def _apply_entra_auth(*, target_class: type, target_type: str, params: dict[str,
         """
         new_params = dict(params)
         if "api_key" in new_params:
-            # User error: api_key isn't used in Entra mode. Don't log the value.
             logger.debug("Discarding 'api_key' from params because auth_mode='entra'.")
             new_params.pop("api_key", None)
 

From 91e891249844495067070876f44b57409f751d61 Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Wed, 20 May 2026 11:10:21 -0700
Subject: [PATCH 4/8] cleanup/formatting changes

---
 .../components/Config/CreateTargetDialog.test.tsx    | 12 ++++++------
 pyrit/backend/models/targets.py                      |  4 ++--
 pyrit/backend/services/target_service.py             |  9 +++------
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx
index 70966a9a9b..743f229e88 100644
--- a/frontend/src/components/Config/CreateTargetDialog.test.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx
@@ -529,7 +529,7 @@ describe("CreateTargetDialog", () => {
     // Select Entra radio.
     await user.click(
       screen.getByRole("radio", {
-        name: /Microsoft Entra ID/,
+        name: /Microsoft Entra Authentication/,
       })
     );
 
@@ -582,7 +582,7 @@ describe("CreateTargetDialog", () => {
     );
 
     await user.click(
-      screen.getByRole("radio", { name: /Microsoft Entra ID/ })
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
     );
 
     await user.click(screen.getByText("Create Target"));
@@ -611,11 +611,11 @@ describe("CreateTargetDialog", () => {
     fireEvent.change(endpointInput, { target: { value: "https://api.openai.com/" } });
 
     await user.click(
-      screen.getByRole("radio", { name: /Microsoft Entra ID/ })
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
     );
 
     expect(
-      screen.getByText(/Microsoft Entra ID authentication is only supported for Azure endpoints/)
+      screen.getByText(/Entra auth only works with Azure OpenAI/)
     ).toBeInTheDocument();
   });
 
@@ -638,11 +638,11 @@ describe("CreateTargetDialog", () => {
     });
 
     await user.click(
-      screen.getByRole("radio", { name: /Microsoft Entra ID/ })
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
     );
 
     expect(
-      screen.queryByText(/only supported for Azure endpoints/)
+      screen.queryByText(/Entra auth only works with Azure OpenAI/)
     ).not.toBeInTheDocument();
   });
 });
diff --git a/pyrit/backend/models/targets.py b/pyrit/backend/models/targets.py
index 1bd319c117..925cc68ef0 100644
--- a/pyrit/backend/models/targets.py
+++ b/pyrit/backend/models/targets.py
@@ -83,8 +83,8 @@ class CreateTargetRequest(BaseModel):
     auth_mode: Literal["api_key", "entra"] = Field(
         "api_key",
         description=(
-            "Authentication mode. 'api_key' uses the api_key in params (default, backwards-compatible). "
-            "'entra' uses Microsoft Entra ID (DefaultAzureCredential); requires an Azure endpoint and is "
+            "Authentication mode. 'api_key' uses the api_key in params (default)"
+            "'entra' uses Microsoft Entra ID; requires an Azure endpoint and is "
             "supported by OpenAI-family targets and AzureMLChatTarget."
         ),
     )
diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py
index ab645e06c4..e8eb29724b 100644
--- a/pyrit/backend/services/target_service.py
+++ b/pyrit/backend/services/target_service.py
@@ -50,17 +50,14 @@
 def _is_azure_openai_endpoint(endpoint: str) -> bool:
     """
     Return True if ``endpoint`` resolves to a known Azure OpenAI / AI Foundry host.
-
-    Strict hostname-suffix check (not a substring search) so a bearer token is
-    never issued for an attacker-controlled endpoint whose URL merely contains
-    the word "azure".
+    Uses a strict hostname-suffix check (not a substring search).
 
     Args:
         endpoint (str): The endpoint URL to validate.
 
     Returns:
-        bool: True if the endpoint's hostname ends with a recognised Azure OpenAI /
-        AI Foundry suffix; False otherwise (including for malformed URLs).
+        bool: True if the endpoint's hostname ends with a recognised Azure suffix;
+            False otherwise
     """
     hostname = (urlparse(endpoint).hostname or "").lower()
     return any(hostname.endswith(suffix) for suffix in _AZURE_OPENAI_HOSTNAME_SUFFIXES)

From e086c5a91d7b794b539d15e7f0470a28f9425fa0 Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Wed, 20 May 2026 11:12:49 -0700
Subject: [PATCH 5/8] clean up comments

---
 .../src/components/Config/CreateTargetDialog.test.tsx     | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx
index 743f229e88..fc5121298f 100644
--- a/frontend/src/components/Config/CreateTargetDialog.test.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx
@@ -521,19 +521,19 @@ describe("CreateTargetDialog", () => {
       target: { value: "https://my-resource.openai.azure.com/" },
     });
 
-    // API Key field is visible by default.
+    // check that API Key field is visible by default.
     expect(
       screen.getByPlaceholderText("API key (stored in memory only)")
     ).toBeInTheDocument();
 
-    // Select Entra radio.
+    // Select Entra option.
     await user.click(
       screen.getByRole("radio", {
         name: /Microsoft Entra Authentication/,
       })
     );
 
-    // API Key field is hidden in Entra mode.
+    // check that API Key field is hidden when Entra mode is selected.
     expect(
       screen.queryByPlaceholderText("API key (stored in memory only)")
     ).not.toBeInTheDocument();
@@ -575,7 +575,7 @@ describe("CreateTargetDialog", () => {
       target: { value: "https://my-resource.openai.azure.com/" },
     });
 
-    // Type a key, then switch to Entra.
+    // Type a key, then switch to Entra option.
     fireEvent.change(
       screen.getByPlaceholderText("API key (stored in memory only)"),
       { target: { value: "sk-typed-before-switch" } }

From 3e0b62dc19e0b9d8fd714d4f507cad538b33b53c Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Wed, 20 May 2026 13:06:24 -0700
Subject: [PATCH 6/8] grey out create target if warning shown

---
 .../Config/CreateTargetDialog.test.tsx        | 28 +++++++++++++++
 .../components/Config/CreateTargetDialog.tsx  |  6 +++-
 tests/unit/backend/test_target_service.py     | 35 ++++++++-----------
 3 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx
index fc5121298f..cd096a192b 100644
--- a/frontend/src/components/Config/CreateTargetDialog.test.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx
@@ -645,4 +645,32 @@ describe("CreateTargetDialog", () => {
       screen.queryByText(/Entra auth only works with Azure OpenAI/)
     ).not.toBeInTheDocument();
   });
+
+  it("should disable Create Target and skip API call for Entra + non-Azure OpenAI endpoint", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "OpenAIChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-resource.openai.azure.com/"
+    );
+    fireEvent.change(endpointInput, { target: { value: "https://api.test.com/" } });
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
+    );
+
+    const createButton = screen.getByText("Create Target").closest("button");
+    expect(createButton).toBeDisabled();
+
+    await user.click(screen.getByText("Create Target"));
+
+    expect(mockedTargetsApi.createTarget).not.toHaveBeenCalled();
+  });
 });
diff --git a/frontend/src/components/Config/CreateTargetDialog.tsx b/frontend/src/components/Config/CreateTargetDialog.tsx
index c949c87e1a..94f433c17a 100644
--- a/frontend/src/components/Config/CreateTargetDialog.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.tsx
@@ -320,7 +320,11 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
             <Button appearance="secondary" onClick={handleClose} disabled={submitting}>
               Cancel
             </Button>
-            <Button appearance="primary" onClick={handleSubmit} disabled={submitting || !targetType || !endpoint}>
+            <Button
+              appearance="primary"
+              onClick={handleSubmit}
+              disabled={submitting || !targetType || !endpoint || showNonAzureEntraWarning}
+            >
               {submitting ? 'Creating...' : 'Create Target'}
             </Button>
           </DialogActions>
diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py
index 0a4f8b34eb..5aed329775 100644
--- a/tests/unit/backend/test_target_service.py
+++ b/tests/unit/backend/test_target_service.py
@@ -42,6 +42,11 @@ def _mock_target_identifier(*, class_name: str = "MockTarget", **kwargs) -> Comp
     )
 
 
+async def _test_token_provider() -> str:
+    """Shared async token provider used in Entra authentication tests."""
+    return "test-token"
+
+
 class TestListTargets:
     """Tests for TargetService.list_targets method."""
 
@@ -307,17 +312,14 @@ async def test_create_target_with_different_underlying_model(self, sqlite_instan
 
 
 class TestCreateTargetEntraAuth:
-    """Tests for TargetService.create_target_async with auth_mode='entra'."""
+    """Test that creating targets with Entra auth mode properly authenticates and handles edge cases."""
 
     async def test_create_openai_target_with_entra_injects_token_provider(self, sqlite_instance) -> None:
-        """OpenAI Entra path: api_key is replaced with the async token provider callable."""
-
-        async def sentinel_provider() -> str:
-            return "test-token"
+        """Entra auth path: api_key is replaced with the authentication callable"""
 
         with patch(
             "pyrit.backend.services.target_service.get_azure_openai_auth",
-            return_value=sentinel_provider,
+            return_value=_test_token_provider,
         ) as mock_get_auth:
             service = TargetService()
 
@@ -336,17 +338,14 @@ async def sentinel_provider() -> str:
             target_obj = service.get_target_object(target_registry_name=result.target_registry_name)
             assert target_obj is not None
             # OpenAI target preserves async callables verbatim through ensure_async_token_provider.
-            assert target_obj._api_key is sentinel_provider  # type: ignore[attr-defined]
+            assert target_obj._api_key is _test_token_provider  # type: ignore[attr-defined]
 
     async def test_create_openai_target_with_entra_drops_user_api_key(self, sqlite_instance) -> None:
         """Any api_key supplied alongside auth_mode='entra' must be discarded."""
 
-        async def sentinel_provider() -> str:
-            return "test-token"
-
         with patch(
             "pyrit.backend.services.target_service.get_azure_openai_auth",
-            return_value=sentinel_provider,
+            return_value=_test_token_provider,
         ):
             service = TargetService()
 
@@ -364,19 +363,16 @@ async def sentinel_provider() -> str:
 
             target_obj = service.get_target_object(target_registry_name=result.target_registry_name)
             assert target_obj is not None
-            assert target_obj._api_key is sentinel_provider  # type: ignore[attr-defined]
+            assert target_obj._api_key is _test_token_provider  # type: ignore[attr-defined]
             # The literal "should-be-ignored" string must never appear.
             assert target_obj._api_key != "should-be-ignored"  # type: ignore[attr-defined]
 
     async def test_create_openai_target_with_entra_does_not_mutate_request_params(self, sqlite_instance) -> None:
         """The CreateTargetRequest.params object must remain unchanged after creation."""
 
-        async def sentinel_provider() -> str:
-            return "test-token"
-
         with patch(
             "pyrit.backend.services.target_service.get_azure_openai_auth",
-            return_value=sentinel_provider,
+            return_value=_test_token_provider,
         ):
             service = TargetService()
 
@@ -439,12 +435,9 @@ async def test_create_openai_target_with_entra_missing_endpoint_raises(self, sql
     async def test_create_azureml_target_with_entra_injects_token_provider(self, sqlite_instance) -> None:
         """AzureML Entra path: api_key is replaced with the ML scope token provider."""
 
-        async def sentinel_provider() -> str:
-            return "aml-token"
-
         with patch(
             "pyrit.backend.services.target_service.get_azure_async_token_provider",
-            return_value=sentinel_provider,
+            return_value=_test_token_provider,
         ) as mock_get_provider:
             service = TargetService()
 
@@ -460,7 +453,7 @@ async def sentinel_provider() -> str:
             target_obj = service.get_target_object(target_registry_name=result.target_registry_name)
             assert target_obj is not None
             # AzureMLChatTarget stores the provider on _api_key_provider; static _api_key is cleared.
-            assert target_obj._api_key_provider is sentinel_provider  # type: ignore[attr-defined]
+            assert target_obj._api_key_provider is _test_token_provider  # type: ignore[attr-defined]
             assert target_obj._api_key == ""  # type: ignore[attr-defined]
 
     async def test_create_target_entra_unsupported_type_raises(self, sqlite_instance) -> None:

From c7282df387d1b773d39316e1080aabdfaf373d7e Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Wed, 20 May 2026 13:14:45 -0700
Subject: [PATCH 7/8] clean up

---
 pyrit/prompt_target/azure_ml_chat_target.py | 28 ++++++++++-----------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/pyrit/prompt_target/azure_ml_chat_target.py b/pyrit/prompt_target/azure_ml_chat_target.py
index f325b48191..04256b58c5 100644
--- a/pyrit/prompt_target/azure_ml_chat_target.py
+++ b/pyrit/prompt_target/azure_ml_chat_target.py
@@ -196,19 +196,21 @@ def _initialize_vars(
         Set the endpoint and key for accessing the Azure ML model. Use this function to manually
         pass in your own endpoint uri and api key. Defaults to the values in the .env file for the variables
         stored in self.endpoint_uri_environment_variable and self.api_key_environment_variable (which default to
-        "AZURE_ML_MANAGED_ENDPOINT" and "AZURE_ML_KEY" respectively).
+        "AZURE_ML_MANAGED_ENDPOINT" and "AZURE_ML_KEY" respectively). It is recommended to set these variables
+        in the .env file and call _set_env_configuration_vars rather than passing the uri and key directly to
+        this function or the target constructor.
 
-        If ``api_key`` is a callable (synchronous or asynchronous), it is treated as an Entra ID
-        token provider. The callable is stored on ``self._api_key_provider`` and resolved per-request
+        If ``api_key`` is a callable, it is treated as an Entra ID token provider.
+        The callable is stored on ``self._api_key_provider`` and resolved per-request
         inside ``_get_headers_async``. Synchronous providers are wrapped via
-        ``ensure_async_token_provider``. In callable mode the ``AZURE_ML_KEY`` environment variable
-        is intentionally NOT consulted, so the caller-provided provider always wins.
+        ``ensure_async_token_provider``.
 
         Args:
             endpoint (str | None): The endpoint uri for the deployed Azure ML model.
-            api_key (str | Callable[[], str | Awaitable[str]] | None): A static API key string,
-                an async/sync callable returning a bearer token, or None to fall back to the
-                ``AZURE_ML_KEY`` environment variable.
+            api_key (str | Callable[[], str | Awaitable[str]] | None):
+                The API key for accessing the Azure ML endpoint, or a callable
+                which returns a bearer token, or None to fall back to the
+                ``AZURE_ML_KEY`` env variable.
         """
         self._endpoint = default_values.get_required_value(
             env_var_name=self.endpoint_uri_environment_variable, passed_value=endpoint
@@ -216,10 +218,6 @@ def _initialize_vars(
 
         if callable(api_key):
             normalized = ensure_async_token_provider(api_key)
-            # ``ensure_async_token_provider`` returns the input unchanged when it is None
-            # or a string and returns a callable when the input is callable. We already
-            # verified the input is callable above, so narrow the return type for ty.
-            assert callable(normalized), "ensure_async_token_provider must return a callable for callable input"
             provider = cast("Callable[[], Awaitable[str]]", normalized)
             self._api_key_provider: Callable[[], Awaitable[str]] | None = provider
             self._api_key = ""
@@ -347,9 +345,9 @@ def _get_headers(self) -> dict[str, str]:
         """
         Headers for accessing inference endpoint deployed in AML using a static API key.
 
-        Only valid when this target was configured with a string ``api_key``. When an Entra ID
-        token provider was supplied, callers must use ``_get_headers_async`` instead, because
-        the token must be fetched asynchronously and refreshed on each request.
+        Only valid when this target was configured with a string ``api_key``.
+        When an Entra ID token provider supplied, callers must use ``_get_headers_async``
+        instead, because the token must be fetched asynchronously and refreshed on each request.
 
         Returns:
             headers(dict): contains bearer token as AML key and content-type: JSON

From 6c3a1131690f7a9222edaddc46406b629c02d269 Mon Sep 17 00:00:00 2001
From: jbolor21 <86250273+jbolor21@users.noreply.github.com>
Date: Tue, 26 May 2026 15:36:07 -0700
Subject: [PATCH 8/8] address feedback

---
 .../Config/CreateTargetDialog.test.tsx        | 112 +++++++++++++++
 .../components/Config/CreateTargetDialog.tsx  | 102 +++++++++-----
 pyrit/backend/models/targets.py               |   2 +-
 pyrit/backend/services/target_service.py      | 103 +++++++++++++-
 tests/unit/backend/test_target_service.py     | 128 ++++++++++++++++++
 5 files changed, 409 insertions(+), 38 deletions(-)

diff --git a/frontend/src/components/Config/CreateTargetDialog.test.tsx b/frontend/src/components/Config/CreateTargetDialog.test.tsx
index cd096a192b..2918ceb919 100644
--- a/frontend/src/components/Config/CreateTargetDialog.test.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.test.tsx
@@ -72,6 +72,30 @@ describe("CreateTargetDialog", () => {
     expect(createButton.closest("button")).toBeDisabled();
   });
 
+  it("should hide the Authentication field until a target type is selected", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    // No type chosen yet — no auth radios visible, but plain API Key input is.
+    expect(
+      screen.queryByRole("radio", { name: /Microsoft Entra Authentication/ })
+    ).not.toBeInTheDocument();
+    expect(
+      screen.getByPlaceholderText("API key (stored in memory only)")
+    ).toBeInTheDocument();
+
+    // Selecting an Entra-capable type reveals the Authentication field.
+    await selectTargetType(user, "OpenAIChatTarget");
+    expect(
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
+    ).toBeInTheDocument();
+  });
+
   it("should call onClose when Cancel is clicked", async () => {
     const onClose = jest.fn();
     const user = userEvent.setup();
@@ -673,4 +697,92 @@ describe("CreateTargetDialog", () => {
 
     expect(mockedTargetsApi.createTarget).not.toHaveBeenCalled();
   });
+
+  it("should warn the user when Entra is selected for a non-AML endpoint on AzureMLChatTarget", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "AzureMLChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-model.region.inference.ml.azure.com/score"
+    );
+    fireEvent.change(endpointInput, {
+      target: { value: "https://example.com/score" },
+    });
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
+    );
+
+    expect(
+      screen.getByText(
+        /Entra auth for AzureMLChatTarget only works with Azure ML managed online endpoints/
+      )
+    ).toBeInTheDocument();
+  });
+
+  it("should NOT warn when Entra is selected for a recognized AML endpoint", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "AzureMLChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-model.region.inference.ml.azure.com/score"
+    );
+    fireEvent.change(endpointInput, {
+      target: { value: "https://my-llama.eastus.inference.ml.azure.com/score" },
+    });
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
+    );
+
+    expect(
+      screen.queryByText(
+        /Entra auth for AzureMLChatTarget only works with Azure ML managed online endpoints/
+      )
+    ).not.toBeInTheDocument();
+  });
+
+  it("should disable Create Target and skip API call for Entra + non-AML endpoint on AzureMLChatTarget", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <TestWrapper>
+        <CreateTargetDialog {...defaultProps} />
+      </TestWrapper>
+    );
+
+    await selectTargetType(user, "AzureMLChatTarget");
+
+    const endpointInput = screen.getByPlaceholderText(
+      "https://your-model.region.inference.ml.azure.com/score"
+    );
+    fireEvent.change(endpointInput, {
+      target: { value: "https://api.test.com/score" },
+    });
+
+    await user.click(
+      screen.getByRole("radio", { name: /Microsoft Entra Authentication/ })
+    );
+
+    const createButton = screen.getByText("Create Target").closest("button");
+    expect(createButton).toBeDisabled();
+
+    await user.click(screen.getByText("Create Target"));
+
+    expect(mockedTargetsApi.createTarget).not.toHaveBeenCalled();
+  });
 });
diff --git a/frontend/src/components/Config/CreateTargetDialog.tsx b/frontend/src/components/Config/CreateTargetDialog.tsx
index 94f433c17a..c795ced516 100644
--- a/frontend/src/components/Config/CreateTargetDialog.tsx
+++ b/frontend/src/components/Config/CreateTargetDialog.tsx
@@ -22,14 +22,19 @@ import {
 import { targetsApi } from '@/services/api'
 import { useCreateTargetDialogStyles } from './CreateTargetDialog.styles'
 
-const TARGET_TYPE_CONFIG: Record<string, 'openai' | 'azureml'> = {
-  OpenAIChatTarget: 'openai',
-  OpenAICompletionTarget: 'openai',
-  OpenAIImageTarget: 'openai',
-  OpenAIVideoTarget: 'openai',
-  OpenAITTSTarget: 'openai',
-  OpenAIResponseTarget: 'openai',
-  AzureMLChatTarget: 'azureml',
+interface TargetTypeConfig {
+  readonly kind: 'openai' | 'azureml'
+  readonly supportsEntra: boolean
+}
+
+const TARGET_TYPE_CONFIG: Record<string, TargetTypeConfig> = {
+  OpenAIChatTarget: { kind: 'openai', supportsEntra: true },
+  OpenAICompletionTarget: { kind: 'openai', supportsEntra: true },
+  OpenAIImageTarget: { kind: 'openai', supportsEntra: true },
+  OpenAIVideoTarget: { kind: 'openai', supportsEntra: true },
+  OpenAITTSTarget: { kind: 'openai', supportsEntra: true },
+  OpenAIResponseTarget: { kind: 'openai', supportsEntra: true },
+  AzureMLChatTarget: { kind: 'azureml', supportsEntra: true },
 }
 
 const SUPPORTED_TARGET_TYPES = Object.keys(TARGET_TYPE_CONFIG)
@@ -45,6 +50,11 @@ const AZURE_OPENAI_HOSTNAME_SUFFIXES = [
   '.cognitiveservices.azure.com',
 ]
 
+// Mirrors backend's hostname-suffix check for Azure ML managed online endpoints
+// (list in target_service.py). Used to warn the user when Microsoft Entra
+// authentication is selected with a non-AML endpoint for AzureMLChatTarget.
+const AZURE_ML_HOSTNAME_SUFFIXES = ['.inference.ml.azure.com']
+
 function isAzureOpenAiEndpoint(endpoint: string): boolean {
   try {
     const host = new URL(endpoint).hostname.toLowerCase()
@@ -54,6 +64,15 @@ function isAzureOpenAiEndpoint(endpoint: string): boolean {
   }
 }
 
+function isAzureMlEndpoint(endpoint: string): boolean {
+  try {
+    const host = new URL(endpoint).hostname.toLowerCase()
+    return AZURE_ML_HOSTNAME_SUFFIXES.some((s) => host.endsWith(s))
+  } catch {
+    return false
+  }
+}
+
 interface CreateTargetDialogProps {
   open: boolean
   onClose: () => void
@@ -77,11 +96,23 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
   const [error, setError] = useState<string | null>(null)
   const [fieldErrors, setFieldErrors] = useState<{ targetType?: string; endpoint?: string }>({})
 
-  const targetKind = TARGET_TYPE_CONFIG[targetType]
-  const isAzureML = targetKind === 'azureml'
-  const isOpenAi = targetKind === 'openai'
-  const isEntra = authMode === 'entra'
-  const showNonAzureEntraWarning = isEntra && isOpenAi && endpoint !== '' && !isAzureOpenAiEndpoint(endpoint)
+  const targetConfig = TARGET_TYPE_CONFIG[targetType]
+  const isAzureML = targetConfig?.kind === 'azureml'
+  const isOpenAi = targetConfig?.kind === 'openai'
+  const supportsEntra = targetConfig?.supportsEntra ?? false
+  const showAuthField = targetType !== '' && supportsEntra
+  const isEntra = showAuthField && authMode === 'entra'
+  const entraEndpointError: string | null = (() => {
+    if (!isEntra || endpoint === '') return null
+    if (isOpenAi && !isAzureOpenAiEndpoint(endpoint)) {
+      return 'Error: Entra auth only works with Azure OpenAI / AI Foundry endpoints (for example, *.openai.azure.com or *.ai.azure.com).'
+    }
+    if (isAzureML && !isAzureMlEndpoint(endpoint)) {
+      return 'Error: Entra auth for AzureMLChatTarget only works with Azure ML managed online endpoints (for example, *.inference.ml.azure.com).'
+    }
+    return null
+  })()
+  const showEntraEndpointError = entraEndpointError !== null
 
   const resetForm = () => {
     setTargetType('')
@@ -176,7 +207,13 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
               >
                 <Select
                   value={targetType}
-                  onChange={(_, data) => setTargetType(data.value)}
+                  onChange={(_, data) => {
+                    const next = data.value
+                    setTargetType(next)
+                    if (!(TARGET_TYPE_CONFIG[next]?.supportsEntra ?? false)) {
+                      setAuthMode('api_key')
+                    }
+                  }}
                 >
                   <option value="">Select a target type</option>
                   {SUPPORTED_TARGET_TYPES.map((type) => (
@@ -272,25 +309,26 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
                 </>
               )}
 
-              <Field label="Authentication">
-                <RadioGroup
-                  value={authMode}
-                  onChange={(_, data) => {
-                    const next = data.value as AuthMode
-                    setAuthMode(next)
-                    if (next === 'entra') setApiKey('')
-                  }}
-                >
-                  <Radio value="api_key" label="API Key" />
-                  <Radio value="entra" label="Microsoft Entra Authentication" />
-                </RadioGroup>
-              </Field>
+              {showAuthField && (
+                <Field label="Authentication">
+                  <RadioGroup
+                    value={authMode}
+                    onChange={(_, data) => {
+                      const next = data.value as AuthMode
+                      setAuthMode(next)
+                      if (next === 'entra') setApiKey('')
+                    }}
+                  >
+                    <Radio value="api_key" label="API Key" />
+                    <Radio value="entra" label="Microsoft Entra Authentication" />
+                  </RadioGroup>
+                </Field>
+              )}
 
-              {showNonAzureEntraWarning && (
-                <MessageBar intent="warning" className={styles.warningMessage}>
+              {showEntraEndpointError && (
+                <MessageBar intent="error" className={styles.warningMessage}>
                   <MessageBarBody className={styles.warningMessageBody}>
-                    Error: Entra auth only works with Azure OpenAI / AI Foundry endpoints (for example,
-                    *.openai.azure.com or *.ai.azure.com).
+                    {entraEndpointError}
                   </MessageBarBody>
                 </MessageBar>
               )}
@@ -323,7 +361,7 @@ export default function CreateTargetDialog({ open, onClose, onCreated }: CreateT
             <Button
               appearance="primary"
               onClick={handleSubmit}
-              disabled={submitting || !targetType || !endpoint || showNonAzureEntraWarning}
+              disabled={submitting || !targetType || !endpoint || showEntraEndpointError}
             >
               {submitting ? 'Creating...' : 'Create Target'}
             </Button>
diff --git a/pyrit/backend/models/targets.py b/pyrit/backend/models/targets.py
index 925cc68ef0..da512da155 100644
--- a/pyrit/backend/models/targets.py
+++ b/pyrit/backend/models/targets.py
@@ -83,7 +83,7 @@ class CreateTargetRequest(BaseModel):
     auth_mode: Literal["api_key", "entra"] = Field(
         "api_key",
         description=(
-            "Authentication mode. 'api_key' uses the api_key in params (default)"
+            "Authentication mode. 'api_key' uses the api_key in params (default). "
             "'entra' uses Microsoft Entra ID; requires an Azure endpoint and is "
             "supported by OpenAI-family targets and AzureMLChatTarget."
         ),
diff --git a/pyrit/backend/services/target_service.py b/pyrit/backend/services/target_service.py
index e8eb29724b..3995213075 100644
--- a/pyrit/backend/services/target_service.py
+++ b/pyrit/backend/services/target_service.py
@@ -13,6 +13,7 @@
 """
 
 import logging
+import os
 from functools import lru_cache
 from typing import Any
 from urllib.parse import urlparse
@@ -46,6 +47,11 @@
     ".cognitiveservices.azure.com",
 )
 
+# Recognised Azure Machine Learning managed online endpoint hostname suffixes.
+# Used for the same strict endpoint validation when issuing Entra ID tokens
+# against an AML scope.
+_AZURE_ML_HOSTNAME_SUFFIXES = (".inference.ml.azure.com",)
+
 
 def _is_azure_openai_endpoint(endpoint: str) -> bool:
     """
@@ -63,6 +69,46 @@ def _is_azure_openai_endpoint(endpoint: str) -> bool:
     return any(hostname.endswith(suffix) for suffix in _AZURE_OPENAI_HOSTNAME_SUFFIXES)
 
 
+def _is_azure_ml_endpoint(endpoint: str) -> bool:
+    """
+    Return True if ``endpoint`` resolves to a known AML managed host.
+    Uses a strict hostname-suffix check (not a substring search).
+
+    Args:
+        endpoint (str): The endpoint URL to validate.
+
+    Returns:
+        bool: True if the endpoint's hostname ends with a recognised AML suffix;
+            False otherwise.
+    """
+    hostname = (urlparse(endpoint).hostname or "").lower()
+    return any(hostname.endswith(suffix) for suffix in _AZURE_ML_HOSTNAME_SUFFIXES)
+
+
+def _resolve_api_key_env_var(target_class: type) -> str | None:
+    """
+    Return the api_key environment variable name for a target class.
+
+    Args:
+        target_class (type): The target class to inspect.
+
+    Returns:
+        str | None: The env var name, or None if the class does not declare one.
+    """
+    if issubclass(target_class, AzureMLChatTarget):
+        env_var = getattr(target_class, "api_key_environment_variable", None)
+        return env_var if isinstance(env_var, str) and env_var else None
+    if issubclass(target_class, OpenAITarget):
+        try:
+            instance = target_class.__new__(target_class)
+            instance._set_openai_env_configuration_vars()
+        except Exception:
+            return None
+        env_var = getattr(instance, "api_key_environment_variable", None)
+        return env_var if isinstance(env_var, str) and env_var else None
+    return None
+
+
 def _build_target_class_registry() -> dict[str, type]:
     """
     Build a registry mapping target class names to their classes.
@@ -208,9 +254,12 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn
             TargetInstance with the new target's details.
 
         Raises:
-            ValueError: If the target type is not found, if Entra ID is requested
-                for an unsupported target type, or if Entra ID is requested for an
-                OpenAI target against a non-Azure endpoint.
+            ValueError: if any of the following occur:
+                - Target type in request is not found in the class registry;
+                - Entra ID auth is requested but the target type does not support it;
+                - Entra ID auth is requested for an OpenAI target or AzureMLChatTarget
+                    but the endpoint is not valid (not managed by correct hosts);
+                - If auth_mode='api_key' is set for a target but no key is supplied
         """
         target_class = self._get_target_class(target_type=request.type)
 
@@ -219,6 +268,8 @@ async def create_target_async(self, *, request: CreateTargetRequest) -> TargetIn
 
         if request.auth_mode == "entra":
             params = self._apply_entra_auth(target_class=target_class, target_type=request.type, params=params)
+        else:
+            self._validate_api_key_auth(target_class=target_class, params=params)
 
         target_obj = target_class(**params)
         self._registry.register_instance(target_obj)
@@ -242,8 +293,9 @@ def _apply_entra_auth(*, target_class: type, target_type: str, params: dict[str,
             token-provider callable suitable for the target class.
 
         Raises:
-            ValueError: If the target type does not support Entra ID, or if an
-                OpenAI target is given a non-Azure endpoint.
+            ValueError: If the target type does not support Entra ID, if an
+                OpenAI target is given a non-Azure endpoint, or if an
+                AzureMLChatTarget is given a non-AML endpoint.
         """
         new_params = dict(params)
         if "api_key" in new_params:
@@ -263,6 +315,14 @@ def _apply_entra_auth(*, target_class: type, target_type: str, params: dict[str,
             return new_params
 
         if issubclass(target_class, AzureMLChatTarget):
+            endpoint = new_params.get("endpoint")
+            if not isinstance(endpoint, str) or not endpoint:
+                raise ValueError("Entra ID authentication requires an 'endpoint' in params.")
+            if not _is_azure_ml_endpoint(endpoint):
+                raise ValueError(
+                    "Entra ID authentication for AzureMLChatTarget requires an AML endpoint "
+                    f"(*.inference.ml.azure.com). Got: {endpoint}"
+                )
             new_params["api_key"] = get_azure_async_token_provider(_AZURE_ML_SCOPE)
             return new_params
 
@@ -271,6 +331,39 @@ def _apply_entra_auth(*, target_class: type, target_type: str, params: dict[str,
             "Supported types are OpenAI-family targets and AzureMLChatTarget."
         )
 
+    @staticmethod
+    def _validate_api_key_auth(*, target_class: type, params: dict[str, Any]) -> None:
+        """
+        Enforce that ``auth_mode='api_key'`` actually has a usable key.
+
+        Targets that do not authenticate via an api_key (e.g. ``TextTarget``)
+        are skipped since they have no env var and the underlying
+        constructor does not take any ``api_key`` arguments.
+
+        Args:
+            target_class (type): The target class being instantiated.
+            params (dict[str, Any]): The constructor parameters from the request.
+
+        Raises:
+            ValueError: If no API key is provided in params or in the relevant
+                environment variable for a target class that authenticates via
+                an API key.
+        """
+        env_var = _resolve_api_key_env_var(target_class)
+        if env_var is None:
+            return
+
+        if params.get("api_key"):
+            return
+        if os.environ.get(env_var):
+            return
+
+        raise ValueError(
+            f"auth_mode='api_key' requires an API key but none was provided. "
+            f"Pass 'api_key' in params or set the {env_var} environment variable. "
+            "To authenticate with Microsoft Entra ID instead, set auth_mode='entra'."
+        )
+
 
 @lru_cache(maxsize=1)
 def get_target_service() -> TargetService:
diff --git a/tests/unit/backend/test_target_service.py b/tests/unit/backend/test_target_service.py
index 5aed329775..a4b536cfdc 100644
--- a/tests/unit/backend/test_target_service.py
+++ b/tests/unit/backend/test_target_service.py
@@ -456,6 +456,46 @@ async def test_create_azureml_target_with_entra_injects_token_provider(self, sql
             assert target_obj._api_key_provider is _test_token_provider  # type: ignore[attr-defined]
             assert target_obj._api_key == ""  # type: ignore[attr-defined]
 
+    async def test_create_azureml_target_with_entra_non_aml_endpoint_raises(self, sqlite_instance) -> None:
+        """Entra ID for AzureMLChatTarget requires a known AML hostname suffix."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="AzureMLChatTarget",
+            params={"endpoint": "https://example.com/score"},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="AML endpoint"):
+            await service.create_target_async(request=request)
+
+    async def test_create_azureml_target_with_entra_substring_lookalike_endpoint_raises(self, sqlite_instance) -> None:
+        """Substring 'inference.ml.azure.com' in the hostname must not be enough to pass AML validation."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="AzureMLChatTarget",
+            # Hostname contains the AML suffix as a substring but does NOT end with it.
+            params={"endpoint": "https://evil-inference.ml.azure.com.attacker.com/score"},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="AML endpoint"):
+            await service.create_target_async(request=request)
+
+    async def test_create_azureml_target_with_entra_missing_endpoint_raises(self, sqlite_instance) -> None:
+        """Entra ID for AzureMLChatTarget must reject a missing endpoint with a clear error."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="AzureMLChatTarget",
+            params={},
+            auth_mode="entra",
+        )
+
+        with pytest.raises(ValueError, match="endpoint"):
+            await service.create_target_async(request=request)
+
     async def test_create_target_entra_unsupported_type_raises(self, sqlite_instance) -> None:
         """Entra ID is only supported for OpenAI-family and AzureMLChatTarget."""
         service = TargetService()
@@ -470,6 +510,94 @@ async def test_create_target_entra_unsupported_type_raises(self, sqlite_instance
             await service.create_target_async(request=request)
 
 
+class TestCreateTargetApiKeyAuth:
+    """Test that auth_mode='api_key' strictly requires a key in params or environment."""
+
+    async def test_create_openai_target_api_key_mode_without_key_raises(self, sqlite_instance) -> None:
+        """Without an api_key (params or env), OpenAITarget would silently fall back to Entra;
+        the service must reject this so the user's explicit choice is honored."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="OpenAIChatTarget",
+            params={
+                "model_name": "gpt-4o",
+                "endpoint": "https://test.openai.azure.com/",
+            },
+            auth_mode="api_key",
+        )
+
+        with patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("OPENAI_CHAT_KEY", None)
+            with pytest.raises(ValueError, match="auth_mode='api_key' requires an API key"):
+                await service.create_target_async(request=request)
+
+    async def test_create_openai_target_api_key_mode_with_env_var_succeeds(self, sqlite_instance) -> None:
+        """An env-var-supplied key satisfies the api_key requirement."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="OpenAIChatTarget",
+            params={
+                "model_name": "gpt-4o",
+                "endpoint": "https://test.openai.azure.com/",
+            },
+            auth_mode="api_key",
+        )
+
+        with patch.dict(os.environ, {"OPENAI_CHAT_KEY": "env-test-key"}):
+            result = await service.create_target_async(request=request)
+
+        assert result.target_type == "OpenAIChatTarget"
+
+    async def test_create_openai_target_api_key_mode_rejects_empty_key(self, sqlite_instance) -> None:
+        """An empty-string api_key counts as missing and must be rejected."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="OpenAIChatTarget",
+            params={
+                "model_name": "gpt-4o",
+                "endpoint": "https://test.openai.azure.com/",
+                "api_key": "",
+            },
+            auth_mode="api_key",
+        )
+
+        with patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("OPENAI_CHAT_KEY", None)
+            with pytest.raises(ValueError, match="auth_mode='api_key' requires an API key"):
+                await service.create_target_async(request=request)
+
+    async def test_create_azureml_target_api_key_mode_without_key_raises(self, sqlite_instance) -> None:
+        """AzureMLChatTarget in api_key mode also requires an explicit key."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="AzureMLChatTarget",
+            params={"endpoint": "https://my-endpoint.eastus.inference.ml.azure.com/score"},
+            auth_mode="api_key",
+        )
+
+        with patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("AZURE_ML_KEY", None)
+            with pytest.raises(ValueError, match="auth_mode='api_key' requires an API key"):
+                await service.create_target_async(request=request)
+
+    async def test_create_text_target_api_key_mode_skips_validation(self, sqlite_instance) -> None:
+        """Targets without an api_key_environment_variable (e.g. TextTarget) are unaffected."""
+        service = TargetService()
+
+        request = CreateTargetRequest(
+            type="TextTarget",
+            params={},
+            auth_mode="api_key",
+        )
+
+        result = await service.create_target_async(request=request)
+        assert result.target_type == "TextTarget"
+
+
 class TestTargetServiceSingleton:
     """Tests for get_target_service singleton function."""