langchain-relayshield 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- langchain_relayshield-0.1.0/.github/workflows/publish.yml +65 -0
- langchain_relayshield-0.1.0/.gitignore +7 -0
- langchain_relayshield-0.1.0/PKG-INFO +88 -0
- langchain_relayshield-0.1.0/README.md +58 -0
- langchain_relayshield-0.1.0/langchain_relayshield/__init__.py +4 -0
- langchain_relayshield-0.1.0/langchain_relayshield/client.py +64 -0
- langchain_relayshield-0.1.0/langchain_relayshield/middleware.py +263 -0
- langchain_relayshield-0.1.0/langchain_relayshield/tools.py +66 -0
- langchain_relayshield-0.1.0/pyproject.toml +46 -0
- langchain_relayshield-0.1.0/tests/test_middleware.py +145 -0
- langchain_relayshield-0.1.0/tests/test_tools.py +71 -0
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
build:
|
|
9
|
+
name: Build distribution
|
|
10
|
+
runs-on: ubuntu-latest
|
|
11
|
+
steps:
|
|
12
|
+
- uses: actions/checkout@v4
|
|
13
|
+
|
|
14
|
+
- name: Set up Python
|
|
15
|
+
uses: actions/setup-python@v5
|
|
16
|
+
with:
|
|
17
|
+
python-version: "3.12"
|
|
18
|
+
|
|
19
|
+
- name: Install hatchling
|
|
20
|
+
run: pip install hatchling
|
|
21
|
+
|
|
22
|
+
- name: Build package
|
|
23
|
+
run: python -m hatchling build
|
|
24
|
+
|
|
25
|
+
- name: Upload build artifacts
|
|
26
|
+
uses: actions/upload-artifact@v4
|
|
27
|
+
with:
|
|
28
|
+
name: dist
|
|
29
|
+
path: dist/
|
|
30
|
+
|
|
31
|
+
test:
|
|
32
|
+
name: Run tests
|
|
33
|
+
runs-on: ubuntu-latest
|
|
34
|
+
steps:
|
|
35
|
+
- uses: actions/checkout@v4
|
|
36
|
+
|
|
37
|
+
- name: Set up Python
|
|
38
|
+
uses: actions/setup-python@v5
|
|
39
|
+
with:
|
|
40
|
+
python-version: "3.12"
|
|
41
|
+
|
|
42
|
+
- name: Install package with test extras
|
|
43
|
+
run: pip install -e ".[test]"
|
|
44
|
+
|
|
45
|
+
- name: Run tests
|
|
46
|
+
run: pytest tests/ -v
|
|
47
|
+
|
|
48
|
+
publish:
|
|
49
|
+
name: Publish to PyPI
|
|
50
|
+
needs: [build, test]
|
|
51
|
+
runs-on: ubuntu-latest
|
|
52
|
+
environment: release
|
|
53
|
+
permissions:
|
|
54
|
+
id-token: write # Required for OIDC Trusted Publishing
|
|
55
|
+
|
|
56
|
+
steps:
|
|
57
|
+
- name: Download build artifacts
|
|
58
|
+
uses: actions/download-artifact@v4
|
|
59
|
+
with:
|
|
60
|
+
name: dist
|
|
61
|
+
path: dist/
|
|
62
|
+
|
|
63
|
+
- name: Publish to PyPI
|
|
64
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
65
|
+
# No API token needed — OIDC handles authentication
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: langchain-relayshield
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: LangChain tools and a mandatory pre-execution gate for RelayShield's MCP registry risk and prompt-injection breach checks.
|
|
5
|
+
Project-URL: Homepage, https://relayshield.net
|
|
6
|
+
Project-URL: Documentation, https://api.relayshield.net/developers
|
|
7
|
+
Project-URL: Repository, https://github.com/nzdsf2-gif/langchain-relayshield
|
|
8
|
+
Author-email: RelayShield <relayshieldadmin@gmail.com>
|
|
9
|
+
License: MIT
|
|
10
|
+
Keywords: agents,ai,breach,langchain,mcp,security
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Topic :: Internet
|
|
19
|
+
Classifier: Topic :: Security
|
|
20
|
+
Requires-Python: >=3.10
|
|
21
|
+
Requires-Dist: langchain-core>=0.3.0
|
|
22
|
+
Provides-Extra: middleware
|
|
23
|
+
Requires-Dist: langchain>=1.0.0; extra == 'middleware'
|
|
24
|
+
Requires-Dist: langgraph>=0.2.0; extra == 'middleware'
|
|
25
|
+
Provides-Extra: test
|
|
26
|
+
Requires-Dist: langchain>=1.0.0; extra == 'test'
|
|
27
|
+
Requires-Dist: langgraph>=0.2.0; extra == 'test'
|
|
28
|
+
Requires-Dist: pytest>=8.0.0; extra == 'test'
|
|
29
|
+
Description-Content-Type: text/markdown
|
|
30
|
+
|
|
31
|
+
# langchain-relayshield
|
|
32
|
+
|
|
33
|
+
LangChain tools and a mandatory pre-execution gate for [RelayShield](https://relayshield.net)'s agentic-security endpoints — MCP server registry risk and AI-agent-sourced credential breach detection.
|
|
34
|
+
|
|
35
|
+
## Install
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
pip install langchain-relayshield
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
The mandatory-gate middleware needs LangChain's agent runtime; install the extra if you want it:
|
|
42
|
+
|
|
43
|
+
```bash
|
|
44
|
+
pip install "langchain-relayshield[middleware]"
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
## Tools
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
from langchain_relayshield import check_mcp_server_risk, check_prompt_injection_breach
|
|
51
|
+
|
|
52
|
+
result = check_mcp_server_risk.invoke({
|
|
53
|
+
"api_key": "rs_live_...",
|
|
54
|
+
"server_url": "https://mcp.example.com/sse",
|
|
55
|
+
})
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
- **`check_mcp_server_risk`** — flags known-malicious IOC matches, typosquat domains, and newly-registered domains hosting an MCP server, before an agent connects to or installs it.
|
|
59
|
+
- **`check_prompt_injection_breach`** — checks whether an email appears in RelayShield's stolen-session corpus with a suspected-agentic-source marker (a session/token exposure that shows signs of having been captured via a compromised AI agent).
|
|
60
|
+
|
|
61
|
+
Get a key at [api.relayshield.net/developers](https://api.relayshield.net/developers).
|
|
62
|
+
|
|
63
|
+
## Mandatory gate middleware
|
|
64
|
+
|
|
65
|
+
Most "AI agent security" checks are optional — the agent *can* call them, but nothing stops it skipping the call and taking the risky action anyway. `RelayShieldMCPGateMiddleware` is the other kind: a gate the agent's framework enforces *before* a protected action (connecting to or installing an MCP server) can happen at all.
|
|
66
|
+
|
|
67
|
+
```python
|
|
68
|
+
from langchain_relayshield.middleware import RelayShieldMCPGateMiddleware
|
|
69
|
+
|
|
70
|
+
agent = create_agent(
|
|
71
|
+
model=...,
|
|
72
|
+
tools=[...],
|
|
73
|
+
middleware=[RelayShieldMCPGateMiddleware()],
|
|
74
|
+
)
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
Properties, all non-negotiable by design:
|
|
78
|
+
|
|
79
|
+
- A hook exception defaults to `defer`, never silently to `allow`.
|
|
80
|
+
- Bounded retry applies only to transient upstream failures (timeout/429/5xx) — auth failures, malformed responses, and payment-required states are terminal after one attempt.
|
|
81
|
+
- The audit log records decision, reason codes, check version, target, and timestamp — never keys, payment proofs, or session material.
|
|
82
|
+
- The raw connect/install tool should never be bound to the model directly — only a composite tool that routes through this gate.
|
|
83
|
+
|
|
84
|
+
Full design writeup and a 12-case protected-action acceptance suite: [relayshield-langchain-gate](https://github.com/nzdsf2-gif/relayshield-langchain-gate) (this middleware's original standalone reference implementation, now also packaged here).
|
|
85
|
+
|
|
86
|
+
## License
|
|
87
|
+
|
|
88
|
+
MIT
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
# langchain-relayshield
|
|
2
|
+
|
|
3
|
+
LangChain tools and a mandatory pre-execution gate for [RelayShield](https://relayshield.net)'s agentic-security endpoints — MCP server registry risk and AI-agent-sourced credential breach detection.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install langchain-relayshield
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
The mandatory-gate middleware needs LangChain's agent runtime; install the extra if you want it:
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
pip install "langchain-relayshield[middleware]"
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
## Tools
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from langchain_relayshield import check_mcp_server_risk, check_prompt_injection_breach
|
|
21
|
+
|
|
22
|
+
result = check_mcp_server_risk.invoke({
|
|
23
|
+
"api_key": "rs_live_...",
|
|
24
|
+
"server_url": "https://mcp.example.com/sse",
|
|
25
|
+
})
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
- **`check_mcp_server_risk`** — flags known-malicious IOC matches, typosquat domains, and newly-registered domains hosting an MCP server, before an agent connects to or installs it.
|
|
29
|
+
- **`check_prompt_injection_breach`** — checks whether an email appears in RelayShield's stolen-session corpus with a suspected-agentic-source marker (a session/token exposure that shows signs of having been captured via a compromised AI agent).
|
|
30
|
+
|
|
31
|
+
Get a key at [api.relayshield.net/developers](https://api.relayshield.net/developers).
|
|
32
|
+
|
|
33
|
+
## Mandatory gate middleware
|
|
34
|
+
|
|
35
|
+
Most "AI agent security" checks are optional — the agent *can* call them, but nothing stops it skipping the call and taking the risky action anyway. `RelayShieldMCPGateMiddleware` is the other kind: a gate the agent's framework enforces *before* a protected action (connecting to or installing an MCP server) can happen at all.
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
from langchain_relayshield.middleware import RelayShieldMCPGateMiddleware
|
|
39
|
+
|
|
40
|
+
agent = create_agent(
|
|
41
|
+
model=...,
|
|
42
|
+
tools=[...],
|
|
43
|
+
middleware=[RelayShieldMCPGateMiddleware()],
|
|
44
|
+
)
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Properties, all non-negotiable by design:
|
|
48
|
+
|
|
49
|
+
- A hook exception defaults to `defer`, never silently to `allow`.
|
|
50
|
+
- Bounded retry applies only to transient upstream failures (timeout/429/5xx) — auth failures, malformed responses, and payment-required states are terminal after one attempt.
|
|
51
|
+
- The audit log records decision, reason codes, check version, target, and timestamp — never keys, payment proofs, or session material.
|
|
52
|
+
- The raw connect/install tool should never be bound to the model directly — only a composite tool that routes through this gate.
|
|
53
|
+
|
|
54
|
+
Full design writeup and a 12-case protected-action acceptance suite: [relayshield-langchain-gate](https://github.com/nzdsf2-gif/relayshield-langchain-gate) (this middleware's original standalone reference implementation, now also packaged here).
|
|
55
|
+
|
|
56
|
+
## License
|
|
57
|
+
|
|
58
|
+
MIT
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Shared HTTP client for RelayShield's agentic-security endpoints.
|
|
3
|
+
|
|
4
|
+
Deliberately dependency-light (stdlib urllib only) so this module has zero
|
|
5
|
+
transitive footprint beyond langchain-core itself -- the same convention as
|
|
6
|
+
relayshield-mcp and relayshield_langchain_gate.py.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import json
|
|
12
|
+
import os
|
|
13
|
+
import urllib.error
|
|
14
|
+
import urllib.request
|
|
15
|
+
|
|
16
|
+
RELAYSHIELD_API_BASE = os.environ.get(
|
|
17
|
+
"RELAYSHIELD_API_BASE", "https://atq6wtkp6k.execute-api.us-east-1.amazonaws.com/prod"
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class RelayShieldAPIError(Exception):
|
|
22
|
+
"""Raised for any non-2xx / malformed response. Callers decide how to
|
|
23
|
+
surface this to the agent -- Tools return it as content, the gate
|
|
24
|
+
middleware maps it to a CheckState (see middleware.py)."""
|
|
25
|
+
|
|
26
|
+
def __init__(self, message: str, status: int | None = None):
|
|
27
|
+
self.status = status
|
|
28
|
+
super().__init__(message)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def relayshield_post(path: str, body: dict, api_key: str | None = None, timeout: float = 8.0) -> dict:
|
|
32
|
+
"""POST to a RelayShield metered endpoint and return the unwrapped `data`
|
|
33
|
+
payload. Raises RelayShieldAPIError on any failure -- never returns a
|
|
34
|
+
partial/guessed result on error."""
|
|
35
|
+
api_key = api_key or os.environ.get("RELAYSHIELD_API_KEY", "")
|
|
36
|
+
if not api_key:
|
|
37
|
+
raise RelayShieldAPIError("No RelayShield API key configured (RELAYSHIELD_API_KEY or api_key arg)", status=None)
|
|
38
|
+
|
|
39
|
+
payload = json.dumps({k: v for k, v in body.items() if v is not None}).encode("utf-8")
|
|
40
|
+
req = urllib.request.Request(
|
|
41
|
+
f"{RELAYSHIELD_API_BASE}{path}",
|
|
42
|
+
data=payload,
|
|
43
|
+
headers={"Content-Type": "application/json", "X-RS-API-KEY": api_key},
|
|
44
|
+
method="POST",
|
|
45
|
+
)
|
|
46
|
+
try:
|
|
47
|
+
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
48
|
+
raw, status = resp.read(), resp.status
|
|
49
|
+
except urllib.error.HTTPError as exc:
|
|
50
|
+
raise RelayShieldAPIError(f"RelayShield API returned HTTP {exc.code}", status=exc.code) from exc
|
|
51
|
+
except (urllib.error.URLError, TimeoutError) as exc:
|
|
52
|
+
raise RelayShieldAPIError(f"RelayShield API request failed: {exc}", status=None) from exc
|
|
53
|
+
|
|
54
|
+
if not raw:
|
|
55
|
+
raise RelayShieldAPIError("RelayShield API returned an empty body", status=status)
|
|
56
|
+
try:
|
|
57
|
+
data = json.loads(raw)
|
|
58
|
+
except (json.JSONDecodeError, UnicodeDecodeError) as exc:
|
|
59
|
+
raise RelayShieldAPIError(f"RelayShield API returned malformed JSON: {exc}", status=status) from exc
|
|
60
|
+
|
|
61
|
+
if not isinstance(data, dict) or not data.get("ok"):
|
|
62
|
+
raise RelayShieldAPIError(f"RelayShield API error: {data}", status=status)
|
|
63
|
+
|
|
64
|
+
return data.get("data", {})
|
|
@@ -0,0 +1,263 @@
|
|
|
1
|
+
"""
|
|
2
|
+
RelayShieldMCPGateMiddleware -- a mandatory, pre-execution gate for
|
|
3
|
+
connect/install actions on an MCP server, built on LangChain's
|
|
4
|
+
`wrap_tool_call` middleware hook.
|
|
5
|
+
|
|
6
|
+
This is the packaged version of relayshield-langchain-gate
|
|
7
|
+
(github.com/nzdsf2-gif/relayshield-langchain-gate) -- same logic, same
|
|
8
|
+
12-case acceptance suite, now distributed as part of this pip-installable
|
|
9
|
+
package instead of a standalone reference repo. See that repo's README for
|
|
10
|
+
the full design rationale (John6666's normalized gate policy,
|
|
11
|
+
HF discuss.huggingface.co, smolagents-agent-security-tools-v2 thread).
|
|
12
|
+
|
|
13
|
+
Requires the `middleware` extra: pip install langchain-relayshield[middleware]
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from __future__ import annotations
|
|
17
|
+
|
|
18
|
+
import json
|
|
19
|
+
import logging
|
|
20
|
+
import os
|
|
21
|
+
import time
|
|
22
|
+
import urllib.error
|
|
23
|
+
import urllib.parse
|
|
24
|
+
import urllib.request
|
|
25
|
+
from dataclasses import dataclass, field
|
|
26
|
+
from datetime import datetime, timezone
|
|
27
|
+
from enum import Enum
|
|
28
|
+
from typing import Any
|
|
29
|
+
|
|
30
|
+
logger = logging.getLogger("langchain_relayshield.gate")
|
|
31
|
+
|
|
32
|
+
RELAYSHIELD_API_BASE = os.environ.get(
|
|
33
|
+
"RELAYSHIELD_API_BASE", "https://atq6wtkp6k.execute-api.us-east-1.amazonaws.com/prod"
|
|
34
|
+
)
|
|
35
|
+
RELAYSHIELD_CHECK_PATH = "/v1/metered/mcp-registry-risk"
|
|
36
|
+
|
|
37
|
+
MAX_RETRIES = 2
|
|
38
|
+
RETRY_BACKOFF_SECONDS = 1.0
|
|
39
|
+
|
|
40
|
+
# The raw connect/install tools are intentionally never bound to the model
|
|
41
|
+
# directly in a real deployment -- only a composite tool that routes through
|
|
42
|
+
# this gate should be exposed. See relayshield-langchain-gate's README for
|
|
43
|
+
# why this is an architectural requirement, not just a naming convention.
|
|
44
|
+
PROTECTED_TOOLS = {
|
|
45
|
+
"connect_mcp_server",
|
|
46
|
+
"install_mcp_package",
|
|
47
|
+
}
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
class CheckState(str, Enum):
|
|
51
|
+
FINDING = "finding"
|
|
52
|
+
NO_KNOWN_FINDING = "no_known_finding"
|
|
53
|
+
UNKNOWN = "unknown"
|
|
54
|
+
PARTIAL = "partial"
|
|
55
|
+
STALE = "stale"
|
|
56
|
+
AUTH_FAILURE = "auth_failure"
|
|
57
|
+
UPSTREAM_FAILURE = "upstream_failure"
|
|
58
|
+
MALFORMED = "malformed"
|
|
59
|
+
PAYMENT_REQUIRED = "payment_required"
|
|
60
|
+
MISSING = "missing"
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
class GateAction(str, Enum):
|
|
64
|
+
ALLOW = "allow"
|
|
65
|
+
REVIEW = "review"
|
|
66
|
+
DENY = "deny"
|
|
67
|
+
DEFER = "defer"
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
@dataclass(frozen=True)
|
|
71
|
+
class GateDecision:
|
|
72
|
+
action: GateAction
|
|
73
|
+
reason_codes: list[str]
|
|
74
|
+
check_version: str
|
|
75
|
+
observed_at: str
|
|
76
|
+
detail: str = ""
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
@dataclass(frozen=True)
|
|
80
|
+
class CheckResult:
|
|
81
|
+
state: CheckState
|
|
82
|
+
verdict: str | None = None
|
|
83
|
+
findings: list[dict] = field(default_factory=list)
|
|
84
|
+
raw_status: int | None = None
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
class RelayShieldCheckError(Exception):
|
|
88
|
+
def __init__(self, kind: str, detail: str = "", status: int | None = None):
|
|
89
|
+
self.kind = kind
|
|
90
|
+
self.detail = detail
|
|
91
|
+
self.status = status
|
|
92
|
+
super().__init__(f"{kind}: {detail}")
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def check_mcp_registry_risk(
|
|
96
|
+
server_url: str | None = None,
|
|
97
|
+
package_name: str | None = None,
|
|
98
|
+
api_key: str | None = None,
|
|
99
|
+
timeout: float = 8.0,
|
|
100
|
+
) -> CheckResult:
|
|
101
|
+
api_key = api_key or os.environ.get("RELAYSHIELD_API_KEY", "")
|
|
102
|
+
if not api_key:
|
|
103
|
+
raise RelayShieldCheckError("auth_failure", "no RELAYSHIELD_API_KEY configured")
|
|
104
|
+
|
|
105
|
+
body = json.dumps(
|
|
106
|
+
{k: v for k, v in {"server_url": server_url, "package_name": package_name}.items() if v}
|
|
107
|
+
).encode("utf-8")
|
|
108
|
+
req = urllib.request.Request(
|
|
109
|
+
f"{RELAYSHIELD_API_BASE}{RELAYSHIELD_CHECK_PATH}",
|
|
110
|
+
data=body,
|
|
111
|
+
headers={"Content-Type": "application/json", "X-RS-API-KEY": api_key},
|
|
112
|
+
method="POST",
|
|
113
|
+
)
|
|
114
|
+
try:
|
|
115
|
+
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
116
|
+
raw, status = resp.read(), resp.status
|
|
117
|
+
except urllib.error.HTTPError as exc:
|
|
118
|
+
raw, status = exc.read(), exc.code
|
|
119
|
+
if status == 401:
|
|
120
|
+
raise RelayShieldCheckError("auth_failure", detail=raw.decode(errors="replace"), status=status) from exc
|
|
121
|
+
if status == 402:
|
|
122
|
+
raise RelayShieldCheckError("payment_required", detail=raw.decode(errors="replace"), status=status) from exc
|
|
123
|
+
if status == 429:
|
|
124
|
+
raise RelayShieldCheckError("upstream_failure", detail="rate_limited", status=status) from exc
|
|
125
|
+
if status >= 500:
|
|
126
|
+
raise RelayShieldCheckError("upstream_failure", detail=raw.decode(errors="replace"), status=status) from exc
|
|
127
|
+
raise RelayShieldCheckError("malformed", detail=f"unexpected status {status}", status=status) from exc
|
|
128
|
+
except (urllib.error.URLError, TimeoutError) as exc:
|
|
129
|
+
raise RelayShieldCheckError("upstream_failure", detail=str(exc)) from exc
|
|
130
|
+
|
|
131
|
+
if not raw:
|
|
132
|
+
raise RelayShieldCheckError("malformed", detail="empty body", status=status)
|
|
133
|
+
try:
|
|
134
|
+
data = json.loads(raw)
|
|
135
|
+
except (ValueError, UnicodeDecodeError) as exc:
|
|
136
|
+
raise RelayShieldCheckError("malformed", detail=str(exc), status=status) from exc
|
|
137
|
+
if not isinstance(data, dict) or "verdict" not in data:
|
|
138
|
+
raise RelayShieldCheckError("malformed", detail="response missing 'verdict'", status=status)
|
|
139
|
+
|
|
140
|
+
verdict = data.get("verdict", "LOW")
|
|
141
|
+
findings = data.get("findings", [])
|
|
142
|
+
|
|
143
|
+
if findings:
|
|
144
|
+
return CheckResult(state=CheckState.FINDING, verdict=verdict, findings=findings, raw_status=status)
|
|
145
|
+
if server_url is None and package_name is not None:
|
|
146
|
+
return CheckResult(state=CheckState.PARTIAL, verdict=verdict, raw_status=status)
|
|
147
|
+
return CheckResult(state=CheckState.NO_KNOWN_FINDING, verdict=verdict, raw_status=status)
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
def check_mcp_registry_risk_with_retry(**kwargs) -> tuple[CheckState, CheckResult | None, RelayShieldCheckError | None]:
|
|
151
|
+
last_err: RelayShieldCheckError | None = None
|
|
152
|
+
for attempt in range(MAX_RETRIES + 1):
|
|
153
|
+
try:
|
|
154
|
+
result = check_mcp_registry_risk(**kwargs)
|
|
155
|
+
return result.state, result, None
|
|
156
|
+
except RelayShieldCheckError as exc:
|
|
157
|
+
last_err = exc
|
|
158
|
+
if exc.kind != "upstream_failure" or attempt == MAX_RETRIES:
|
|
159
|
+
return CheckState(exc.kind), None, exc
|
|
160
|
+
time.sleep(RETRY_BACKOFF_SECONDS * (attempt + 1))
|
|
161
|
+
return CheckState.UPSTREAM_FAILURE, None, last_err
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
_POLICY: dict[CheckState, GateAction] = {
|
|
165
|
+
CheckState.FINDING: GateAction.DENY,
|
|
166
|
+
CheckState.NO_KNOWN_FINDING: GateAction.ALLOW,
|
|
167
|
+
CheckState.UNKNOWN: GateAction.REVIEW,
|
|
168
|
+
CheckState.PARTIAL: GateAction.REVIEW,
|
|
169
|
+
CheckState.STALE: GateAction.REVIEW,
|
|
170
|
+
CheckState.AUTH_FAILURE: GateAction.DEFER,
|
|
171
|
+
CheckState.UPSTREAM_FAILURE: GateAction.DEFER,
|
|
172
|
+
CheckState.MALFORMED: GateAction.DEFER,
|
|
173
|
+
CheckState.PAYMENT_REQUIRED: GateAction.DEFER,
|
|
174
|
+
CheckState.MISSING: GateAction.DEFER,
|
|
175
|
+
}
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def evaluate_gate_policy(state: CheckState, result: CheckResult | None = None) -> GateDecision:
|
|
179
|
+
action = _POLICY.get(state, GateAction.DEFER)
|
|
180
|
+
reason_codes = [state.value]
|
|
181
|
+
if result and result.findings:
|
|
182
|
+
reason_codes += [f["type"] for f in result.findings]
|
|
183
|
+
return GateDecision(
|
|
184
|
+
action=action,
|
|
185
|
+
reason_codes=reason_codes,
|
|
186
|
+
check_version="mcp-registry-risk-v1",
|
|
187
|
+
observed_at=datetime.now(timezone.utc).isoformat(),
|
|
188
|
+
detail=(result.verdict if result else state.value),
|
|
189
|
+
)
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def record_gate_decision(target: dict, decision: GateDecision) -> None:
|
|
193
|
+
logger.info(
|
|
194
|
+
"relayshield_gate_decision target=%s action=%s reasons=%s check_version=%s observed_at=%s",
|
|
195
|
+
{k: v for k, v in target.items() if v},
|
|
196
|
+
decision.action.value,
|
|
197
|
+
decision.reason_codes,
|
|
198
|
+
decision.check_version,
|
|
199
|
+
decision.observed_at,
|
|
200
|
+
)
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
def _deny_message(tool_call: dict, decision: GateDecision):
|
|
204
|
+
from langchain_core.messages import ToolMessage
|
|
205
|
+
|
|
206
|
+
return ToolMessage(
|
|
207
|
+
content=(
|
|
208
|
+
f"Blocked by RelayShield registry-risk gate: action={decision.action.value} "
|
|
209
|
+
f"reasons={decision.reason_codes}. The connection/install did not occur. "
|
|
210
|
+
"Local reading, analysis, and drafting remain available."
|
|
211
|
+
),
|
|
212
|
+
name=tool_call["name"],
|
|
213
|
+
tool_call_id=tool_call["id"],
|
|
214
|
+
status="error",
|
|
215
|
+
)
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
def build_gate_middleware():
|
|
219
|
+
"""Constructs the middleware class lazily so this module imports cleanly
|
|
220
|
+
even without langchain/langgraph installed (base package has no hard
|
|
221
|
+
dependency on them -- only the `middleware` extra does)."""
|
|
222
|
+
from langchain.agents.middleware import AgentMiddleware
|
|
223
|
+
|
|
224
|
+
class RelayShieldMCPGateMiddleware(AgentMiddleware):
|
|
225
|
+
"""Mandatory pre-execution gate for connect_mcp_server / install_mcp_package.
|
|
226
|
+
|
|
227
|
+
A hook exception defaults to defer, never silently to allow --
|
|
228
|
+
failing open would defeat the point of a mandatory gate.
|
|
229
|
+
"""
|
|
230
|
+
|
|
231
|
+
def wrap_tool_call(self, request, handler):
|
|
232
|
+
tool_call = request.tool_call
|
|
233
|
+
if tool_call["name"] not in PROTECTED_TOOLS:
|
|
234
|
+
return handler(request)
|
|
235
|
+
|
|
236
|
+
args = tool_call.get("args", {})
|
|
237
|
+
target = {"server_url": args.get("server_url"), "package_name": args.get("package_name")}
|
|
238
|
+
|
|
239
|
+
try:
|
|
240
|
+
state, result, err = check_mcp_registry_risk_with_retry(**target)
|
|
241
|
+
decision = evaluate_gate_policy(state, result)
|
|
242
|
+
except Exception: # noqa: BLE001 -- gate failure must not become allow
|
|
243
|
+
logger.exception("relayshield_gate: unhandled exception, defaulting to defer")
|
|
244
|
+
decision = GateDecision(
|
|
245
|
+
action=GateAction.DEFER,
|
|
246
|
+
reason_codes=["gate_internal_exception"],
|
|
247
|
+
check_version="mcp-registry-risk-v1",
|
|
248
|
+
observed_at=datetime.now(timezone.utc).isoformat(),
|
|
249
|
+
)
|
|
250
|
+
|
|
251
|
+
record_gate_decision(target, decision)
|
|
252
|
+
|
|
253
|
+
if decision.action == GateAction.ALLOW:
|
|
254
|
+
return handler(request)
|
|
255
|
+
return _deny_message(tool_call, decision)
|
|
256
|
+
|
|
257
|
+
return RelayShieldMCPGateMiddleware
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
try:
|
|
261
|
+
RelayShieldMCPGateMiddleware = build_gate_middleware()
|
|
262
|
+
except ImportError:
|
|
263
|
+
RelayShieldMCPGateMiddleware = None # type: ignore[assignment]
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
"""
|
|
2
|
+
LangChain tools wrapping RelayShield's agentic-security endpoints.
|
|
3
|
+
|
|
4
|
+
Both tools take `api_key` as a call argument rather than reading it from the
|
|
5
|
+
environment implicitly -- same pattern as relayshield-mcp and the
|
|
6
|
+
smolagents.Tool classes, since a shared agent process may act on behalf of
|
|
7
|
+
multiple callers with different RelayShield keys.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
from langchain_core.tools import tool
|
|
13
|
+
|
|
14
|
+
from .client import RelayShieldAPIError, relayshield_post
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@tool
|
|
18
|
+
def check_mcp_server_risk(
|
|
19
|
+
api_key: str,
|
|
20
|
+
server_url: str | None = None,
|
|
21
|
+
package_name: str | None = None,
|
|
22
|
+
) -> dict:
|
|
23
|
+
"""Check an MCP server or package for registry risk before connecting to
|
|
24
|
+
or installing it. Flags known-malicious IOC matches, typosquat domains
|
|
25
|
+
close to well-known MCP ecosystem names, and newly-registered domains
|
|
26
|
+
hosting the server. Provide server_url, package_name, or both -- checks
|
|
27
|
+
are more limited with package_name alone (no domain to check).
|
|
28
|
+
|
|
29
|
+
Returns a dict with `verdict` (CRITICAL/HIGH/MEDIUM/LOW) and `findings`
|
|
30
|
+
(a list of specific reasons). Absence of a finding is not proof of
|
|
31
|
+
safety -- MCP registry security coverage is minimal industry-wide.
|
|
32
|
+
"""
|
|
33
|
+
if not server_url and not package_name:
|
|
34
|
+
return {"error": "server_url or package_name is required"}
|
|
35
|
+
try:
|
|
36
|
+
return relayshield_post(
|
|
37
|
+
"/v1/metered/mcp-registry-risk",
|
|
38
|
+
{"server_url": server_url, "package_name": package_name},
|
|
39
|
+
api_key=api_key,
|
|
40
|
+
)
|
|
41
|
+
except RelayShieldAPIError as exc:
|
|
42
|
+
return {"error": str(exc)}
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
@tool
|
|
46
|
+
def check_prompt_injection_breach(api_key: str, email: str) -> dict:
|
|
47
|
+
"""Check whether an email address appears in RelayShield's stolen-session
|
|
48
|
+
corpus with a suspected-agentic-source marker -- i.e. a session/token
|
|
49
|
+
exposure that shows signs of having been captured via a compromised or
|
|
50
|
+
hijacked AI agent, not a conventional phishing/malware infostealer path.
|
|
51
|
+
|
|
52
|
+
Returns a dict with `found` (bool), `session_count`, and `sessions`
|
|
53
|
+
(domain/severity/service_category per match). This is a heuristic
|
|
54
|
+
keyword classifier over dump-announcement text, not confirmed
|
|
55
|
+
attribution -- a clean result does not rule out agent involvement.
|
|
56
|
+
"""
|
|
57
|
+
if not email or "@" not in email:
|
|
58
|
+
return {"error": "email is required and must be a valid address"}
|
|
59
|
+
try:
|
|
60
|
+
return relayshield_post(
|
|
61
|
+
"/v1/metered/prompt-injection-breach",
|
|
62
|
+
{"email": email},
|
|
63
|
+
api_key=api_key,
|
|
64
|
+
)
|
|
65
|
+
except RelayShieldAPIError as exc:
|
|
66
|
+
return {"error": str(exc)}
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "langchain-relayshield"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "LangChain tools and a mandatory pre-execution gate for RelayShield's MCP registry risk and prompt-injection breach checks."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "MIT" }
|
|
11
|
+
authors = [{ name = "RelayShield", email = "relayshieldadmin@gmail.com" }]
|
|
12
|
+
keywords = ["langchain", "security", "mcp", "agents", "ai", "breach"]
|
|
13
|
+
classifiers = [
|
|
14
|
+
"Development Status :: 4 - Beta",
|
|
15
|
+
"Intended Audience :: Developers",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Programming Language :: Python :: 3",
|
|
18
|
+
"Programming Language :: Python :: 3.10",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Programming Language :: Python :: 3.12",
|
|
21
|
+
"Topic :: Security",
|
|
22
|
+
"Topic :: Internet",
|
|
23
|
+
]
|
|
24
|
+
requires-python = ">=3.10"
|
|
25
|
+
dependencies = [
|
|
26
|
+
"langchain-core>=0.3.0",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
[project.optional-dependencies]
|
|
30
|
+
middleware = [
|
|
31
|
+
"langchain>=1.0.0",
|
|
32
|
+
"langgraph>=0.2.0",
|
|
33
|
+
]
|
|
34
|
+
test = [
|
|
35
|
+
"langchain>=1.0.0",
|
|
36
|
+
"langgraph>=0.2.0",
|
|
37
|
+
"pytest>=8.0.0",
|
|
38
|
+
]
|
|
39
|
+
|
|
40
|
+
[project.urls]
|
|
41
|
+
Homepage = "https://relayshield.net"
|
|
42
|
+
Documentation = "https://api.relayshield.net/developers"
|
|
43
|
+
Repository = "https://github.com/nzdsf2-gif/langchain-relayshield"
|
|
44
|
+
|
|
45
|
+
[tool.hatch.build.targets.wheel]
|
|
46
|
+
packages = ["langchain_relayshield"]
|
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Acceptance suite for langchain_relayshield.middleware -- ported from
|
|
3
|
+
relayshield-langchain-gate's test suite (12/12 passing there against real
|
|
4
|
+
langchain/langgraph). Asserts on the *protected action* (whether handler()
|
|
5
|
+
-- the real MCP connect/install -- was invoked), not merely the returned
|
|
6
|
+
label, per the original design's core requirement.
|
|
7
|
+
|
|
8
|
+
Requires: pip install langchain-relayshield[test]
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from unittest.mock import patch
|
|
12
|
+
|
|
13
|
+
import pytest
|
|
14
|
+
from langchain_core.messages import ToolMessage
|
|
15
|
+
from langgraph.prebuilt.tool_node import ToolCallRequest
|
|
16
|
+
|
|
17
|
+
from langchain_relayshield.middleware import (
|
|
18
|
+
CheckResult,
|
|
19
|
+
CheckState,
|
|
20
|
+
RelayShieldCheckError,
|
|
21
|
+
RelayShieldMCPGateMiddleware,
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def _make_request(tool_name="connect_mcp_server", server_url="https://example-mcp-server.test"):
|
|
26
|
+
return ToolCallRequest(
|
|
27
|
+
tool_call={"name": tool_name, "args": {"server_url": server_url}, "id": "call_1"},
|
|
28
|
+
tool=None,
|
|
29
|
+
state={},
|
|
30
|
+
runtime=None,
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _run_gate(mock_return):
|
|
35
|
+
middleware = RelayShieldMCPGateMiddleware()
|
|
36
|
+
request = _make_request()
|
|
37
|
+
handler_called = {"value": False}
|
|
38
|
+
|
|
39
|
+
def handler(req):
|
|
40
|
+
handler_called["value"] = True
|
|
41
|
+
return ToolMessage(content="connected", name=req.tool_call["name"], tool_call_id=req.tool_call["id"])
|
|
42
|
+
|
|
43
|
+
with patch("langchain_relayshield.middleware.check_mcp_registry_risk_with_retry") as mock_check:
|
|
44
|
+
if isinstance(mock_return, RelayShieldCheckError):
|
|
45
|
+
mock_check.return_value = (CheckState(mock_return.kind), None, mock_return)
|
|
46
|
+
else:
|
|
47
|
+
mock_check.return_value = mock_return
|
|
48
|
+
result = middleware.wrap_tool_call(request, handler)
|
|
49
|
+
return handler_called["value"], result
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def test_synthetic_known_risk_server_blocks_connection():
|
|
53
|
+
result = CheckResult(state=CheckState.FINDING, verdict="CRITICAL", findings=[{"type": "typosquat_suspected"}])
|
|
54
|
+
called, msg = _run_gate((CheckState.FINDING, result, None))
|
|
55
|
+
assert called is False
|
|
56
|
+
assert msg.status == "error"
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def test_benign_fresh_sufficient_coverage_allows():
|
|
60
|
+
result = CheckResult(state=CheckState.NO_KNOWN_FINDING, verdict="LOW")
|
|
61
|
+
called, _ = _run_gate((CheckState.NO_KNOWN_FINDING, result, None))
|
|
62
|
+
assert called is True
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def test_no_known_finding_with_partial_coverage_reviews_not_allows():
|
|
66
|
+
result = CheckResult(state=CheckState.PARTIAL, verdict="LOW")
|
|
67
|
+
called, _ = _run_gate((CheckState.PARTIAL, result, None))
|
|
68
|
+
assert called is False
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def test_stale_result_defers_to_review():
|
|
72
|
+
result = CheckResult(state=CheckState.STALE, verdict="LOW")
|
|
73
|
+
called, _ = _run_gate((CheckState.STALE, result, None))
|
|
74
|
+
assert called is False
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def test_401_blocks_protected_action():
|
|
78
|
+
called, _ = _run_gate(RelayShieldCheckError("auth_failure", status=401))
|
|
79
|
+
assert called is False
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def test_429_defers_no_false_allow():
|
|
83
|
+
called, _ = _run_gate(RelayShieldCheckError("upstream_failure", status=429))
|
|
84
|
+
assert called is False
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def test_500_blocks_protected_action():
|
|
88
|
+
called, _ = _run_gate(RelayShieldCheckError("upstream_failure", status=500))
|
|
89
|
+
assert called is False
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def test_malformed_response_blocks_protected_action():
|
|
93
|
+
called, _ = _run_gate(RelayShieldCheckError("malformed", detail="invalid JSON"))
|
|
94
|
+
assert called is False
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def test_network_disconnect_blocks_protected_action():
|
|
98
|
+
called, _ = _run_gate(RelayShieldCheckError("upstream_failure", detail="connection refused"))
|
|
99
|
+
assert called is False
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
def test_402_payment_required_defers_no_protected_action_yet():
|
|
103
|
+
called, _ = _run_gate(RelayShieldCheckError("payment_required", status=402))
|
|
104
|
+
assert called is False
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_gate_internal_exception_blocks_protected_action_not_allow():
|
|
108
|
+
middleware = RelayShieldMCPGateMiddleware()
|
|
109
|
+
request = _make_request()
|
|
110
|
+
handler_called = {"value": False}
|
|
111
|
+
|
|
112
|
+
def handler(req):
|
|
113
|
+
handler_called["value"] = True
|
|
114
|
+
return ToolMessage(content="connected", name=req.tool_call["name"], tool_call_id=req.tool_call["id"])
|
|
115
|
+
|
|
116
|
+
with patch("langchain_relayshield.middleware.check_mcp_registry_risk_with_retry", side_effect=RuntimeError("boom")):
|
|
117
|
+
result = middleware.wrap_tool_call(request, handler)
|
|
118
|
+
assert handler_called["value"] is False
|
|
119
|
+
assert result.status == "error"
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def test_unrelated_tool_is_unaffected_by_gate():
|
|
123
|
+
middleware = RelayShieldMCPGateMiddleware()
|
|
124
|
+
request = ToolCallRequest(
|
|
125
|
+
tool_call={"name": "read_local_file", "args": {"path": "/tmp/x"}, "id": "call_2"},
|
|
126
|
+
tool=None,
|
|
127
|
+
state={},
|
|
128
|
+
runtime=None,
|
|
129
|
+
)
|
|
130
|
+
handler_called = {"value": False}
|
|
131
|
+
|
|
132
|
+
def handler(req):
|
|
133
|
+
handler_called["value"] = True
|
|
134
|
+
return ToolMessage(content="ok", name=req.tool_call["name"], tool_call_id=req.tool_call["id"])
|
|
135
|
+
|
|
136
|
+
with patch("langchain_relayshield.middleware.check_mcp_registry_risk_with_retry") as mock_check:
|
|
137
|
+
result = middleware.wrap_tool_call(request, handler)
|
|
138
|
+
mock_check.assert_not_called()
|
|
139
|
+
assert handler_called["value"] is True
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
if __name__ == "__main__":
|
|
143
|
+
import sys
|
|
144
|
+
|
|
145
|
+
sys.exit(pytest.main([__file__, "-v"]))
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Tests for the two LangChain tools -- only the RelayShield HTTP call is
|
|
3
|
+
mocked (via client.relayshield_post), not any langchain_core internals.
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from unittest.mock import patch
|
|
7
|
+
|
|
8
|
+
from langchain_relayshield.client import RelayShieldAPIError
|
|
9
|
+
from langchain_relayshield.tools import check_mcp_server_risk, check_prompt_injection_breach
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def test_mcp_server_risk_requires_target():
|
|
13
|
+
result = check_mcp_server_risk.invoke({"api_key": "rs_live_test"})
|
|
14
|
+
assert "error" in result
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def test_mcp_server_risk_returns_finding():
|
|
18
|
+
with patch("langchain_relayshield.tools.relayshield_post") as mock_post:
|
|
19
|
+
mock_post.return_value = {"verdict": "CRITICAL", "findings": [{"type": "typosquat_suspected"}]}
|
|
20
|
+
result = check_mcp_server_risk.invoke({
|
|
21
|
+
"api_key": "rs_live_test",
|
|
22
|
+
"server_url": "https://suspicious-mcp.test",
|
|
23
|
+
})
|
|
24
|
+
assert result["verdict"] == "CRITICAL"
|
|
25
|
+
mock_post.assert_called_once()
|
|
26
|
+
args, kwargs = mock_post.call_args
|
|
27
|
+
assert args[0] == "/v1/metered/mcp-registry-risk"
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def test_mcp_server_risk_surfaces_api_error_as_content_not_exception():
|
|
31
|
+
"""A RelayShieldAPIError must come back as tool content (`{"error": ...}`),
|
|
32
|
+
not propagate as an uncaught exception -- an agent framework should be
|
|
33
|
+
able to show the model what went wrong rather than crashing the run."""
|
|
34
|
+
with patch("langchain_relayshield.tools.relayshield_post", side_effect=RelayShieldAPIError("HTTP 500", status=500)):
|
|
35
|
+
result = check_mcp_server_risk.invoke({
|
|
36
|
+
"api_key": "rs_live_test",
|
|
37
|
+
"server_url": "https://example.test",
|
|
38
|
+
})
|
|
39
|
+
assert "error" in result
|
|
40
|
+
assert "500" in result["error"]
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def test_prompt_injection_breach_requires_email():
|
|
44
|
+
result = check_prompt_injection_breach.invoke({"api_key": "rs_live_test", "email": "not-an-email"})
|
|
45
|
+
assert "error" in result
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_prompt_injection_breach_returns_no_finding():
|
|
49
|
+
with patch("langchain_relayshield.tools.relayshield_post") as mock_post:
|
|
50
|
+
mock_post.return_value = {"email": "user@example.com", "found": False, "session_count": 0, "sessions": []}
|
|
51
|
+
result = check_prompt_injection_breach.invoke({
|
|
52
|
+
"api_key": "rs_live_test",
|
|
53
|
+
"email": "user@example.com",
|
|
54
|
+
})
|
|
55
|
+
assert result["found"] is False
|
|
56
|
+
args, kwargs = mock_post.call_args
|
|
57
|
+
assert args[0] == "/v1/metered/prompt-injection-breach"
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def test_prompt_injection_breach_returns_finding():
|
|
61
|
+
with patch("langchain_relayshield.tools.relayshield_post") as mock_post:
|
|
62
|
+
mock_post.return_value = {
|
|
63
|
+
"email": "user@example.com", "found": True, "session_count": 1,
|
|
64
|
+
"sessions": [{"domain": "example.com", "severity": "HIGH", "service_category": "cloud"}],
|
|
65
|
+
}
|
|
66
|
+
result = check_prompt_injection_breach.invoke({
|
|
67
|
+
"api_key": "rs_live_test",
|
|
68
|
+
"email": "user@example.com",
|
|
69
|
+
})
|
|
70
|
+
assert result["found"] is True
|
|
71
|
+
assert result["session_count"] == 1
|