Agent Builder - improve rejected response detection based on agent response (#21850)

* fix: feat: add litellm_system_prompt support * feat: support new 'litellm_agent' model provider * feat: ui/ - new agent builder ui * fix(anthropic/chat/transformation.py): normalize max_tokens if decimal * feat(agentbuilderview.tsx): run compliance datasets against litellm agent * feat: new response rejection detector * fix: multiple fixes * feat: add mcp tools support to agent builder create an agent with access to llm's + mcp servers
2026-07-10 17:04:47 +00:00 · 2026-02-21 15:34:42 -08:00
parent 9fc6fd647c
commit 1f7eeb274c
10 changed files with 890 additions and 253 deletions
@@ -2,6 +2,9 @@

 This module allows users to write custom guardrail logic using Python-like code
 that runs in a sandboxed environment with access to LiteLLM-provided primitives.
+
+Pre-built custom code for common guardrails (e.g. response rejection detection)
+is available in response_rejection_code.py.
 """

 from typing import TYPE_CHECKING
@@ -9,6 +12,8 @@ from typing import TYPE_CHECKING
 from litellm.types.guardrails import SupportedGuardrailIntegrations

 from .custom_code_guardrail import CustomCodeGuardrail
+from .response_rejection_code import (DEFAULT_REJECTION_PHRASES,
+                                      RESPONSE_REJECTION_GUARDRAIL_CODE)

 if TYPE_CHECKING:
    from litellm.types.guardrails import Guardrail, LitellmParams
@@ -61,5 +66,7 @@ guardrail_class_registry = {

 __all__ = [
    "CustomCodeGuardrail",
+    "DEFAULT_REJECTION_PHRASES",
+    "RESPONSE_REJECTION_GUARDRAIL_CODE",
    "initialize_guardrail",
 ]
@@ -26,6 +26,12 @@ Example custom code (async with HTTP):
            if response["success"] and response["body"].get("flagged"):
                return block("Content flagged by moderation API")
        return allow()
+
+Example: block when response rejects the user (input_type response only):
+
+    Use RESPONSE_REJECTION_GUARDRAIL_CODE from .response_rejection_code — it
+    checks response texts for phrases like "That's not something I can help with"
+    and returns block() so the guardrail raises a block error.
 """

 import asyncio
@@ -35,18 +41,18 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast
 from fastapi import HTTPException

 from litellm._logging import verbose_proxy_logger
-from litellm.integrations.custom_guardrail import (
-    CustomGuardrail,
-    log_guardrail_information,
-)
+from litellm.integrations.custom_guardrail import (CustomGuardrail,
+                                                   log_guardrail_information)
 from litellm.types.guardrails import GuardrailEventHooks
-from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
+from litellm.types.proxy.guardrails.guardrail_hooks.base import \
+    GuardrailConfigModel
 from litellm.types.utils import GenericGuardrailAPIInputs

 from .primitives import get_custom_code_primitives

 if TYPE_CHECKING:
-    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.litellm_core_utils.litellm_logging import \
+        Logging as LiteLLMLoggingObj


 class CustomCodeGuardrailError(Exception):
@@ -0,0 +1,76 @@
+"""
+Custom code for a response guardrail that blocks when the model response
+indicates it is rejecting the user request (e.g. "That's not something I can help with").
+
+Use this with the Custom Code Guardrail (custom_code) by setting litellm_params.custom_code
+to RESPONSE_REJECTION_GUARDRAIL_CODE. The guardrail runs only on input_type "response"
+and raises a block error if any response text matches known rejection phrases.
+"""
+
+# Default phrases that indicate the model is refusing the user request (lowercase for case-insensitive match).
+# Custom code guardrails can override by defining rejection_phrases in the code.
+DEFAULT_REJECTION_PHRASES = [
+    "that's not something i can help with",
+    "that is not something i can help with",
+    "i can't help with that",
+    "i cannot help with that",
+    "i'm not able to help",
+    "i am not able to help",
+    "i'm unable to help",
+    "i cannot assist",
+    "i can't assist",
+    "i'm not allowed to",
+    "i'm not permitted to",
+    "i won't be able to help",
+    "i'm sorry, i can't",
+    "i'm sorry, i cannot",
+    "as an ai, i can't",
+    "as an ai, i cannot",
+]
+
+# Custom code string for the Custom Code Guardrail. Only runs on input_type "response".
+# Uses primitives: allow(), block(), lower(), contains()
+RESPONSE_REJECTION_GUARDRAIL_CODE = '''
+def apply_guardrail(inputs, request_data, input_type):
+    """Block responses that indicate the model rejected the user request."""
+    if input_type != "response":
+        return allow()
+
+    texts = inputs.get("texts") or []
+    # All lowercase for case-insensitive matching (text is lowercased before check)
+    rejection_phrases = [
+        "that's not something i can help with",
+        "that is not something i can help with",
+        "i can't help with that",
+        "i cannot help with that",
+        "i'm not able to help",
+        "i am not able to help",
+        "i'm unable to help",
+        "i cannot assist",
+        "i can't assist",
+        "i'm not allowed to",
+        "i'm not permitted to",
+        "i won't be able to help",
+        "i'm sorry, i can't",
+        "i'm sorry, i cannot",
+        "as an ai, i can't",
+        "as an ai, i cannot",
+    ]
+
+    for text in texts:
+        if not text:
+            continue
+        text_lower = lower(text)
+        for phrase in rejection_phrases:
+            if contains(text_lower, phrase):
+                return block(
+                    "Response indicates the model rejected the user request.",
+                    detection_info={"matched_phrase": phrase, "input_type": "response"},
+                )
+    return allow()
+'''
+
+__all__ = [
+    "DEFAULT_REJECTION_PHRASES",
+    "RESPONSE_REJECTION_GUARDRAIL_CODE",
+]
@@ -12,10 +12,11 @@ All /policy management endpoints
 import copy
 import json
 import os
-from typing import TYPE_CHECKING, AsyncIterator, List, Literal, Optional, cast
+from typing import (TYPE_CHECKING, Any, AsyncIterator, List, Literal, Optional,
+                    cast)

 from fastapi import APIRouter, Depends, HTTPException, Request
-from fastapi.responses import StreamingResponse
+from fastapi.responses import Response, StreamingResponse
 from pydantic import BaseModel, Field
 from typing_extensions import TypedDict

@@ -24,8 +25,12 @@ from litellm.constants import (COMPETITOR_LLM_TEMPERATURE,
                               DEFAULT_COMPETITOR_DISCOVERY_MODEL,
                               MAX_COMPETITOR_NAMES)
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.llms.openai.chat.guardrail_translation.handler import \
+    OpenAIChatCompletionsHandler
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.guardrails.guardrail_hooks.custom_code import (
+    RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail)
 from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
 from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
 from litellm.proxy.policy_engine.policy_registry import get_policy_registry
@@ -39,7 +44,7 @@ from litellm.types.proxy.policy_engine import (PolicyGuardrailsResponse,
                                               PolicyTestResponse,
                                               PolicyValidateRequest,
                                               PolicyValidationResponse)
-from litellm.types.utils import GenericGuardrailAPIInputs
+from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse

 if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import \
@@ -69,20 +74,32 @@ class GuardrailErrorEntry(TypedDict):
    message: str


-class ApplyPoliciesResult(TypedDict):
-    """Result of apply_policies: inputs plus any guardrail failures."""
+class _ApplyPoliciesResultBase(TypedDict):
+    """Base result of apply_policies: inputs plus any guardrail failures."""

    inputs: GenericGuardrailAPIInputs
    guardrail_errors: List[GuardrailErrorEntry]


-class ApplyPoliciesPerItemResult(TypedDict):
-    """Result for one input when using inputs_list."""
+class ApplyPoliciesResult(_ApplyPoliciesResultBase, total=False):
+    """Result of apply_policies. agent_response set when agent_id provided."""
+
+    agent_response: Any
+
+
+class _ApplyPoliciesPerItemResultBase(TypedDict):
+    """Base result for one input when using inputs_list."""

    inputs: GenericGuardrailAPIInputs
    guardrail_errors: List[GuardrailErrorEntry]


+class ApplyPoliciesPerItemResult(_ApplyPoliciesPerItemResultBase, total=False):
+    """Result for one input when using inputs_list. agent_response set when agent_id provided."""
+
+    agent_response: Any
+
+
 class ApplyPoliciesListResult(TypedDict):
    """Result when using inputs_list: one result per input."""

@@ -185,21 +202,78 @@ async def apply_policies(
    return {"inputs": current_inputs, "guardrail_errors": guardrail_errors}


+def _chat_body_from_inputs(
+    inputs: GenericGuardrailAPIInputs, agent_id: str, request_data: dict
+) -> dict:
+    """Build a chat completion request body from guardrail inputs and agent_id."""
+    messages: List[dict]
+    structured = inputs.get("structured_messages")
+    texts = inputs.get("texts")
+    if structured:
+        messages = list(structured)  # type: ignore[arg-type]
+    elif texts:
+        if len(texts) == 1:
+            messages = [{"role": "user", "content": texts[0]}]
+        else:
+            messages = [{"role": "user", "content": "\n".join(texts)}]
+    else:
+        messages = [{"role": "user", "content": "Hello"}]
+    body: dict = {"model": agent_id, "messages": messages, "stream": False}
+    if request_data:
+        body.setdefault("metadata", {}).update(request_data)
+    return body
+
+
+def _request_with_json_body(body: dict) -> Request:
+    """Create a Starlette Request that will return the given dict as parsed JSON body."""
+    body_bytes = json.dumps(body).encode()
+    received: List[bool] = [False]
+
+    async def receive() -> dict:
+        if received[0]:
+            return {"type": "http.disconnect"}
+        received[0] = True
+        return {"type": "http.request", "body": body_bytes, "more_body": False}
+
+    scope: dict = {
+        "type": "http",
+        "method": "POST",
+        "path": "/v1/chat/completions",
+        "query_string": b"",
+        "headers": [(b"content-type", b"application/json")],
+        "scheme": "http",
+        "server": ("localhost", 8000),
+        "client": ("127.0.0.1", 0),
+        "root_path": "",
+        "app": None,
+        "asgi": {"version": "3.0", "spec_version": "2.0"},
+    }
+    return Request(scope, receive=receive)
+
+
 class TestPoliciesAndGuardrailsRequest(BaseModel):
    """Request body for POST /utils/test_policies_and_guardrails."""

-    policy_names: Optional[List[str]] = Field(default=None, description="Policy names to resolve guardrails from")
-    guardrail_names: Optional[List[str]] = Field(default=None, description="Guardrail names to apply directly")
-    inputs: Optional[dict] = Field(
-        default=None,
-        description="GenericGuardrailAPIInputs, e.g. { \"texts\": [\"...\"] }. Use inputs_list for per-input processing.",
+    policy_names: Optional[List[str]] = Field(
+        default=None, description="Policy names to resolve guardrails from"
    )
-    inputs_list: Optional[List[dict]] = Field(
-        default=None,
+    guardrail_names: Optional[List[str]] = Field(
+        default=None, description="Guardrail names to apply directly"
+    )
+    inputs_list: List[GenericGuardrailAPIInputs] = Field(
+        default=[],
        description="List of GenericGuardrailAPIInputs; each item processed separately (for batch compliance testing).",
    )
-    request_data: dict = Field(default_factory=dict, description="Request context (model, user_id, etc.)")
-    input_type: Literal["request", "response"] = Field(default="request", description="Whether inputs are request or response")
+    request_data: dict = Field(
+        default_factory=dict, description="Request context (model, user_id, etc.)"
+    )
+    input_type: Literal["request", "response"] = Field(
+        default="request", description="Whether inputs are request or response"
+    )
+    agent_id: Optional[str] = Field(
+        default=None,
+        description="When set, call chat completion with this model/agent for each input and include the response in the result.",
+    )


@router.post(
@@ -223,40 +297,86 @@ async def test_policies_and_guardrails(
    """
    from litellm.litellm_core_utils.litellm_logging import \
        Logging as LiteLLMLoggingObj
-    from litellm.proxy.proxy_server import proxy_logging_obj
+    from litellm.proxy.proxy_server import chat_completion, proxy_logging_obj
    from litellm.proxy.utils import handle_exception_on_proxy

+    def _serialize_chat_response(response: Any) -> Any:
+        if hasattr(response, "model_dump"):
+            return response.model_dump(exclude_unset=True)
+        if isinstance(response, dict):
+            return response
+        return response
+
+    async def _get_agent_response(
+        inputs: GenericGuardrailAPIInputs,
+        agent_id: str,
+        user_api_key_dict: UserAPIKeyAuth,
+    ) -> Any:
+        body = _chat_body_from_inputs(inputs, agent_id, data.request_data)
+        req = _request_with_json_body(body)
+        resp = Response()
+        result = await chat_completion(
+            request=req,
+            fastapi_response=resp,
+            model=agent_id,
+            user_api_key_dict=user_api_key_dict,
+        )
+        return _serialize_chat_response(result)
+
    try:
        logging_obj = cast(LiteLLMLoggingObj, proxy_logging_obj)
-        if data.inputs_list is not None:
-            results: List[ApplyPoliciesPerItemResult] = []
-            for inp in data.inputs_list:
-                inputs_typed = cast(GenericGuardrailAPIInputs, inp)
-                item_result = await apply_policies(
-                    policy_names=data.policy_names,
-                    inputs=inputs_typed,
-                    request_data=data.request_data,
-                    input_type=data.input_type,
-                    proxy_logging_obj=logging_obj,
-                    guardrail_names=data.guardrail_names,
-                )
-                results.append(
-                    ApplyPoliciesPerItemResult(
-                        inputs=item_result["inputs"],
-                        guardrail_errors=item_result["guardrail_errors"],
-                    )
-                )
-            return ApplyPoliciesListResult(results=results)
-        if data.inputs is not None:
-            inputs_typed = cast(GenericGuardrailAPIInputs, data.inputs)
-            return await apply_policies(
+
+        results: List[ApplyPoliciesPerItemResult] = []
+        for inp in data.inputs_list:
+            item_result = await apply_policies(
                policy_names=data.policy_names,
-                inputs=inputs_typed,
+                inputs=inp,
                request_data=data.request_data,
                input_type=data.input_type,
                proxy_logging_obj=logging_obj,
                guardrail_names=data.guardrail_names,
            )
+            item: ApplyPoliciesPerItemResult = {
+                "inputs": item_result["inputs"],
+                "guardrail_errors": item_result["guardrail_errors"],
+            }
+            if data.agent_id is not None:
+                item["agent_response"] = await _get_agent_response(
+                    item_result["inputs"],
+                    data.agent_id,
+                    user_api_key_dict,
+                )
+                # run response through response_rejection_guardrail (reuses handler extraction + apply)
+                response_rejection_guardrail = CustomCodeGuardrail(
+                    custom_code=RESPONSE_REJECTION_GUARDRAIL_CODE,
+                    guardrail_name="response_rejection",
+                )
+                try:
+                    model_response = ModelResponse.model_validate(
+                        item["agent_response"]
+                    )
+                    handler = OpenAIChatCompletionsHandler()
+                    await handler.process_output_response(
+                        response=model_response,
+                        guardrail_to_apply=response_rejection_guardrail,
+                        litellm_logging_obj=logging_obj,
+                        user_api_key_dict=user_api_key_dict,
+                    )
+                except Exception as guardrail_err:
+                    item["guardrail_errors"] = list(item["guardrail_errors"])
+                    detail = getattr(guardrail_err, "detail", None)
+                    if isinstance(detail, dict) and "error" in detail:
+                        err_msg = detail["error"]
+                    else:
+                        err_msg = str(detail if detail is not None else guardrail_err)
+                    item["guardrail_errors"].append(
+                        GuardrailErrorEntry(
+                            guardrail_name="response_rejection",
+                            message=err_msg,
+                        )
+                    )
+            results.append(item)
+        return ApplyPoliciesListResult(results=results)
        raise ValueError("Either inputs or inputs_list must be provided")
    except Exception as e:
        raise handle_exception_on_proxy(e)
@@ -581,11 +701,15 @@ def _validate_enrichment_request(data: EnrichTemplateRequest) -> tuple[dict, dic
    templates = _load_policy_templates_from_local_backup()
    template = next((t for t in templates if t.get("id") == data.template_id), None)
    if template is None:
-        raise HTTPException(status_code=404, detail=f"Template '{data.template_id}' not found")
+        raise HTTPException(
+            status_code=404, detail=f"Template '{data.template_id}' not found"
+        )

    llm_enrichment = template.get("llm_enrichment")
    if llm_enrichment is None:
-        raise HTTPException(status_code=400, detail="Template does not support LLM enrichment")
+        raise HTTPException(
+            status_code=400, detail="Template does not support LLM enrichment"
+        )

    # Validate competitors list size if provided
    if data.competitors and len(data.competitors) > MAX_COMPETITOR_NAMES:
@@ -695,7 +819,11 @@ async def _stream_llm_competitor_names(
        while "\n" in buffer:
            line, buffer = buffer.split("\n", 1)
            name = _clean_competitor_line(line)
-            if name and name.lower() not in existing_lower and count < MAX_COMPETITOR_NAMES:
+            if (
+                name
+                and name.lower() not in existing_lower
+                and count < MAX_COMPETITOR_NAMES
+            ):
                existing_lower.add(name.lower())
                count += 1
                yield name, False
@@ -744,9 +872,7 @@ async def _stream_competitor_events(
            "{{" + llm_enrichment["parameter"] + "}}", brand_name
        )
        try:
-            async for name, _ in _stream_llm_competitor_names(
-                prompt, model, []
-            ):
+            async for name, _ in _stream_llm_competitor_names(prompt, model, []):
                if name:
                    competitors.append(name)
                    yield f"data: {json.dumps({'type': 'competitor', 'name': name})}\n\n"
@@ -893,10 +1019,7 @@ def _build_all_names_per_competitor(
    competitors: list[str], variations_map: dict[str, list[str]]
 ) -> dict[str, list[str]]:
    """Build canonical + variation name lists for each competitor."""
-    return {
-        comp: [comp] + variations_map.get(comp, [])
-        for comp in competitors
-    }
+    return {comp: [comp] + variations_map.get(comp, []) for comp in competitors}


 def _build_competitor_guardrail_definitions(
@@ -912,7 +1035,9 @@ def _build_competitor_guardrail_definitions(

    output_blocked = _build_name_blocked_words(competitors, all_names)
    recommendation_blocked = _build_recommendation_blocked_words(competitors, all_names)
-    comparison_blocked = _build_comparison_blocked_words(competitors, all_names, brand_name)
+    comparison_blocked = _build_comparison_blocked_words(
+        competitors, all_names, brand_name
+    )

    blocked_words_map = {
        "competitor-output-blocker": output_blocked,
@@ -943,7 +1068,11 @@ def _build_name_blocked_words(
    result = []
    for comp in competitors:
        for name in all_names[comp]:
-            desc = f"Competitor: {comp}" if name == comp else f"Competitor variation ({comp}): {name}"
+            desc = (
+                f"Competitor: {comp}"
+                if name == comp
+                else f"Competitor variation ({comp}): {name}"
+            )
            result.append({"keyword": name, "action": "BLOCK", "description": desc})
    return result

@@ -956,11 +1085,13 @@ def _build_recommendation_blocked_words(
    for comp in competitors:
        for name in all_names[comp]:
            for prefix in ["try", "use", "switch to", "consider"]:
-                result.append({
-                    "keyword": f"{prefix} {name}",
-                    "action": "BLOCK",
-                    "description": f"Recommendation to competitor ({comp})",
-                })
+                result.append(
+                    {
+                        "keyword": f"{prefix} {name}",
+                        "action": "BLOCK",
+                        "description": f"Recommendation to competitor ({comp})",
+                    }
+                )
    return result


@@ -971,23 +1102,29 @@ def _build_comparison_blocked_words(
    result = []
    for comp in competitors:
        for name in all_names[comp]:
-            result.append({
-                "keyword": f"{name} is better",
-                "action": "BLOCK",
-                "description": f"Unfavorable comparison ({comp})",
-            })
+            result.append(
+                {
+                    "keyword": f"{name} is better",
+                    "action": "BLOCK",
+                    "description": f"Unfavorable comparison ({comp})",
+                }
+            )

    # Brand-level comparisons (only need one entry each, not per-competitor)
-    result.append({
-        "keyword": f"better than {brand_name}",
-        "action": "BLOCK",
-        "description": "Unfavorable comparison",
-    })
-    result.append({
-        "keyword": f"{brand_name} is worse",
-        "action": "BLOCK",
-        "description": "Unfavorable comparison",
-    })
+    result.append(
+        {
+            "keyword": f"better than {brand_name}",
+            "action": "BLOCK",
+            "description": "Unfavorable comparison",
+        }
+    )
+    result.append(
+        {
+            "keyword": f"{brand_name} is worse",
+            "action": "BLOCK",
+            "description": "Unfavorable comparison",
+        }
+    )

    return result

@@ -1121,7 +1258,9 @@ async def _test_guardrail_definitions(
                request_data={},
                input_type="request",
            )
-            output_text = output.get("texts", [text])[0] if output.get("texts") else text
+            output_text = (
+                output.get("texts", [text])[0] if output.get("texts") else text
+            )

            if output_text != text:
                action = "masked"
@@ -1,61 +1,47 @@
 import json
 import time
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union
+from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional,
+                    Union)

 from openai._models import BaseModel as OpenAIObject
-from openai.types.audio.transcription_create_params import (
-    FileTypes as FileTypes,  # type: ignore
-)
+from openai.types.audio.transcription_create_params import \
+    FileTypes as FileTypes  # type: ignore
 from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion
-from openai.types.completion_usage import (
-    CompletionTokensDetails,
-    CompletionUsage,
-    PromptTokensDetails,
-)
+from openai.types.completion_usage import (CompletionTokensDetails,
+                                           CompletionUsage,
+                                           PromptTokensDetails)
 from openai.types.moderation import Categories as Categories
-from openai.types.moderation import (
-    CategoryAppliedInputTypes as CategoryAppliedInputTypes,
-)
+from openai.types.moderation import \
+    CategoryAppliedInputTypes as CategoryAppliedInputTypes
 from openai.types.moderation import CategoryScores as CategoryScores
 from openai.types.moderation_create_response import Moderation as Moderation
-from openai.types.moderation_create_response import (
-    ModerationCreateResponse as ModerationCreateResponse,
-)
+from openai.types.moderation_create_response import \
+    ModerationCreateResponse as ModerationCreateResponse
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
 from typing_extensions import Required, TypedDict

 from litellm._uuid import uuid
-from litellm.types.llms.base import (
-    BaseLiteLLMOpenAIResponseObject,
-    LiteLLMPydanticObjectBase,
-)
+from litellm.types.llms.base import (BaseLiteLLMOpenAIResponseObject,
+                                     LiteLLMPydanticObjectBase)
 from litellm.types.mcp import MCPServerCostInfo

 from ..litellm_core_utils.core_helpers import map_finish_reason
 from .agents import LiteLLMSendMessageResponse
 from .guardrails import GuardrailEventHooks
-from .llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
+from .llms.anthropic_messages.anthropic_response import \
+    AnthropicMessagesResponse
 from .llms.base import HiddenParams
-from .llms.openai import (
-    AllMessageValues,
-    Batch,
-    ChatCompletionAnnotation,
-    ChatCompletionRedactedThinkingBlock,
-    ChatCompletionThinkingBlock,
-    ChatCompletionToolCallChunk,
-    ChatCompletionToolParam,
-    ChatCompletionUsageBlock,
-    FileSearchTool,
-    FineTuningJob,
-    ImageURLListItem,
-    OpenAIChatCompletionChunk,
-    OpenAIChatCompletionFinishReason,
-    OpenAIFileObject,
-    OpenAIRealtimeStreamList,
-    ResponsesAPIResponse,
-    WebSearchOptions,
-)
+from .llms.openai import (AllMessageValues, Batch, ChatCompletionAnnotation,
+                          ChatCompletionRedactedThinkingBlock,
+                          ChatCompletionThinkingBlock,
+                          ChatCompletionToolCallChunk, ChatCompletionToolParam,
+                          ChatCompletionUsageBlock, FileSearchTool,
+                          FineTuningJob, ImageURLListItem,
+                          OpenAIChatCompletionChunk,
+                          OpenAIChatCompletionFinishReason, OpenAIFileObject,
+                          OpenAIRealtimeStreamList, ResponsesAPIResponse,
+                          WebSearchOptions)
 from .rerank import RerankResponse as RerankResponse

 if TYPE_CHECKING:
@@ -3173,6 +3159,7 @@ class LlmProviders(str, Enum):
    POE = "poe"
    CHUTES = "chutes"
    XIAOMI_MIMO = "xiaomi_mimo"
+    LITELLM_AGENT = "litellm_agent"


 # Create a set of all provider values for quick lookup
@@ -3203,6 +3190,7 @@ class SearchProviders(str, Enum):
    LINKUP = "linkup"
    DUCKDUCKGO = "duckduckgo"

+
 # Create a set of all search provider values for quick lookup
 SearchProvidersSet = {provider.value for provider in SearchProviders}

@@ -0,0 +1,78 @@
+"""Tests for the response-rejection custom guardrail code (input_type response, block on refusal)."""
+
+import pytest
+from fastapi import HTTPException
+
+from litellm.proxy.guardrails.guardrail_hooks.custom_code import (
+    RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail)
+
+
+@pytest.fixture
+def response_rejection_guardrail():
+    """Guardrail instance using the response-rejection custom code."""
+    return CustomCodeGuardrail(
+        guardrail_name="response_rejection",
+        custom_code=RESPONSE_REJECTION_GUARDRAIL_CODE,
+    )
+
+
+@pytest.mark.asyncio
+async def test_response_rejection_allows_request_input_type(response_rejection_guardrail):
+    """Should allow when input_type is 'request' (no response check)."""
+    result = await response_rejection_guardrail.apply_guardrail(
+        inputs={"texts": ["some user message"]},
+        request_data={},
+        input_type="request",
+    )
+    assert result == {"texts": ["some user message"]}
+
+
+@pytest.mark.asyncio
+async def test_response_rejection_allows_helpful_response(response_rejection_guardrail):
+    """Should allow when response text does not contain rejection phrases."""
+    result = await response_rejection_guardrail.apply_guardrail(
+        inputs={"texts": ["Here is how you can do that: step 1, step 2."]},
+        request_data={},
+        input_type="response",
+    )
+    assert result["texts"] == ["Here is how you can do that: step 1, step 2."]
+
+
+@pytest.mark.asyncio
+async def test_response_rejection_blocks_refusal_phrase(response_rejection_guardrail):
+    """Should block when response contains a known rejection phrase."""
+    with pytest.raises(HTTPException) as exc_info:
+        await response_rejection_guardrail.apply_guardrail(
+            inputs={"texts": ["That's not something I can help with."]},
+            request_data={},
+            input_type="response",
+        )
+    assert exc_info.value.status_code == 400
+    detail = exc_info.value.detail
+    assert isinstance(detail, dict)
+    assert "error" in detail
+    assert "rejected" in detail["error"].lower() or "reject" in detail["error"].lower()
+    assert detail.get("guardrail") == "response_rejection"
+    assert detail.get("detection_info", {}).get("matched_phrase") is not None
+
+
+@pytest.mark.asyncio
+async def test_response_rejection_blocks_case_insensitive(response_rejection_guardrail):
+    """Should block on refusal phrase regardless of case."""
+    with pytest.raises(HTTPException):
+        await response_rejection_guardrail.apply_guardrail(
+            inputs={"texts": ["I'M SORRY, I CAN'T do that."]},
+            request_data={},
+            input_type="response",
+        )
+
+
+@pytest.mark.asyncio
+async def test_response_rejection_empty_texts_allowed(response_rejection_guardrail):
+    """Should allow when texts is empty or missing."""
+    result = await response_rejection_guardrail.apply_guardrail(
+        inputs={},
+        request_data={},
+        input_type="response",
+    )
+    assert result == {}
@@ -5453,6 +5453,8 @@ export interface TestPoliciesAndGuardrailsRequest {
  inputs_list?: GuardrailInputs[] | null;
  request_data?: Record<string, unknown>;
  input_type?: "request" | "response";
+  /** When set, backend runs chat completion with this model/agent per input and includes agent_response in each result. */
+  agent_id?: string | null;
 }

 export interface GuardrailErrorEntry {
@@ -5460,16 +5462,24 @@ export interface GuardrailErrorEntry {
  message: string;
 }

+export interface TestPoliciesAndGuardrailsResultItem {
+  inputs: Record<string, unknown>;
+  guardrail_errors: GuardrailErrorEntry[];
+  /** Present when request included agent_id; serialized chat completion response. */
+  agent_response?: Record<string, unknown>;
+}
+
 export interface TestPoliciesAndGuardrailsResponse {
  inputs?: Record<string, unknown>;
  guardrail_errors?: GuardrailErrorEntry[];
  /** Present when inputs_list was used; one result per input. */
-  results?: Array<{ inputs: Record<string, unknown>; guardrail_errors: GuardrailErrorEntry[] }>;
+  results?: TestPoliciesAndGuardrailsResultItem[];
 }

 export const testPoliciesAndGuardrails = async (
  accessToken: string,
-  body: TestPoliciesAndGuardrailsRequest
+  body: TestPoliciesAndGuardrailsRequest,
+  signal?: AbortSignal
 ): Promise<TestPoliciesAndGuardrailsResponse> => {
  try {
    const url = proxyBaseUrl
@@ -5477,6 +5487,7 @@ export const testPoliciesAndGuardrails = async (
      : `/utils/test_policies_and_guardrails`;
    const response = await fetch(url, {
      method: "POST",
+      signal,
      headers: {
        [globalLitellmHeaderName]: `Bearer ${accessToken}`,
        "Content-Type": "application/json",
@@ -5488,6 +5499,7 @@ export const testPoliciesAndGuardrails = async (
        inputs_list: body.inputs_list ?? null,
        request_data: body.request_data ?? {},
        input_type: body.input_type ?? "request",
+        agent_id: body.agent_id ?? null,
      }),
    });

@@ -1,11 +1,14 @@
 "use client";

-import { CommentOutlined, ExperimentOutlined, PlusOutlined, RobotOutlined, SaveOutlined } from "@ant-design/icons";
-import { Button, Input, Select, Spin, Tabs } from "antd";
+import { CommentOutlined, DeleteOutlined, ExperimentOutlined, LinkOutlined, PlusOutlined, RobotOutlined, SaveOutlined } from "@ant-design/icons";
+import { Button, Input, Modal, Select, Spin, Tabs } from "antd";
 import React, { useCallback, useEffect, useState } from "react";
+import CodeBlock from "@/app/(dashboard)/api-reference/components/CodeBlock";
 import NotificationsManager from "../../molecules/notifications_manager";
-import { modelCreateCall } from "../../networking";
-import { AgentModel, fetchAvailableAgentModels } from "../llm_calls/fetch_agents";
+import { keyCreateCall, modelCreateCall, modelDeleteCall, modelPatchUpdateCall, proxyBaseUrl } from "../../networking";
+import { fetchMCPServers } from "../../networking";
+import { MCPServer } from "../../mcp_tools/types";
+import { AgentModel, fetchAvailableAgentModels, MCPToolEntry } from "../llm_calls/fetch_agents";
 import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models";
 import ComplianceUI from "../complianceUI/ComplianceUI";
 import ChatUI from "./ChatUI";
@@ -28,6 +31,133 @@ export interface AgentBuilderViewProps {

 const NEW_AGENT_ID = "__new__";

+function getConnectTabBaseUrl(
+  proxySettings: AgentBuilderViewProps["proxySettings"],
+  customProxyBaseUrl?: string,
+): string {
+  const customDocBaseUrl = proxySettings?.LITELLM_UI_API_DOC_BASE_URL;
+  if (customDocBaseUrl && customDocBaseUrl.trim()) return customDocBaseUrl;
+  if (proxySettings?.PROXY_BASE_URL) return proxySettings.PROXY_BASE_URL;
+  if (customProxyBaseUrl?.trim()) return customProxyBaseUrl;
+  return "<your_proxy_base_url>";
+}
+
+interface ConnectTabContentProps {
+  agentName: string;
+  proxySettings: AgentBuilderViewProps["proxySettings"];
+  customProxyBaseUrl?: string;
+  accessToken: string | null;
+  userID: string | null;
+  disabledPersonalKeyCreation: boolean;
+  creatingKey: boolean;
+  createdKeyValue: string | null;
+  onCreateKey: () => void;
+}
+
+function ConnectTabContent({
+  agentName,
+  proxySettings,
+  customProxyBaseUrl,
+  disabledPersonalKeyCreation,
+  creatingKey,
+  createdKeyValue,
+  onCreateKey,
+}: ConnectTabContentProps) {
+  const baseUrl = proxyBaseUrl ?? getConnectTabBaseUrl(proxySettings, customProxyBaseUrl);
+  const apiKeyForCurl =
+    createdKeyValue ?
+      createdKeyValue.startsWith("Bearer ") ? createdKeyValue : `Bearer ${createdKeyValue}`
+    : "Bearer sk-1234";
+  const curlExample = `curl -L -X POST '${baseUrl}/v1/chat/completions' \\
+-H 'x-litellm-api-key: ${apiKeyForCurl}' \\
+-d '{
+  "model": "${agentName}",
+  "stream": true,
+  "stream_options": {
+    "include_usage": true
+  },
+  "messages": [
+    {
+      "role": "user",
+      "content": "hey"
+    }
+  ]
+}'`;
+  return (
+    <div className="mx-auto max-w-3xl space-y-6">
+      <div>
+        <h3 className="text-sm font-semibold text-gray-900 mb-1">Proxy base URL</h3>
+        <p className="text-sm text-gray-600 font-mono bg-gray-50 px-2 py-1.5 rounded border border-gray-200 break-all">
+          {baseUrl}
+        </p>
+      </div>
+      <div>
+        <h3 className="text-sm font-semibold text-gray-900 mb-2">Call your agent (cURL)</h3>
+        <CodeBlock code={curlExample} language="bash" />
+      </div>
+      <div className="rounded-lg border border-gray-200 bg-gray-50 p-4">
+        <h3 className="text-sm font-semibold text-gray-900 mb-2">Create a key for this agent</h3>
+        <p className="text-sm text-gray-600 mb-3">
+          Create a virtual key that can only call this agent. The key will be scoped to you (user_id) and restricted to
+          the model <span className="font-mono text-gray-800">{agentName}</span>.
+        </p>
+        <Button
+          type="primary"
+          onClick={onCreateKey}
+          loading={creatingKey}
+          disabled={disabledPersonalKeyCreation}
+        >
+          Create key for this agent
+        </Button>
+        {disabledPersonalKeyCreation && (
+          <p className="text-xs text-amber-600 mt-2">Key creation is disabled for your account.</p>
+        )}
+        {createdKeyValue && (
+          <p className="text-xs text-green-700 mt-2">
+            Key created. It is shown in the cURL example above — copy the snippet to use it.
+          </p>
+        )}
+      </div>
+    </div>
+  );
+}
+
+function getAgentModelId(agent: AgentModel): string | null {
+  const info = agent.model_info as { id?: string } | null | undefined;
+  return info?.id ?? null;
+}
+
+function parseUnderlyingModel(litellmModel: string | undefined): string | undefined {
+  if (!litellmModel || !litellmModel.startsWith("litellm_agent/")) return undefined;
+  return litellmModel.slice("litellm_agent/".length) || undefined;
+}
+
+const MCP_TOOLS_PREFIX = "litellm_proxy/mcp/";
+
+function buildToolsFromServerIds(serverIds: string[], servers: MCPServer[]): MCPToolEntry[] {
+  return serverIds.map((serverId) => {
+    const server = servers.find((s) => s.server_id === serverId);
+    const serverName = server?.alias || server?.server_name || serverId;
+    return {
+      type: "mcp",
+      server_label: "litellm",
+      server_url: `${MCP_TOOLS_PREFIX}${serverName}`,
+      require_approval: "never",
+    };
+  });
+}
+
+function getServerIdsFromTools(tools: MCPToolEntry[], servers: MCPServer[]): string[] {
+  return tools
+    .filter((t) => t.type === "mcp" && t.server_url?.startsWith(MCP_TOOLS_PREFIX))
+    .map((t) => {
+      const suffix = t.server_url.slice(MCP_TOOLS_PREFIX.length);
+      const server = servers.find((s) => (s.alias || s.server_name || s.server_id) === suffix);
+      return server?.server_id;
+    })
+    .filter((id): id is string => id != null);
+}
+
 export default function AgentBuilderView({
  accessToken,
  token,
@@ -42,7 +172,9 @@ export default function AgentBuilderView({
  const [modelGroups, setModelGroups] = useState<ModelGroup[]>([]);
  const [loadingAgents, setLoadingAgents] = useState(true);
  const [selectedId, setSelectedId] = useState<string | null>(null);
-  const [activeTab, setActiveTab] = useState<"configure" | "chat" | "test">("configure");
+  const [activeTab, setActiveTab] = useState<"configure" | "chat" | "test" | "connect">("configure");
+  const [creatingKey, setCreatingKey] = useState(false);
+  const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);

  // Draft for new agent
  const [draftName, setDraftName] = useState("");
@@ -50,12 +182,18 @@ export default function AgentBuilderView({
  const [draftUnderlyingModel, setDraftUnderlyingModel] = useState<string | undefined>(undefined);
  const [draftTemperature, setDraftTemperature] = useState(0.7);
  const [draftMaxTokens, setDraftMaxTokens] = useState(4096);
+  const [draftTools, setDraftTools] = useState<MCPToolEntry[]>([]);
+
+  const [mcpServers, setMCPServers] = useState<MCPServer[]>([]);
+  const [loadingMCPServers, setLoadingMCPServers] = useState(false);

  const [saving, setSaving] = useState(false);
+  const [deleting, setDeleting] = useState(false);

  const effectiveApiKey = apiKey || accessToken || "";
  const selectedAgent = selectedId === NEW_AGENT_ID ? null : agentModels.find((a) => a.model_name === selectedId) ?? null;
  const isNewAgent = selectedId === NEW_AGENT_ID;
+  const selectedAgentModelId = selectedAgent ? getAgentModelId(selectedAgent) : null;

  const loadAgents = useCallback(async () => {
    if (!accessToken || !userID || !userRole) return;
@@ -95,6 +233,52 @@ export default function AgentBuilderView({
    loadModels();
  }, [loadModels]);

+  const loadMCPServers = useCallback(async () => {
+    if (!effectiveApiKey) return;
+    setLoadingMCPServers(true);
+    try {
+      const servers = await fetchMCPServers(effectiveApiKey);
+      setMCPServers(Array.isArray(servers) ? servers : (servers as { data?: MCPServer[] })?.data ?? []);
+    } catch (e) {
+      console.error("Error fetching MCP servers:", e);
+    } finally {
+      setLoadingMCPServers(false);
+    }
+  }, [effectiveApiKey]);
+
+  useEffect(() => {
+    loadMCPServers();
+  }, [loadMCPServers]);
+
+  // Clear created key when switching to another agent
+  useEffect(() => {
+    setCreatedKeyValue(null);
+  }, [selectedId]);
+
+  // Sync draft fields when selecting an existing agent
+  useEffect(() => {
+    if (selectedAgent && !isNewAgent) {
+      setDraftName(selectedAgent.model_name);
+      setDraftSystemPrompt(selectedAgent.litellm_params?.litellm_system_prompt ?? "");
+      const underlying = parseUnderlyingModel(selectedAgent.litellm_params?.model);
+      setDraftUnderlyingModel(underlying ?? modelGroups[0]?.model_group);
+      const p = selectedAgent.litellm_params as { temperature?: number; max_tokens?: number } | undefined;
+      setDraftTemperature(typeof p?.temperature === "number" ? p.temperature : 0.7);
+      setDraftMaxTokens(typeof p?.max_tokens === "number" ? p.max_tokens : 4096);
+      const rawTools = selectedAgent.litellm_params?.tools;
+      const tools: MCPToolEntry[] = Array.isArray(rawTools)
+        ? rawTools.filter((t): t is MCPToolEntry => t && typeof t === "object" && (t as MCPToolEntry).type === "mcp" && typeof (t as MCPToolEntry).server_url === "string")
+        : [];
+      setDraftTools(tools);
+    }
+  }, [selectedId, isNewAgent, selectedAgent?.model_name, selectedAgent?.litellm_params?.tools]);
+
+  const selectedMCPServerIds = getServerIdsFromTools(draftTools, mcpServers);
+
+  const handleMCPServerChange = (serverIds: string[]) => {
+    setDraftTools(buildToolsFromServerIds(serverIds, mcpServers));
+  };
+
  const handleAddAgent = () => {
    setSelectedId(NEW_AGENT_ID);
    setDraftName("");
@@ -102,6 +286,7 @@ export default function AgentBuilderView({
    setDraftUnderlyingModel(modelGroups[0]?.model_group);
    setDraftTemperature(0.7);
    setDraftMaxTokens(4096);
+    setDraftTools([]);
    setActiveTab("configure");
  };

@@ -119,6 +304,7 @@ export default function AgentBuilderView({
          litellm_system_prompt: draftSystemPrompt.trim() || undefined,
          temperature: draftTemperature,
          max_tokens: draftMaxTokens,
+          tools: draftTools,
        },
        model_info: {},
      });
@@ -133,6 +319,86 @@ export default function AgentBuilderView({
    }
  };

+  const handleUpdateAgent = async () => {
+    if (!accessToken || !selectedAgent || !selectedAgentModelId || !draftName?.trim() || !draftUnderlyingModel) {
+      NotificationsManager.fromBackend("Name and underlying model are required");
+      return;
+    }
+    setSaving(true);
+    try {
+      await modelPatchUpdateCall(
+        accessToken,
+        {
+          model_name: draftName.trim(),
+          litellm_params: {
+            model: `litellm_agent/${draftUnderlyingModel}`,
+            litellm_system_prompt: draftSystemPrompt.trim() || undefined,
+            temperature: draftTemperature,
+            max_tokens: draftMaxTokens,
+            tools: draftTools,
+          },
+          model_info: selectedAgent.model_info ?? {},
+        },
+        selectedAgentModelId,
+      );
+      NotificationsManager.success("Agent updated successfully");
+      await loadAgents();
+      setSelectedId(draftName.trim());
+    } catch (e) {
+      NotificationsManager.fromBackend("Failed to update agent");
+    } finally {
+      setSaving(false);
+    }
+  };
+
+  const handleCreateKeyForAgent = async () => {
+    if (!accessToken || !userID || !selectedAgent) return;
+    setCreatingKey(true);
+    setCreatedKeyValue(null);
+    try {
+      const response = await keyCreateCall(accessToken, userID, {
+        models: [selectedAgent.model_name],
+        key_alias: `Agent: ${selectedAgent.model_name}`,
+      });
+      const keyValue = response?.key ?? null;
+      if (keyValue) {
+        setCreatedKeyValue(keyValue);
+        NotificationsManager.success("Virtual key created. Use it in the curl example below.");
+      } else {
+        NotificationsManager.fromBackend("Key created but value not returned");
+      }
+    } catch (e) {
+      NotificationsManager.fromBackend("Failed to create key for agent");
+    } finally {
+      setCreatingKey(false);
+    }
+  };
+
+  const handleDeleteAgent = () => {
+    if (!selectedAgent || !selectedAgentModelId || !accessToken) return;
+    Modal.confirm({
+      title: "Delete agent",
+      content: `Are you sure you want to delete "${selectedAgent.model_name}"? This cannot be undone.`,
+      okText: "Delete",
+      okType: "danger",
+      cancelText: "Cancel",
+      onOk: async () => {
+        setDeleting(true);
+        try {
+          await modelDeleteCall(accessToken, selectedAgentModelId);
+          NotificationsManager.success("Agent deleted");
+          await loadAgents();
+          const remaining = agentModels.filter((a) => a.model_name !== selectedAgent.model_name);
+          setSelectedId(remaining.length > 0 ? remaining[0].model_name : null);
+        } catch (e) {
+          NotificationsManager.fromBackend("Failed to delete agent");
+        } finally {
+          setDeleting(false);
+        }
+      },
+    });
+  };
+
  if (!accessToken || !userID || !userRole) {
    return (
      <div className="flex h-full items-center justify-center p-8 text-gray-500">
@@ -163,7 +429,11 @@ export default function AgentBuilderView({
        <div className="flex items-center gap-2 border-t border-amber-200 bg-amber-50 px-4 py-2 text-xs text-amber-800">
          <ExperimentOutlined className="flex-shrink-0 text-amber-600" />
          <span>
-            Agent Builder is experimental and may change or be removed without notice.
+            Agent Builder is experimental and may change or be removed without notice. We’d love your feedback—email us at{" "}
+            <a href="mailto:product@berri.ai" className="font-medium text-amber-900 underline hover:text-amber-700">
+              product@berri.ai
+            </a>
+            .
          </span>
        </div>
      </div>
@@ -220,7 +490,7 @@ export default function AgentBuilderView({
            <>
              <Tabs
                activeKey={activeTab}
-                onChange={(k) => setActiveTab(k as "configure" | "chat" | "test")}
+                onChange={(k) => setActiveTab(k as "configure" | "chat" | "test" | "connect")}
                className="flex-1 overflow-hidden [&_.ant-tabs-content]:h-full [&_.ant-tabs-tabpane]:h-full [&_.ant-tabs-nav]:pl-4"
                items={[
                  {
@@ -232,8 +502,13 @@ export default function AgentBuilderView({
                    ),
                    children: (
                      <div className="h-full overflow-y-auto p-6">
-                        {isNewAgent ? (
+                        {(isNewAgent || selectedAgent) ? (
                          <div className="mx-auto max-w-xl space-y-4">
+                            {!selectedAgentModelId && selectedAgent && (
+                              <div className="rounded border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-800">
+                                This agent cannot be updated or deleted here (missing model id). Manage it from Models &amp; Endpoints.
+                              </div>
+                            )}
                            <div>
                              <label className="mb-1 block text-sm font-medium text-gray-700">Agent name</label>
                              <Input
@@ -283,30 +558,58 @@ export default function AgentBuilderView({
                                />
                              </div>
                            </div>
-                          </div>
-                        ) : selectedAgent ? (
-                          <div className="mx-auto max-w-xl space-y-4">
                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">Name</label>
-                              <div className="rounded border border-gray-200 bg-gray-50 px-3 py-2 text-sm">
-                                {selectedAgent.model_name}
-                              </div>
+                              <label className="mb-1 block text-sm font-medium text-gray-700">MCP servers</label>
+                              <Select
+                                mode="multiple"
+                                placeholder="Select MCP servers to attach (same format as chat completions API)"
+                                value={selectedMCPServerIds}
+                                onChange={handleMCPServerChange}
+                                loading={loadingMCPServers}
+                                className="w-full"
+                                allowClear
+                                showSearch
+                                optionFilterProp="label"
+                                options={mcpServers.map((s) => ({
+                                  value: s.server_id,
+                                  label: s.alias || s.server_name || s.server_id,
+                                }))}
+                              />
+                              {selectedAgent && draftTools.length > 0 && (
+                                <p className="mt-1 text-xs text-gray-500">
+                                  {draftTools.length} MCP server{draftTools.length !== 1 ? "s" : ""} saved. Use the same <code className="rounded bg-gray-100 px-1">tools</code> array in chat completions when calling this agent.
+                                </p>
+                              )}
                            </div>
-                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">System prompt</label>
-                              <div className="rounded border border-gray-200 bg-gray-50 px-3 py-2 text-sm whitespace-pre-wrap">
-                                {selectedAgent.litellm_params?.litellm_system_prompt || "(none)"}
+                            {selectedAgent && (
+                              <div className="flex flex-wrap items-center gap-2 pt-2">
+                                {selectedAgentModelId && (
+                                  <>
+                                    <Button
+                                      type="primary"
+                                      icon={<SaveOutlined />}
+                                      onClick={handleUpdateAgent}
+                                      loading={saving}
+                                      disabled={!draftName?.trim() || !draftUnderlyingModel}
+                                    >
+                                      Update Agent
+                                    </Button>
+                                    <Button
+                                      type="default"
+                                      danger
+                                      icon={<DeleteOutlined />}
+                                      onClick={handleDeleteAgent}
+                                      loading={deleting}
+                                    >
+                                      Delete
+                                    </Button>
+                                  </>
+                                )}
+                                <Button type="primary" icon={<CommentOutlined />} onClick={() => setActiveTab("chat")}>
+                                  Test in Chat
+                                </Button>
                              </div>
-                            </div>
-                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">Underlying model</label>
-                              <div className="rounded border border-gray-200 bg-gray-50 px-3 py-2 text-sm font-mono">
-                                {selectedAgent.litellm_params?.model ?? ""}
-                              </div>
-                            </div>
-                            <Button type="primary" icon={<CommentOutlined />} onClick={() => setActiveTab("chat")}>
-                              Test in Chat
-                            </Button>
+                            )}
                          </div>
                        ) : null}
                      </div>
@@ -368,6 +671,36 @@ export default function AgentBuilderView({
                      </div>
                    ),
                  },
+                  {
+                    key: "connect",
+                    label: (
+                      <span>
+                        <LinkOutlined className="mr-1" /> Connect
+                      </span>
+                    ),
+                    disabled: isNewAgent,
+                    children: (
+                      <div className="h-full overflow-y-auto p-6">
+                        {selectedAgent ? (
+                          <ConnectTabContent
+                            agentName={selectedAgent.model_name}
+                            proxySettings={proxySettings}
+                            customProxyBaseUrl={customProxyBaseUrl}
+                            accessToken={accessToken}
+                            userID={userID}
+                            disabledPersonalKeyCreation={disabledPersonalKeyCreation}
+                            creatingKey={creatingKey}
+                            createdKeyValue={createdKeyValue}
+                            onCreateKey={handleCreateKeyForAgent}
+                          />
+                        ) : (
+                          <div className="flex flex-1 items-center justify-center text-gray-500">
+                            Select an agent to see how to connect.
+                          </div>
+                        )}
+                      </div>
+                    ),
+                  },
                ]}
              />
            </>
@@ -39,6 +39,7 @@ import {
  Send,
  Shield,
  Smile,
+  Square,
  Trash2,
  TrendingDown,
  Upload,
@@ -161,6 +162,7 @@ export default function ComplianceUI({
  const [isRunning, setIsRunning] = useState(false);
  const [resultFilter, setResultFilter] = useState<ResultFilter>("all");
  const [expandedResults, setExpandedResults] = useState<Set<string>>(new Set());
+  const batchAbortControllerRef = useRef<AbortController | null>(null);

  useEffect(() => {
    if (!accessToken) return;
@@ -570,6 +572,9 @@ export default function ComplianceUI({

  const runTests = useCallback(async () => {
    if (selectedPromptIds.size === 0 || !accessToken) return;
+    const controller = new AbortController();
+    batchAbortControllerRef.current = controller;
+    const signal = controller.signal;
    setIsRunning(true);
    setResultFilter("all");
    setRightTab("batch-results");
@@ -590,100 +595,62 @@ export default function ComplianceUI({
    }));
    setTestResults(pendingResults);
    try {
-      if (backendMode === "chat_completions" && fixedModel) {
-        const newResults = [...pendingResults];
-        for (let index = 0; index < selected.length; index++) {
-          const row = pendingResults[index];
-          const prompt = allTexts[index];
-          let responseText = "";
-          try {
-            await makeOpenAIChatCompletionRequest(
-              [{ role: "user", content: prompt }],
-              (chunk: string) => {
-                responseText += chunk;
-              },
-          fixedModel,
-          accessToken,
-          undefined,
-          undefined,
-          undefined,
-          undefined,
-          undefined,
-          undefined,
-          undefined,
-          selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
-              selectedPolicies.length > 0 ? selectedPolicies : undefined,
-              undefined,
-              undefined,
-              undefined,
-              undefined,
-              undefined,
-              undefined,
-              requestProxyBaseUrl,
-              undefined
-            );
-            const actualResult: "blocked" | "allowed" = "allowed";
-            newResults[index] = {
-              ...row,
-              actualResult,
-              returnedText: responseText,
-              isMatch: row.expectedResult === "pass",
-              status: "complete" as const,
-            };
-          } catch (err) {
-            const errorMessage = err instanceof Error ? err.message : String(err);
-            newResults[index] = {
-              ...row,
-              actualResult: "blocked" as const,
-              isMatch: false,
-              triggeredBy: errorMessage,
-              status: "complete" as const,
-            };
-          }
-          setTestResults([...newResults]);
-        }
-      } else {
-        const inputsList = allTexts.map((text) => ({ texts: [text] }));
-        const response = await testPoliciesAndGuardrails(accessToken, {
+      const useAgentId = backendMode === "chat_completions" && fixedModel;
+      const response = await testPoliciesAndGuardrails(
+        accessToken,
+        {
          policy_names:
            selectedPolicies.length > 0 ? selectedPolicies : undefined,
          guardrail_names:
            selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
-          inputs_list: inputsList,
+          inputs_list: allTexts.map((text) => ({ texts: [text] })),
          request_data: {},
          input_type: "request",
-        });
-        const results = response.results ?? [];
-        setTestResults(
-          pendingResults.map((row, index) => {
-            const item = results[index];
-            const guardrail_errors = item?.guardrail_errors ?? [];
-            const actualResult: "blocked" | "allowed" =
-              guardrail_errors.length > 0 ? "blocked" : "allowed";
-            const triggeredBy =
-              guardrail_errors.length > 0
-                ? guardrail_errors
-                    .map((e) => `${e.guardrail_name}: ${e.message}`)
-                    .join("; ")
+          ...(useAgentId ? { agent_id: fixedModel } : {}),
+        },
+        signal
+      );
+      const results = response.results ?? [];
+      setTestResults(
+        pendingResults.map((row, index) => {
+          const item = results[index];
+          const guardrail_errors = item?.guardrail_errors ?? [];
+          const actualResult: "blocked" | "allowed" =
+            guardrail_errors.length > 0 ? "blocked" : "allowed";
+          const triggeredBy =
+            guardrail_errors.length > 0
+              ? guardrail_errors
+                  .map((e) => `${e.guardrail_name}: ${e.message}`)
+                  .join("; ")
+              : undefined;
+          let returnedText: string | undefined;
+          if (item?.agent_response != null) {
+            const choices = (item.agent_response as { choices?: Array<{ message?: { content?: string } }> }).choices;
+            returnedText =
+              Array.isArray(choices) && choices[0]?.message?.content != null
+                ? String(choices[0].message.content)
                : undefined;
-            const returnedText =
-              Array.isArray(item?.inputs?.texts) && item.inputs.texts.length > 0
-                ? item.inputs.texts[0]
-                : undefined;
-            return {
-              ...row,
-              actualResult,
-              isMatch:
-                (row.expectedResult === "fail" && actualResult === "blocked") ||
-                (row.expectedResult === "pass" && actualResult === "allowed"),
-              triggeredBy,
-              returnedText,
-              status: "complete" as const,
-            };
-          })
-        );
-      }
+          }
+          if (returnedText === undefined && Array.isArray(item?.inputs?.texts) && item.inputs.texts.length > 0) {
+            returnedText = item.inputs.texts[0] as string;
+          }
+          return {
+            ...row,
+            actualResult,
+            isMatch:
+              (row.expectedResult === "fail" && actualResult === "blocked") ||
+              (row.expectedResult === "pass" && actualResult === "allowed"),
+            triggeredBy,
+            returnedText,
+            status: "complete" as const,
+          };
+        })
+      );
    } catch (err) {
+      if (err instanceof Error && err.name === "AbortError") {
+        // Stopped by user; leave partial results as-is (already set in loop)
+        return;
+      }
      const errorMessage = err instanceof Error ? err.message : String(err);
      setTestResults(
        pendingResults.map((row) => ({
@@ -694,8 +661,10 @@ export default function ComplianceUI({
          status: "complete" as const,
        }))
      );
+    } finally {
+      setIsRunning(false);
+      batchAbortControllerRef.current = null;
    }
-    setIsRunning(false);
  }, [
    accessToken,
    selectedPromptIds,
@@ -959,23 +928,30 @@ export default function ComplianceUI({
          </div>

          <div className="flex flex-col gap-1.5 pt-6 flex-shrink-0">
-            <button
-              type="button"
-              onClick={runTests}
-              disabled={selectedPromptIds.size === 0 || isRunning || disabledPersonalKeyCreation}
-              className={`flex items-center gap-1.5 px-4 py-2 rounded-lg text-sm font-medium transition-colors whitespace-nowrap ${selectedPromptIds.size === 0 || isRunning || disabledPersonalKeyCreation ? "bg-gray-100 text-gray-400 cursor-not-allowed" : "bg-blue-600 text-white hover:bg-blue-700"}`}
-            >
-              {isRunning ? (
-                <>
-                  <Loader2 className="w-3.5 h-3.5 animate-spin" /> Running...
-                </>
-              ) : (
-                <>
-                  <Play className="w-3.5 h-3.5" /> Simulate (
-                  {selectedPromptIds.size})
-                </>
-              )}
-            </button>
+            {isRunning ? (
+              <button
+                type="button"
+                onClick={() => batchAbortControllerRef.current?.abort()}
+                className="flex items-center gap-1.5 px-4 py-2 rounded-lg text-sm font-medium transition-colors whitespace-nowrap bg-red-600 text-white hover:bg-red-700"
+              >
+                <Square className="w-3.5 h-3.5" /> Stop
+              </button>
+            ) : (
+              <button
+                type="button"
+                onClick={runTests}
+                disabled={selectedPromptIds.size === 0 || disabledPersonalKeyCreation}
+                className={`flex items-center gap-1.5 px-4 py-2 rounded-lg text-sm font-medium transition-colors whitespace-nowrap ${selectedPromptIds.size === 0 || disabledPersonalKeyCreation ? "bg-gray-100 text-gray-400 cursor-not-allowed" : "bg-blue-600 text-white hover:bg-blue-700"}`}
+              >
+                <Play className="w-3.5 h-3.5" /> Simulate (
+                {selectedPromptIds.size})
+              </button>
+            )}
+            {isRunning && (
+              <span className="text-[11px] text-gray-500 flex items-center gap-1">
+                <Loader2 className="w-3 h-3 animate-spin" /> Running...
+              </span>
+            )}
            <button
              type="button"
              onClick={() => {
@@ -1738,6 +1714,16 @@ export default function ComplianceUI({
                                        : "False positive — incorrectly blocked"}
                                  </span>
                                </div>
+                                {result.returnedText != null && result.returnedText !== "" && (
+                                  <div className="mt-1.5">
+                                    <span className="text-gray-400 block mb-0.5">
+                                      LLM response:
+                                    </span>
+                                    <div className="text-gray-700 bg-gray-50 rounded px-2 py-1.5 border border-gray-100 max-h-32 overflow-y-auto whitespace-pre-wrap break-words">
+                                      {result.returnedText}
+                                    </div>
+                                  </div>
+                                )}
                              </div>
                            )}
                          </div>
@@ -13,12 +13,23 @@ export interface Agent {
  };
 }

+/** MCP tool entry in the same format as chat completions API (litellm_params.tools) */
+export interface MCPToolEntry {
+  type: "mcp";
+  server_label?: string;
+  server_url: string;
+  require_approval?: string;
+  allowed_tools?: string[];
+}
+
 /** Agent model from /model/info where litellm_params.model starts with "litellm_agent/" */
 export interface AgentModel {
  model_name: string;
  litellm_params: {
    model: string;
    litellm_system_prompt?: string;
+    /** Saved MCP tools array (same shape as chat completions API tools) */
+    tools?: MCPToolEntry[];
    [key: string]: unknown;
  };
  model_info?: Record<string, unknown> | null;
@@ -93,6 +104,7 @@ export const fetchAvailableAgentModels = async (
          ...m.litellm_params,
          model: m.litellm_params.model,
          litellm_system_prompt: m.litellm_params?.litellm_system_prompt,
+          tools: Array.isArray(m.litellm_params?.tools) ? m.litellm_params.tools : undefined,
        },
        model_info: m.model_info ?? null,
      }));