diff --git a/.cursor/rules/ai-guard.mdc b/.cursor/rules/ai-guard.mdc index f205ab9912d..acf91965cd2 100644 --- a/.cursor/rules/ai-guard.mdc +++ b/.cursor/rules/ai-guard.mdc @@ -77,7 +77,9 @@ ddtrace/aiguard/ ├── _anthropic_streaming.py # Anthropic stream reconstruction ├── _anthropic_errors.py # Anthropic-compatible abort error class ├── _langchain.py # LangChain agent/chatmodel/llm listeners - └── _mcp.py # MCP metadata builder, OpenAI MCP approval decision cache + ├── _mcp.py # MCP metadata builder, model tool-call registry (real call + │ # IDs + conversation), OpenAI MCP approval decision cache + └── _openai_agents.py # openai-agents MCP server tool listeners ``` The old import path `ddtrace.appsec.ai_guard` remains as a deprecated re-export @@ -137,6 +139,7 @@ Events currently consumed: | OpenAI (stream wrappers) | `openai.patch` / `.unpatch` | | Anthropic | `anthropic.messages.create.before` / `.after`, `anthropic.patch` / `.unpatch` | | LangChain | `langchain.{chatmodel,llm}.{generate,agenerate,stream}.before` / `.finally`, `langchain.{chatmodel,llm}.stream.started`, `langchain.patch` / `.unpatch` | +| openai-agents MCP (opt-in, DD_AI_GUARD_COLLECT_MCP_ENABLED) | `openai_agents.mcp.invoke_tool.before`, `openai_agents.mcp.call_tool.before` | | HTTP (client IP) | `set_http_meta_for_asm` | ### 3. Provider converters → canonical `Message` shape diff --git a/.riot/requirements/127d746.txt b/.riot/requirements/127d746.txt new file mode 100644 index 00000000000..b87ff1dd53c --- /dev/null +++ b/.riot/requirements/127d746.txt @@ -0,0 +1,54 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +colorama==0.4.6 +coverage[toml]==7.16.2 +cryptography==50.0.2 +distro==1.9.0 +griffe==1.15.0 +h11==0.16.0 +httpcore==1.0.9 +httpx==0.28.1 +httpx-sse==0.4.3 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==1.30.0 +mock==5.2.0 +openai==1.99.9 +openai-agents==0.0.19 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pydantic-settings==2.15.0 +pygments==2.21.0 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-dotenv==1.2.4 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==2026.9.1 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +tqdm==4.70.1 +types-requests==2.33.0.20260906 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 diff --git a/.riot/requirements/189d47d.txt b/.riot/requirements/189d47d.txt new file mode 100644 index 00000000000..b87ff1dd53c --- /dev/null +++ b/.riot/requirements/189d47d.txt @@ -0,0 +1,54 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +colorama==0.4.6 +coverage[toml]==7.16.2 +cryptography==50.0.2 +distro==1.9.0 +griffe==1.15.0 +h11==0.16.0 +httpcore==1.0.9 +httpx==0.28.1 +httpx-sse==0.4.3 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==1.30.0 +mock==5.2.0 +openai==1.99.9 +openai-agents==0.0.19 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pydantic-settings==2.15.0 +pygments==2.21.0 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-dotenv==1.2.4 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==2026.9.1 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +tqdm==4.70.1 +types-requests==2.33.0.20260906 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 diff --git a/.riot/requirements/198d9ea.txt b/.riot/requirements/198d9ea.txt new file mode 100644 index 00000000000..a39512635c4 --- /dev/null +++ b/.riot/requirements/198d9ea.txt @@ -0,0 +1,52 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +coverage[toml]==7.16.2 +cryptography==50.0.2 +griffelib==2.3.0 +h11==0.16.0 +httpcore2==2.13.1 +httpx2==2.13.1 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==2.3.0 +mcp-types==2.3.0 +mock==5.2.0 +openai==3.24.0 +openai-agents==0.23.1 +opentelemetry-api==1.45.0 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pygments==2.21.0 +pyjwt==2.15.1 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==2026.9.1 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +truststore==0.10.4 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 +websockets==16.1.1 diff --git a/.riot/requirements/1bb3879.txt b/.riot/requirements/1bb3879.txt new file mode 100644 index 00000000000..dbfca18b898 --- /dev/null +++ b/.riot/requirements/1bb3879.txt @@ -0,0 +1,56 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +colorama==0.4.6 +coverage[toml]==7.16.2 +cryptography==50.0.2 +distro==1.9.0 +exceptiongroup==1.3.1 +griffe==1.15.0 +h11==0.16.0 +httpcore==1.0.9 +httpx==0.28.1 +httpx-sse==0.4.3 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==1.30.0 +mock==5.2.0 +openai==1.99.9 +openai-agents==0.0.19 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pydantic-settings==2.15.0 +pygments==2.21.0 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-dotenv==1.2.4 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==0.30.0 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +tomli==2.4.1 +tqdm==4.70.1 +types-requests==2.33.0.20260906 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 diff --git a/.riot/requirements/53d094a.txt b/.riot/requirements/53d094a.txt new file mode 100644 index 00000000000..a39512635c4 --- /dev/null +++ b/.riot/requirements/53d094a.txt @@ -0,0 +1,52 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +coverage[toml]==7.16.2 +cryptography==50.0.2 +griffelib==2.3.0 +h11==0.16.0 +httpcore2==2.13.1 +httpx2==2.13.1 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==2.3.0 +mcp-types==2.3.0 +mock==5.2.0 +openai==3.24.0 +openai-agents==0.23.1 +opentelemetry-api==1.45.0 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pygments==2.21.0 +pyjwt==2.15.1 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==2026.9.1 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +truststore==0.10.4 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 +websockets==16.1.1 diff --git a/.riot/requirements/63221c8.txt b/.riot/requirements/63221c8.txt new file mode 100644 index 00000000000..3da871353c4 --- /dev/null +++ b/.riot/requirements/63221c8.txt @@ -0,0 +1,53 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +coverage[toml]==7.16.2 +cryptography==50.0.2 +griffelib==2.3.0 +h11==0.16.0 +httpcore2==2.13.1 +httpx2==2.13.1 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==2.3.0 +mcp-types==2.3.0 +mock==5.2.0 +openai==3.24.0 +openai-agents==0.23.1 +opentelemetry-api==1.45.0 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pygments==2.21.0 +pyjwt==2.15.1 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==2026.9.1 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +tomli==2.4.1 +truststore==0.10.4 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 +websockets==16.1.1 diff --git a/.riot/requirements/b61ec3c.txt b/.riot/requirements/b61ec3c.txt new file mode 100644 index 00000000000..a7d3c63346b --- /dev/null +++ b/.riot/requirements/b61ec3c.txt @@ -0,0 +1,55 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +colorama==0.4.6 +coverage[toml]==7.16.2 +cryptography==50.0.2 +distro==1.9.0 +griffe==1.15.0 +h11==0.16.0 +httpcore==1.0.9 +httpx==0.28.1 +httpx-sse==0.4.3 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==1.30.0 +mock==5.2.0 +openai==1.99.9 +openai-agents==0.0.19 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pydantic-settings==2.15.0 +pygments==2.21.0 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-dotenv==1.2.4 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==2026.9.1 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +tomli==2.4.1 +tqdm==4.70.1 +types-requests==2.33.0.20260906 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 diff --git a/.riot/requirements/bdbcb83.txt b/.riot/requirements/bdbcb83.txt new file mode 100644 index 00000000000..720e23395d9 --- /dev/null +++ b/.riot/requirements/bdbcb83.txt @@ -0,0 +1,54 @@ +annotated-types==0.8.0 +anyio==4.15.1 +attrs==26.1.0 +certifi==2026.7.22 +cffi==2.1.1 +charset-normalizer==3.5.2 +click==8.5.0 +coverage[toml]==7.16.2 +cryptography==50.0.2 +exceptiongroup==1.3.1 +griffelib==2.3.0 +h11==0.16.0 +httpcore2==2.13.1 +httpx2==2.13.1 +hypothesis==6.45.0 +idna==3.20 +iniconfig==2.3.0 +jiter==0.17.0 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +mcp==2.3.0 +mcp-types==2.3.0 +mock==5.2.0 +openai==3.24.0 +openai-agents==0.23.1 +opentelemetry-api==1.45.0 +opentracing==2.4.0 +packaging==26.3 +pluggy==1.6.0 +pycparser==3.0 +pydantic==2.13.5 +pydantic-core==2.46.5 +pygments==2.21.0 +pyjwt==2.15.1 +pyjwt[crypto]==2.15.1 +pytest==8.4.2 +pytest-asyncio==0.23.7 +pytest-cov==7.1.0 +pytest-mock==3.16.0 +python-multipart==0.0.32 +referencing==0.37.0 +requests==2.34.2 +rpds-py==0.30.0 +sniffio==1.3.1 +sortedcontainers==2.4.0 +sse-starlette==3.5.0 +starlette==1.7.0 +tomli==2.4.1 +truststore==0.10.4 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +urllib3==2.8.0 +uvicorn==0.54.0 +websockets==16.1.1 diff --git a/ddtrace/aiguard/_constants.py b/ddtrace/aiguard/_constants.py index 87b2485da04..25b5602a656 100644 --- a/ddtrace/aiguard/_constants.py +++ b/ddtrace/aiguard/_constants.py @@ -24,8 +24,12 @@ class AI_GUARD(metaclass=Constant_Class): MCP_SERVER_NAME_TAG: str = TAG + ".mcp.name" MCP_SERVER_URL_TAG: str = TAG + ".mcp.url" MCP_TRANSPORT_TAG: str = TAG + ".mcp.transport" - # OpenAI runs hosted MCP tools itself and does not report whether it reached the server over - # streamable HTTP or SSE, so their calls carry the unknown transport. + # Values of the transport field of the optional mcp object. OpenAI runs hosted MCP tools itself + # and does not report whether it reached the server over streamable HTTP or SSE, so their calls + # carry the unknown transport. + MCP_TRANSPORT_STREAMABLE_HTTP: Literal["streamable_http"] = "streamable_http" + MCP_TRANSPORT_SSE: Literal["sse"] = "sse" + MCP_TRANSPORT_STDIO: Literal["stdio"] = "stdio" MCP_TRANSPORT_UNKNOWN: Literal["unknown"] = "unknown" # Core-context key for the (client_ip, peer_ip) pair, applied to the service-entry @@ -89,6 +93,7 @@ class AI_GUARD(metaclass=Constant_Class): INTEGRATION_LANGCHAIN: Literal["langchain"] = "langchain" INTEGRATION_LITELLM: Literal["litellm"] = "litellm" INTEGRATION_STRANDS: Literal["strands"] = "strands" + INTEGRATION_OPENAI_AGENTS: Literal["openai_agents"] = "openai_agents" # Closed tag sets: anything else reaching the metrics is clamped back to these defaults, # so a bad value from a caller cannot invent telemetry series. @@ -100,6 +105,7 @@ class AI_GUARD(metaclass=Constant_Class): INTEGRATION_LANGCHAIN, INTEGRATION_LITELLM, INTEGRATION_STRANDS, + INTEGRATION_OPENAI_AGENTS, ) # environment variables @@ -117,6 +123,7 @@ class AI_GUARD(metaclass=Constant_Class): ENV_OPENAI_ENABLED: Literal["DD_AI_GUARD_OPENAI_ENABLED"] = "DD_AI_GUARD_OPENAI_ENABLED" ENV_ANTHROPIC_ENABLED: Literal["DD_AI_GUARD_ANTHROPIC_ENABLED"] = "DD_AI_GUARD_ANTHROPIC_ENABLED" ENV_LANGCHAIN_ENABLED: Literal["DD_AI_GUARD_LANGCHAIN_ENABLED"] = "DD_AI_GUARD_LANGCHAIN_ENABLED" - # Adds MCP metadata to tool-call evaluations and evaluates MCP approval requests. Off by - # default until the AI Guard service validates the optional mcp object of the evaluate contract. + # Evaluates MCP tool calls with their MCP metadata: OpenAI hosted MCP calls and approvals, and + # openai-agents MCP server tools before they run. Off by default until the AI Guard service + # validates the optional mcp object of the evaluate contract. ENV_COLLECT_MCP_ENABLED: Literal["DD_AI_GUARD_COLLECT_MCP_ENABLED"] = "DD_AI_GUARD_COLLECT_MCP_ENABLED" diff --git a/ddtrace/aiguard/_listener.py b/ddtrace/aiguard/_listener.py index ae3f414446e..611a139c119 100644 --- a/ddtrace/aiguard/_listener.py +++ b/ddtrace/aiguard/_listener.py @@ -28,6 +28,8 @@ from ddtrace.aiguard.integrations._langchain import _langchain_patch from ddtrace.aiguard.integrations._langchain import _langchain_stream_started from ddtrace.aiguard.integrations._langchain import _langchain_unpatch +from ddtrace.aiguard.integrations._openai_agents import _openai_agents_mcp_call_tool_before +from ddtrace.aiguard.integrations._openai_agents import _openai_agents_mcp_invoke_tool_before from ddtrace.aiguard.integrations._openai_chat import _openai_chat_completion_after from ddtrace.aiguard.integrations._openai_chat import _openai_chat_completion_before from ddtrace.aiguard.integrations._openai_responses import _openai_response_create_after @@ -125,6 +127,11 @@ def _openai_listen(client: AIGuardClient) -> None: core.on("openai.responses.create.after", partial(_openai_response_create_after, client)) core.on("openai.patch", partial(_install_openai_wrappers, client)) core.on("openai.unpatch", _uninstall_openai_wrappers) + # openai-agents MCP server tools, opt-in (DD_AI_GUARD_COLLECT_MCP_ENABLED) until the AI Guard + # service validates the optional mcp object of the evaluate contract. + if aiguard_config._ai_guard_collect_mcp_enabled: + core.on("openai_agents.mcp.invoke_tool.before", partial(_openai_agents_mcp_invoke_tool_before, client)) + core.on("openai_agents.mcp.call_tool.before", partial(_openai_agents_mcp_call_tool_before, client)) def _make_openai_stream_wrappers( diff --git a/ddtrace/aiguard/integrations/_mcp.py b/ddtrace/aiguard/integrations/_mcp.py index 8f3c8dd6fe1..f0ece1f716e 100644 --- a/ddtrace/aiguard/integrations/_mcp.py +++ b/ddtrace/aiguard/integrations/_mcp.py @@ -1,15 +1,23 @@ """Shared helpers for evaluating MCP tool calls detected on the model SDK side. -OpenAI hosted MCP approvals are decided when the approval request is returned, and that decision -is reused when the application sends its approval back. +Model function calls carry no MCP identity, so provider listeners record the tool calls a model +returned and the conversation around them; an MCP adapter that later runs one of those calls looks +it up by call ID to evaluate it with the real ID and history. OpenAI hosted MCP approvals are +decided when the approval request is returned, and that decision is reused when the application +sends its approval back. """ from collections import OrderedDict +from contextvars import ContextVar import threading from typing import NamedTuple from typing import Optional +import uuid from ddtrace.aiguard._types import MCP +from ddtrace.aiguard._types import Message +from ddtrace.aiguard._types import ToolCall +from ddtrace.internal.settings.aiguard import aiguard_config from ddtrace.internal.utils.http import canonicalize_url @@ -24,6 +32,53 @@ def mcp_metadata(transport: str, tool_name: str, name: Optional[str] = None, url return mcp +def local_tool_call_id() -> str: + """ID for a tool call no model issued, prefixed so it is never mistaken for a provider ID.""" + return f"dd_mcp_{uuid.uuid4().hex}" + + +class _ModelToolCall(NamedTuple): + messages: list[Message] + message_index: int + + +# Tool calls of the last model response evaluated in this context, keyed by call ID. +_model_tool_calls: ContextVar[Optional[dict[str, _ModelToolCall]]] = ContextVar( + "ai_guard_model_tool_calls", default=None +) + + +def record_model_tool_calls(messages: list[Message], start: int) -> None: + """Remember the tool calls in messages[start:], the model response part of a conversation.""" + if not aiguard_config._ai_guard_collect_mcp_enabled: + return + calls = { + tool_call["id"]: _ModelToolCall(messages, index) + for index in range(start, len(messages)) + for tool_call in messages[index].get("tool_calls") or [] + if tool_call.get("id") + } + # Replaced even when empty: a newer response supersedes the calls of an older one. + _model_tool_calls.set(calls or None) + + +def tool_call_conversation(tool_call: ToolCall) -> list[Message]: + """Return the conversation evaluated before tool_call runs. + + When the call came from a recorded model response, the real history is kept and tool_call + replaces the model's calls in the assistant turn, since sibling calls are evaluated when they + run. Otherwise only the tool call is sent: no history is fabricated. + """ + calls = _model_tool_calls.get() + model_call = calls.get(tool_call["id"]) if calls else None + if model_call is None: + return [Message(role="assistant", tool_calls=[tool_call])] + messages, message_index = model_call + assistant = messages[message_index].copy() + assistant["tool_calls"] = [tool_call] + return messages[:message_index] + [assistant] + + class ApprovalDecision(NamedTuple): blocked: bool action: str diff --git a/ddtrace/aiguard/integrations/_openai_agents.py b/ddtrace/aiguard/integrations/_openai_agents.py new file mode 100644 index 00000000000..b5eeefe1646 --- /dev/null +++ b/ddtrace/aiguard/integrations/_openai_agents.py @@ -0,0 +1,129 @@ +"""AI Guard evaluation of MCP tool calls run by openai-agents MCP servers. + +openai-agents converts the tools of each configured MCP server into function tools and runs the +ones the model selects through MCPUtil.invoke_mcp_tool, which knows the server, the original tool +name, the model-visible name and the model call ID. Applications can also call server.call_tool +directly. Both are evaluated before the tools/call request is sent; an agent-driven call is +evaluated once, at the adapter, and the lower server.call_tool check is skipped for it. +""" + +from collections.abc import Mapping +from contextvars import ContextVar +import json +from typing import Any +from typing import Optional + +from ddtrace.aiguard._api_client import AIGuardAbortError +from ddtrace.aiguard._api_client import AIGuardClient +from ddtrace.aiguard._api_client import Function +from ddtrace.aiguard._api_client import Message +from ddtrace.aiguard._api_client import ToolCall +from ddtrace.aiguard._common import evaluate_auto +from ddtrace.aiguard._constants import AI_GUARD +from ddtrace.aiguard._types import MCP +from ddtrace.aiguard.integrations._mcp import local_tool_call_id +from ddtrace.aiguard.integrations._mcp import mcp_metadata +from ddtrace.aiguard.integrations._mcp import tool_call_conversation +import ddtrace.internal.logger as ddlogger + + +logger = ddlogger.get_logger(__name__) + +# Matched on the class hierarchy so subclasses of the SDK servers keep their transport. +_TRANSPORT_BY_SERVER_CLASS = { + "MCPServerStdio": AI_GUARD.MCP_TRANSPORT_STDIO, + "MCPServerSse": AI_GUARD.MCP_TRANSPORT_SSE, + "MCPServerStreamableHttp": AI_GUARD.MCP_TRANSPORT_STREAMABLE_HTTP, +} +# Names the SDK generates when none is configured. They embed the stdio command or the raw URL, +# so they are never reported. +_GENERATED_NAME_PREFIXES = ("stdio: ", "sse: ", "streamable_http: ") + + +class _AdapterCall: + """An agent-driven call already evaluated at the adapter, consumed by its server.call_tool.""" + + __slots__ = ("server", "tool_name", "consumed") + + def __init__(self, server: Any, tool_name: str) -> None: + self.server = server + self.tool_name = tool_name + self.consumed = False + + +_adapter_call: ContextVar[Optional[_AdapterCall]] = ContextVar("ai_guard_openai_agents_mcp_call", default=None) + + +def _server_mcp(server: Any, tool_name: str) -> MCP: + transport = next( + ( + _TRANSPORT_BY_SERVER_CLASS[cls.__name__] + for cls in type(server).__mro__ + if cls.__name__ in _TRANSPORT_BY_SERVER_CLASS + ), + AI_GUARD.MCP_TRANSPORT_UNKNOWN, + ) + try: + name = server.name + except Exception: + name = None + if not isinstance(name, str) or name.startswith(_GENERATED_NAME_PREFIXES): + name = None + url = None + if transport in (AI_GUARD.MCP_TRANSPORT_SSE, AI_GUARD.MCP_TRANSPORT_STREAMABLE_HTTP): + params = getattr(server, "params", None) + candidate = params.get("url") if isinstance(params, Mapping) else None + url = candidate if isinstance(candidate, str) else None + return mcp_metadata(transport, tool_name, name=name, url=url) + + +def _evaluate(client: AIGuardClient, messages: list[Message]) -> None: + try: + evaluate_auto(client, messages, AI_GUARD.INTEGRATION_OPENAI_AGENTS) + except AIGuardAbortError: + raise + except Exception: + logger.debug("Failed to evaluate openai-agents MCP tool call", exc_info=True) + + +def _openai_agents_mcp_invoke_tool_before( + client: AIGuardClient, server: Any, tool: Any, context: Any, input_json: Any, tool_display_name: Any +) -> None: + """Listener for openai_agents.mcp.invoke_tool.before.""" + tool_name = getattr(tool, "name", None) + if not isinstance(tool_name, str) or not tool_name: + return + call_id = getattr(context, "tool_call_id", None) + model_name = tool_display_name or getattr(context, "tool_name", None) or tool_name + tool_call = ToolCall( + id=call_id if isinstance(call_id, str) and call_id else local_tool_call_id(), + function=Function( + name=str(model_name), arguments=input_json if isinstance(input_json, str) and input_json else "{}" + ), + mcp=_server_mcp(server, tool_name), + ) + _evaluate(client, tool_call_conversation(tool_call)) + # Set only once the call may proceed, so a blocked call cannot excuse a later direct call. + _adapter_call.set(_AdapterCall(server, tool_name)) + + +def _openai_agents_mcp_call_tool_before(client: AIGuardClient, server: Any, tool_name: Any, arguments: Any) -> None: + """Listener for openai_agents.mcp.call_tool.before.""" + if not isinstance(tool_name, str) or not tool_name: + return + adapter_call = _adapter_call.get() + if ( + adapter_call is not None + and not adapter_call.consumed + and adapter_call.server is server + and adapter_call.tool_name == tool_name + ): + adapter_call.consumed = True + return + # Direct call: no model asked for it, so no history and an ID marked local. + tool_call = ToolCall( + id=local_tool_call_id(), + function=Function(name=tool_name, arguments=json.dumps(arguments or {}, default=str)), + mcp=_server_mcp(server, tool_name), + ) + _evaluate(client, [Message(role="assistant", tool_calls=[tool_call])]) diff --git a/ddtrace/aiguard/integrations/_openai_chat.py b/ddtrace/aiguard/integrations/_openai_chat.py index 2206885abbc..597e3daa64c 100644 --- a/ddtrace/aiguard/integrations/_openai_chat.py +++ b/ddtrace/aiguard/integrations/_openai_chat.py @@ -17,6 +17,7 @@ from ddtrace.aiguard._common import evaluate_auto from ddtrace.aiguard._constants import AI_GUARD from ddtrace.aiguard._context import is_aiguard_context_active +from ddtrace.aiguard.integrations._mcp import record_model_tool_calls from ddtrace.aiguard.integrations._openai import _wrap_abort_error import ddtrace.internal.logger as ddlogger @@ -217,4 +218,5 @@ def _openai_chat_completion_after(client: AIGuardClient, kwargs: dict[str, Any], raise _wrap_abort_error(e) except Exception: logger.debug("Failed to evaluate OpenAI chat completion response", exc_info=True) + record_model_tool_calls(all_messages, len(request_messages)) return None diff --git a/ddtrace/aiguard/integrations/_openai_responses.py b/ddtrace/aiguard/integrations/_openai_responses.py index b8390276c44..d0a647616a9 100644 --- a/ddtrace/aiguard/integrations/_openai_responses.py +++ b/ddtrace/aiguard/integrations/_openai_responses.py @@ -36,6 +36,7 @@ from ddtrace.aiguard.integrations._mcp import ApprovalDecision from ddtrace.aiguard.integrations._mcp import approval_decisions from ddtrace.aiguard.integrations._mcp import mcp_metadata +from ddtrace.aiguard.integrations._mcp import record_model_tool_calls from ddtrace.aiguard.integrations._openai import _wrap_abort_error import ddtrace.internal.logger as ddlogger from ddtrace.internal.settings.aiguard import aiguard_config @@ -626,4 +627,5 @@ def _openai_response_create_after(client: AIGuardClient, kwargs: dict[str, Any], raise _wrap_abort_error(e) except Exception: logger.debug("Failed to evaluate OpenAI responses response", exc_info=True) + record_model_tool_calls(all_messages, len(request_messages)) return None diff --git a/ddtrace/contrib/internal/openai_agents/patch.py b/ddtrace/contrib/internal/openai_agents/patch.py index c2e6c7e574a..e93dd443b8a 100644 --- a/ddtrace/contrib/internal/openai_agents/patch.py +++ b/ddtrace/contrib/internal/openai_agents/patch.py @@ -1,4 +1,6 @@ +from contextvars import ContextVar import importlib +import inspect import agents from agents.tracing import add_trace_processor @@ -7,7 +9,9 @@ from ddtrace.contrib.internal.openai_agents.processor import LLMObsTraceProcessor from ddtrace.contrib.trace_utils import unwrap from ddtrace.contrib.trace_utils import wrap +from ddtrace.internal import core from ddtrace.internal.logger import get_logger +from ddtrace.internal.utils import get_argument_value from ddtrace.llmobs._integrations.openai_agents import OpenAIAgentsIntegration from ddtrace.trace import tracer @@ -65,6 +69,71 @@ def _has_module_level_run_loop() -> bool: ] +async def _patched_invoke_mcp_tool(func, instance, args, kwargs): + # Agent-driven MCP tool run: the model-visible tool and its call context are only known here. + core.dispatch( + "openai_agents.mcp.invoke_tool.before", + ( + get_argument_value(args, kwargs, 0, "server", optional=True), + get_argument_value(args, kwargs, 1, "tool", optional=True), + get_argument_value(args, kwargs, 2, "context", optional=True), + get_argument_value(args, kwargs, 3, "input_json", optional=True), + kwargs.get("tool_display_name"), + ), + allow_raise=True, + ) + return await func(*args, **kwargs) + + +# Set while a wrapped server call_tool runs, so an override calling super() dispatches once. +_in_mcp_server_call_tool: ContextVar[bool] = ContextVar("dd_openai_agents_in_mcp_call_tool", default=False) + + +async def _patched_mcp_server_call_tool(func, instance, args, kwargs): + # Lowest MCP call boundary of the SDK servers, also reached by direct server.call_tool calls. + if _in_mcp_server_call_tool.get(): + return await func(*args, **kwargs) + token = _in_mcp_server_call_tool.set(True) + try: + core.dispatch( + "openai_agents.mcp.call_tool.before", + ( + instance, + get_argument_value(args, kwargs, 0, "tool_name", optional=True), + get_argument_value(args, kwargs, 1, "arguments", optional=True), + ), + allow_raise=True, + ) + return await func(*args, **kwargs) + finally: + _in_mcp_server_call_tool.reset(token) + + +def _mcp_wrap_targets() -> list: + """(owner, attribute, wrapper) for the MCP adapter, empty when agents.mcp is unavailable. + + agents.mcp needs the optional mcp package, which is not installable on every Python version. + """ + try: + from agents.mcp import server as mcp_server + from agents.mcp.util import MCPUtil + except ImportError: + return [] + targets: list = [(MCPUtil, "invoke_mcp_tool", _patched_invoke_mcp_tool)] + # Concrete SDK servers may override call_tool without calling super(), so each class that + # implements it gets wrapped. The abstract MCPServer.call_tool is never reached. + for _, cls in inspect.getmembers(mcp_server, inspect.isclass): + call_tool = cls.__dict__.get("call_tool") + if ( + cls.__module__ == mcp_server.__name__ + and issubclass(cls, mcp_server.MCPServer) + and call_tool is not None + and not getattr(call_tool, "__isabstractmethod__", False) + ): + targets.append((cls, "call_tool", _patched_mcp_server_call_tool)) + return targets + + def patch(): """ Patch the instrumented methods @@ -91,6 +160,9 @@ def patch(): if hasattr(runner_cls, "_run_single_turn_streamed"): wrap(runner_cls, "_run_single_turn_streamed", _patched_run_single_turn) + for owner, attr_name, wrapper in _mcp_wrap_targets(): + wrap(owner, attr_name, wrapper) + def unpatch(): """ @@ -113,3 +185,6 @@ def unpatch(): unwrap(runner_cls, "_run_single_turn") if hasattr(runner_cls, "_run_single_turn_streamed"): unwrap(runner_cls, "_run_single_turn_streamed") + + for owner, attr_name, _ in _mcp_wrap_targets(): + unwrap(owner, attr_name) diff --git a/docs/configuration.rst b/docs/configuration.rst index 0776712cc38..9d733622223 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -849,8 +849,14 @@ AI Guard default: False description: | When set to True and AI Guard is enabled, AI Guard evaluates MCP tool calls with their MCP - server and tool details (server label, sanitized server URL, transport and tool name). - Credentials, headers and URL query strings are never sent. + server and tool details (configured server name, sanitized server URL, transport and + original tool name). Credentials, headers, URL query strings, stdio commands, arguments + and environment are never sent. + + Tools of MCP servers configured in the OpenAI Agents SDK (MCPServerStdio, MCPServerSse, + MCPServerStreamableHttp) are evaluated before each call reaches the server, whether the + agent runs the tool or the application calls server.call_tool directly. A blocking verdict + raises AIGuardAbortError and the call is not sent. For remote MCP tools run by the OpenAI Responses API (tools of type mcp), MCP approval requests are evaluated before your application sees them, and an approval your diff --git a/releasenotes/notes/aiguard-openai-agents-mcp-tool-calls-88cc75f75c24e907.yaml b/releasenotes/notes/aiguard-openai-agents-mcp-tool-calls-88cc75f75c24e907.yaml new file mode 100644 index 00000000000..1efaf08912f --- /dev/null +++ b/releasenotes/notes/aiguard-openai-agents-mcp-tool-calls-88cc75f75c24e907.yaml @@ -0,0 +1,7 @@ +--- +features: + - | + AI Guard: Evaluates the tool calls of MCP servers configured in the OpenAI Agents SDK before + each call reaches the server, including direct ``server.call_tool`` calls, when + ``DD_AI_GUARD_COLLECT_MCP_ENABLED=true``. A blocking verdict raises ``AIGuardAbortError`` and + the call is not sent. diff --git a/tests/aiguard/openai/conftest.py b/tests/aiguard/openai/conftest.py index 41789742a69..e0ffa9e876c 100644 --- a/tests/aiguard/openai/conftest.py +++ b/tests/aiguard/openai/conftest.py @@ -8,6 +8,7 @@ from ddtrace.aiguard._context import reset_aiguard_context_active from ddtrace.aiguard._context import set_aiguard_context_active from ddtrace.aiguard._initialization import load_ai_guard +from ddtrace.aiguard.integrations._mcp import _model_tool_calls from ddtrace.aiguard.integrations._mcp import approval_decisions from ddtrace.contrib.internal.openai.patch import patch from ddtrace.contrib.internal.openai.patch import unpatch @@ -619,9 +620,11 @@ def async_openai_responses_mcp_approval_client(openai_sdk, _require_responses_ap @pytest.fixture(autouse=True) def _reset_mcp_state(): - """MCP approval decisions are kept process-wide: isolate tests.""" + """MCP approval decisions are process-wide and model tool calls per context: isolate tests.""" approval_decisions.clear() + token = _model_tool_calls.set(None) yield + _model_tool_calls.reset(token) approval_decisions.clear() diff --git a/tests/aiguard/openai_agents/conftest.py b/tests/aiguard/openai_agents/conftest.py new file mode 100644 index 00000000000..f3d3236075e --- /dev/null +++ b/tests/aiguard/openai_agents/conftest.py @@ -0,0 +1,30 @@ +"""Fixtures for the AI Guard openai-agents MCP tests.""" + +import pytest + +from ddtrace.aiguard._initialization import load_ai_guard +from ddtrace.contrib.internal.openai_agents.patch import patch +from ddtrace.contrib.internal.openai_agents.patch import unpatch +from tests.aiguard.utils import override_ai_guard_config + + +@pytest.fixture(scope="session", autouse=True) +def _ai_guard_session_init(): + with override_ai_guard_config( + dict( + _ai_guard_enabled=True, + _ai_guard_endpoint="https://api.example.com/ai-guard", + _ai_guard_collect_mcp_enabled=True, + _dd_api_key="test-api-key", + _dd_app_key="test-application-key", + ) + ): + load_ai_guard() + yield + + +@pytest.fixture(autouse=True) +def openai_agents_patched(): + patch() + yield + unpatch() diff --git a/tests/aiguard/openai_agents/test_mcp.py b/tests/aiguard/openai_agents/test_mcp.py new file mode 100644 index 00000000000..3460258bf4d --- /dev/null +++ b/tests/aiguard/openai_agents/test_mcp.py @@ -0,0 +1,407 @@ +"""AI Guard evaluation of MCP tool calls run by openai-agents MCP servers.""" + +import asyncio +import json +import types +from unittest.mock import patch + +from agents.mcp import MCPServerSse +from agents.mcp import MCPServerStdio +from agents.mcp import MCPServerStreamableHttp +from agents.mcp.util import MCPUtil +from mcp.types import CallToolResult +from mcp.types import ListToolsResult +from mcp.types import TextContent +from mcp.types import Tool +import pytest + + +try: + # openai 3.x moved to httpx2, which keeps the httpx transport API. + import httpx2 as httpx +except ImportError: + import httpx + +from ddtrace.aiguard import AIGuardAbortError +from ddtrace.aiguard import new_ai_guard_client +from ddtrace.aiguard._constants import AI_GUARD +from ddtrace.aiguard._listener import _openai_listen +from ddtrace.aiguard.integrations._mcp import _model_tool_calls +from ddtrace.aiguard.integrations._openai_responses import _openai_response_create_after +from tests.aiguard.utils import mock_evaluate_response +from tests.aiguard.utils import override_ai_guard_config + + +EXECUTE_REQUEST = "ddtrace.aiguard._api_client.AIGuardClient._execute_request" +ATLASSIAN_MCP = { + "transport": "streamable_http", + "tool_name": "search_issues", + "name": "atlassian", + "url": "https://mcp.example.com/mcp", +} + + +class FakeSession: + """Stands in for the MCP client session so no server is needed; records the tools/call requests.""" + + def __init__(self): + self.calls = [] + + async def call_tool(self, name, arguments, *args, **kwargs): + self.calls.append((name, arguments)) + return CallToolResult(content=[TextContent(type="text", text="ok")]) + + +@pytest.fixture(autouse=True) +def _reset_model_tool_calls(): + token = _model_tool_calls.set(None) + yield + _model_tool_calls.reset(token) + + +def _connected(server): + server.session = FakeSession() + return server + + +def _atlassian_server(url="https://user:secret@MCP.Example.com:443/mcp?token=abc#frag", name="atlassian"): + return _connected( + MCPServerStreamableHttp(params={"url": url, "headers": {"Authorization": "Bearer secret"}}, name=name) + ) + + +def _tool(name="search_issues"): + return Tool(name=name, inputSchema={"type": "object", "properties": {}}) + + +def _tool_context(tool_call_id="call_1", tool_name="atlassian_search_issues", arguments='{"query": "APPSEC"}'): + try: + from agents.tool_context import ToolContext + + return ToolContext(context=None, tool_name=tool_name, tool_call_id=tool_call_id, tool_arguments=arguments) + except (ImportError, TypeError): + # Older SDKs pass a context without the model call details. + return types.SimpleNamespace(tool_call_id=tool_call_id, tool_name=tool_name) + + +def _invoke(server, tool=None, context=None, input_json='{"query": "APPSEC"}'): + return asyncio.run(MCPUtil.invoke_mcp_tool(server, tool or _tool(), context or _tool_context(), input_json)) + + +def _messages(mock_execute_request, index=-1): + return mock_execute_request.call_args_list[index].args[1]["data"]["attributes"]["messages"] + + +def _single_tool_call(messages): + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + (tool_call,) = messages[0]["tool_calls"] + return tool_call + + +@patch(EXECUTE_REQUEST) +def test_agent_driven_call_is_evaluated_once_with_mcp_identity(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + server = _atlassian_server() + + _invoke(server) + + assert server.session.calls == [("search_issues", {"query": "APPSEC"})] + assert mock_execute_request.call_count == 1 + assert _single_tool_call(_messages(mock_execute_request)) == { + "id": "call_1", + "function": {"name": "atlassian_search_issues", "arguments": '{"query": "APPSEC"}'}, + "mcp": ATLASSIAN_MCP, + } + + +@pytest.mark.parametrize("decision", ["DENY", "ABORT"]) +@patch(EXECUTE_REQUEST) +def test_blocking_verdict_prevents_the_tools_call(mock_execute_request, decision, test_spans): + mock_execute_request.return_value = mock_evaluate_response(decision, block=True) + server = _atlassian_server() + + with pytest.raises(AIGuardAbortError): + _invoke(server) + + assert server.session.calls == [] + (ai_guard_span,) = [span for span in test_spans.spans if span.name == AI_GUARD.RESOURCE_TYPE] + assert ai_guard_span.get_tag(AI_GUARD.BLOCKED_TAG) == "true" + assert ai_guard_span.get_tag(AI_GUARD.MCP_SERVER_NAME_TAG) == "atlassian" + assert ai_guard_span.get_tag(AI_GUARD.MCP_TOOL_NAME_TAG) == "search_issues" + assert ai_guard_span.get_tag(AI_GUARD.MCP_TRANSPORT_TAG) == "streamable_http" + assert ai_guard_span.get_tag(AI_GUARD.MCP_SERVER_URL_TAG) == "https://mcp.example.com/mcp" + + +@patch(EXECUTE_REQUEST) +def test_monitor_mode_and_evaluation_errors_preserve_execution(mock_execute_request): + server = _atlassian_server() + + mock_execute_request.return_value = mock_evaluate_response("DENY", block=False) + _invoke(server) + mock_execute_request.side_effect = ConnectionError("unreachable") + _invoke(server) + + assert len(server.session.calls) == 2 + + +@patch(EXECUTE_REQUEST) +def test_direct_server_call_uses_a_local_id(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + server = _atlassian_server() + + asyncio.run(server.call_tool("search_issues", {"query": "APPSEC"})) + + tool_call = _single_tool_call(_messages(mock_execute_request)) + assert tool_call["id"].startswith("dd_mcp_") + assert tool_call["function"] == {"name": "search_issues", "arguments": '{"query": "APPSEC"}'} + assert tool_call["mcp"] == ATLASSIAN_MCP + + +@patch(EXECUTE_REQUEST) +def test_direct_call_after_an_agent_call_is_evaluated(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + server = _atlassian_server() + + async def _run(): + await MCPUtil.invoke_mcp_tool(server, _tool(), _tool_context(), '{"query": "APPSEC"}') + await server.call_tool("search_issues", {"query": "APPSEC"}) + + asyncio.run(_run()) + + assert mock_execute_request.call_count == 2 + assert _single_tool_call(_messages(mock_execute_request))["id"].startswith("dd_mcp_") + + +@patch(EXECUTE_REQUEST) +def test_blocked_agent_call_does_not_excuse_a_direct_call(mock_execute_request): + server = _atlassian_server() + mock_execute_request.return_value = mock_evaluate_response("DENY", block=True) + + async def _run(): + with pytest.raises(AIGuardAbortError): + await MCPUtil.invoke_mcp_tool(server, _tool(), _tool_context(), '{"query": "APPSEC"}') + await server.call_tool("search_issues", {"query": "APPSEC"}) + + with pytest.raises(AIGuardAbortError): + asyncio.run(_run()) + + assert mock_execute_request.call_count == 2 + assert server.session.calls == [] + + +@patch(EXECUTE_REQUEST) +def test_model_call_reuses_the_conversation(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + response = { + "output": [ + { + "type": "function_call", + "call_id": "call_1", + "name": "atlassian_search_issues", + "arguments": '{"query": "APPSEC"}', + }, + {"type": "function_call", "call_id": "call_2", "name": "other_tool", "arguments": "{}"}, + ] + } + _openai_response_create_after(new_ai_guard_client(), {"input": "Find APPSEC issues"}, response) + + _invoke(_atlassian_server()) + + assert _messages(mock_execute_request) == [ + {"role": "user", "content": "Find APPSEC issues"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "atlassian_search_issues", "arguments": '{"query": "APPSEC"}'}, + "mcp": ATLASSIAN_MCP, + } + ], + }, + ] + + +@patch(EXECUTE_REQUEST) +def test_unknown_call_id_sends_only_the_tool_call(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + _openai_response_create_after( + new_ai_guard_client(), + {"input": "Find APPSEC issues"}, + {"output": [{"type": "function_call", "call_id": "call_9", "name": "x", "arguments": "{}"}]}, + ) + + _invoke(_atlassian_server()) + + assert _single_tool_call(_messages(mock_execute_request))["id"] == "call_1" + + +@patch(EXECUTE_REQUEST) +def test_same_tool_on_two_servers_keeps_each_identity(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + + _invoke(_atlassian_server(name="jira")) + _invoke(_atlassian_server(url="https://wiki.example.com:8443/mcp", name="confluence")) + + first, second = (_single_tool_call(_messages(mock_execute_request, i))["mcp"] for i in range(2)) + assert (first["name"], first["url"]) == ("jira", "https://mcp.example.com/mcp") + assert (second["name"], second["url"]) == ("confluence", "https://wiki.example.com:8443/mcp") + + +@patch(EXECUTE_REQUEST) +def test_stdio_server_reports_no_command_arguments_or_environment(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + server = _connected( + MCPServerStdio(params={"command": "uvx", "args": ["mcp-server", "--token", "secret-arg"], "env": {"K": "v"}}) + ) + + _invoke(server) + + payload = json.dumps(mock_execute_request.call_args_list[-1].args[1]) + assert _single_tool_call(_messages(mock_execute_request))["mcp"] == { + "transport": "stdio", + "tool_name": "search_issues", + } + for secret in ("uvx", "secret-arg", '"K"'): + assert secret not in payload + + +@patch(EXECUTE_REQUEST) +def test_generated_server_name_is_not_reported(mock_execute_request): + mock_execute_request.return_value = mock_evaluate_response("ALLOW") + server = _connected(MCPServerSse(params={"url": "https://user:secret@sse.example.com/sse?key=abc"})) + + _invoke(server) + + payload = json.dumps(mock_execute_request.call_args_list[-1].args[1]) + assert _single_tool_call(_messages(mock_execute_request))["mcp"] == { + "transport": "sse", + "tool_name": "search_issues", + "url": "https://sse.example.com/sse", + } + for secret in ("secret", "key=abc", "Bearer"): + assert secret not in payload + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_listeners_are_gated(enabled): + with override_ai_guard_config(dict(_ai_guard_collect_mcp_enabled=enabled)): + with patch("ddtrace.aiguard._listener.core.on") as core_on: + _openai_listen(new_ai_guard_client()) + + events = {c.args[0] for c in core_on.call_args_list} + mcp_events = {"openai_agents.mcp.invoke_tool.before", "openai_agents.mcp.call_tool.before"} + assert events & mcp_events == (mcp_events if enabled else set()) + + +def _responses_body(output): + return { + "id": "resp-test", + "object": "response", + "created_at": 0, + "model": "gpt-4o-mini", + "status": "completed", + "output": output, + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "metadata": {}, + "parallel_tool_calls": True, + "temperature": 1.0, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + } + + +class _ModelTransport(httpx.AsyncBaseTransport): + """First turn: the model calls the MCP tool. Second turn: it answers.""" + + def __init__(self): + self.turns = 0 + + async def handle_async_request(self, request): + self.turns += 1 + if self.turns == 1: + output = [ + { + "id": "fc_1", + "type": "function_call", + "call_id": "call_run", + "name": "search_issues", + "arguments": '{"query": "APPSEC"}', + "status": "completed", + } + ] + else: + output = [ + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ] + return httpx.Response(200, json=_responses_body(output)) + + +class FakeListingSession(FakeSession): + async def list_tools(self, *args, **kwargs): + return ListToolsResult(tools=[_tool()]) + + +@pytest.mark.parametrize("decision,blocked", [("ALLOW", False), ("DENY", True)]) +@patch(EXECUTE_REQUEST) +def test_runner_evaluates_mcp_tool_before_it_runs(mock_execute_request, decision, blocked): + from agents import Agent + from agents import OpenAIResponsesModel + from agents import Runner + import openai + + server = MCPServerStreamableHttp(params={"url": "https://mcp.example.com/mcp"}, name="atlassian") + server.session = FakeListingSession() + client = openai.AsyncOpenAI(api_key="", http_client=httpx.AsyncClient(transport=_ModelTransport())) + agent = Agent( + name="assistant", + model=OpenAIResponsesModel(model="gpt-4o-mini", openai_client=client), + mcp_servers=[server], + ) + + def _evaluate(url, payload): + tool_calls = payload["data"]["attributes"]["messages"][-1].get("tool_calls") or [] + is_mcp = any("mcp" in tool_call for tool_call in tool_calls) + return mock_evaluate_response(decision if is_mcp else "ALLOW", block=True) + + mock_execute_request.side_effect = _evaluate + + if blocked: + with pytest.raises(AIGuardAbortError): + asyncio.run(Runner.run(agent, "Find APPSEC issues")) + assert server.session.calls == [] + else: + result = asyncio.run(Runner.run(agent, "Find APPSEC issues")) + assert result.final_output == "done" + assert server.session.calls == [("search_issues", {"query": "APPSEC"})] + + mcp_evaluations = [ + c.args[1]["data"]["attributes"]["messages"] + for c in mock_execute_request.call_args_list + if any( + "mcp" in tool_call for tool_call in c.args[1]["data"]["attributes"]["messages"][-1].get("tool_calls") or [] + ) + ] + assert len(mcp_evaluations) == 1 + (mcp_tool_call,) = mcp_evaluations[0][-1]["tool_calls"] + assert mcp_tool_call["id"] == "call_run" + assert mcp_tool_call["mcp"] == { + "transport": "streamable_http", + "tool_name": "search_issues", + "name": "atlassian", + "url": "https://mcp.example.com/mcp", + } diff --git a/tests/aiguard/suitespec.yml b/tests/aiguard/suitespec.yml index 38ebc64be79..6b97f9abb97 100644 --- a/tests/aiguard/suitespec.yml +++ b/tests/aiguard/suitespec.yml @@ -75,6 +75,27 @@ suites: dependencies: ['pytest-asyncio==0.23.7', 'openai==1.102.0'] - name: ai_guard_openai dependencies: ['pytest-asyncio==0.23.7', openai] + ai_guard_openai_agents: + paths: + - '@bootstrap' + - '@core' + - '@tracing' + - '@ai_guard' + - '@openai_agents' + - ddtrace/internal/utils/http.py + - tests/aiguard/openai_agents/* + - tests/aiguard/utils.py + - tests/aiguard/suitespec.yml + matrix: + command: pytest {cmdargs} tests/aiguard/openai_agents/ + env: + DD_TRACE_PY_ENABLE_ITR_TEST_SKIPPING_FOR_JOB: 'true' + python: ['3.10', '3.11', '3.12', '3.13'] + variants: + - name: ai_guard_openai_agents + dependencies: ['pytest-asyncio==0.23.7', 'openai-agents~=0.0.0', 'openai<1.100'] + - name: ai_guard_openai_agents + dependencies: ['pytest-asyncio==0.23.7', openai-agents] ai_guard_anthropic: paths: - '@bootstrap'