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.
@@ -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,7 @@
1
+ __pycache__/
2
+ *.pyc
3
+ .pytest_cache/
4
+ dist/
5
+ build/
6
+ *.egg-info/
7
+ .venv/
@@ -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,4 @@
1
+ from .tools import check_mcp_server_risk, check_prompt_injection_breach
2
+
3
+ __all__ = ["check_mcp_server_risk", "check_prompt_injection_breach"]
4
+ __version__ = "0.1.0"
@@ -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