pattern-agent-sdk 0.1.0__tar.gz → 0.3.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.
- {pattern_agent_sdk-0.1.0 → pattern_agent_sdk-0.3.0}/PKG-INFO +4 -1
- {pattern_agent_sdk-0.1.0 → pattern_agent_sdk-0.3.0}/pyproject.toml +4 -1
- pattern_agent_sdk-0.3.0/src/pattern_agent_sdk/__init__.py +4 -0
- pattern_agent_sdk-0.3.0/src/pattern_agent_sdk/sdk.py +316 -0
- pattern_agent_sdk-0.3.0/tests/test_sdk.py +335 -0
- pattern_agent_sdk-0.1.0/src/pattern_agent_sdk/__init__.py +0 -3
- pattern_agent_sdk-0.1.0/src/pattern_agent_sdk/sdk.py +0 -164
- pattern_agent_sdk-0.1.0/tests/test_sdk.py +0 -138
- {pattern_agent_sdk-0.1.0 → pattern_agent_sdk-0.3.0}/.gitignore +0 -0
- {pattern_agent_sdk-0.1.0 → pattern_agent_sdk-0.3.0}/LICENSE.md +0 -0
- {pattern_agent_sdk-0.1.0 → pattern_agent_sdk-0.3.0}/README.md +0 -0
- {pattern_agent_sdk-0.1.0 → pattern_agent_sdk-0.3.0}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: pattern_agent_sdk
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: SDK for Pattern Agentic agents
|
|
5
5
|
Author-email: Amos Joshua <amos@patternagentic.ai>
|
|
6
6
|
License-File: LICENSE.md
|
|
@@ -9,6 +9,9 @@ Classifier: Operating System :: OS Independent
|
|
|
9
9
|
Classifier: Programming Language :: Python :: 3
|
|
10
10
|
Requires-Python: >=3.11
|
|
11
11
|
Requires-Dist: httpx>=0.27.0
|
|
12
|
+
Requires-Dist: pattern-agentic-messaging>=1.2.0
|
|
13
|
+
Requires-Dist: pydantic>=2.0.0
|
|
14
|
+
Requires-Dist: pyjwt[crypto]>=2.8.0
|
|
12
15
|
Description-Content-Type: text/markdown
|
|
13
16
|
|
|
14
17
|
# pattern_agentic_sdk
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "pattern_agent_sdk"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.3.0"
|
|
8
8
|
description = "SDK for Pattern Agentic agents"
|
|
9
9
|
authors = [
|
|
10
10
|
{ name="Amos Joshua", email="amos@patternagentic.ai" }
|
|
@@ -18,6 +18,9 @@ classifiers = [
|
|
|
18
18
|
]
|
|
19
19
|
dependencies = [
|
|
20
20
|
"httpx>=0.27.0",
|
|
21
|
+
"PyJWT[crypto]>=2.8.0",
|
|
22
|
+
"pydantic>=2.0.0",
|
|
23
|
+
"pattern-agentic-messaging>=1.2.0",
|
|
21
24
|
]
|
|
22
25
|
|
|
23
26
|
[tool.hatch.build.targets.wheel]
|
|
@@ -0,0 +1,316 @@
|
|
|
1
|
+
import asyncio
|
|
2
|
+
import os
|
|
3
|
+
import json
|
|
4
|
+
import time
|
|
5
|
+
import base64
|
|
6
|
+
import logging
|
|
7
|
+
import tempfile
|
|
8
|
+
from typing import Optional
|
|
9
|
+
|
|
10
|
+
import httpx
|
|
11
|
+
import jwt as pyjwt
|
|
12
|
+
from jwt import PyJWKClient
|
|
13
|
+
from pydantic import BaseModel
|
|
14
|
+
from pattern_agentic_messaging import PASlimApp, PASlimConfig
|
|
15
|
+
from pattern_agentic_messaging.session_token import PatternAgentSessionToken
|
|
16
|
+
|
|
17
|
+
logger = logging.getLogger(__name__)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _parse_extra_headers(raw: Optional[str]) -> dict[str, str]:
|
|
21
|
+
if not raw:
|
|
22
|
+
return {}
|
|
23
|
+
try:
|
|
24
|
+
decoded = base64.b64decode(raw)
|
|
25
|
+
headers = json.loads(decoded)
|
|
26
|
+
if isinstance(headers, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in headers.items()):
|
|
27
|
+
return headers
|
|
28
|
+
except Exception:
|
|
29
|
+
pass
|
|
30
|
+
logger.warning("PatternSDK: __PA_AGENT_SERVICES_TOKEN_REQUEST_HEADERS is malformed, ignoring")
|
|
31
|
+
return {}
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _decode_jwt_payload(token: str) -> dict: # no sig check: token just received from local sidecar
|
|
35
|
+
parts = token.split(".")
|
|
36
|
+
if len(parts) != 3:
|
|
37
|
+
raise ValueError("Invalid JWT format")
|
|
38
|
+
payload_b64 = parts[1] + "=" * (-len(parts[1]) % 4)
|
|
39
|
+
return json.loads(base64.urlsafe_b64decode(payload_b64))
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class SessionTokenJWKSConfig(BaseModel):
|
|
43
|
+
jwks_url: str
|
|
44
|
+
issuer: Optional[str] = None
|
|
45
|
+
audience: str
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
class PatternSDK:
|
|
50
|
+
def __init__(self, services_endpoint: Optional[str] = None, token_request_headers: Optional[str] = None):
|
|
51
|
+
self._services_endpoint = (services_endpoint or "").rstrip("/")
|
|
52
|
+
self._enabled = bool(self._services_endpoint)
|
|
53
|
+
if not self._enabled:
|
|
54
|
+
logger.info("PatternSDK: no services endpoint configured, SDK disabled")
|
|
55
|
+
return
|
|
56
|
+
self._mint_headers = _parse_extra_headers(token_request_headers)
|
|
57
|
+
self._token: Optional[str] = None
|
|
58
|
+
self._deployment_id: Optional[str] = None
|
|
59
|
+
self._token_exp: float = 0
|
|
60
|
+
self._renew_failed: bool = False
|
|
61
|
+
self._config_cache: Optional[dict] = None
|
|
62
|
+
self._prompts_cache: Optional[dict] = None
|
|
63
|
+
self._persist_token_to_disk: bool = False
|
|
64
|
+
self._token_persist_path: Optional[str] = None
|
|
65
|
+
self._session_jwks_config: Optional[SessionTokenJWKSConfig] = None
|
|
66
|
+
self._jwks_client: Optional[PyJWKClient] = None
|
|
67
|
+
|
|
68
|
+
def _require_enabled(self):
|
|
69
|
+
if not self._enabled:
|
|
70
|
+
raise RuntimeError("PatternSDK is disabled: services endpoint not specified")
|
|
71
|
+
|
|
72
|
+
@property
|
|
73
|
+
def _endpoint(self) -> str:
|
|
74
|
+
return self._services_endpoint
|
|
75
|
+
|
|
76
|
+
async def _mint_token(self) -> str:
|
|
77
|
+
url = f"{self._endpoint}/mint-agent-token"
|
|
78
|
+
try:
|
|
79
|
+
async with httpx.AsyncClient(timeout=10) as client:
|
|
80
|
+
resp = await client.post(url, headers=self._mint_headers or None)
|
|
81
|
+
except httpx.ConnectError as e:
|
|
82
|
+
logger.error(f"PatternSDK: cannot connect to agent services at {url} — is the sidecar running? {e}")
|
|
83
|
+
raise
|
|
84
|
+
except httpx.TimeoutException:
|
|
85
|
+
logger.error(f"PatternSDK: request to {url} timed out")
|
|
86
|
+
raise
|
|
87
|
+
|
|
88
|
+
if resp.status_code != 200:
|
|
89
|
+
logger.error(f"PatternSDK: mint-agent-token returned {resp.status_code}: {resp.text}")
|
|
90
|
+
raise RuntimeError(
|
|
91
|
+
f"Error authorizing agent for pattern-agentic services: {resp.status_code} {resp.reason_phrase}"
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
try:
|
|
95
|
+
token = resp.json()["token"]
|
|
96
|
+
except (KeyError, json.JSONDecodeError) as e:
|
|
97
|
+
logger.error(f"PatternSDK: mint-agent-token response missing 'token' field: {resp.text}")
|
|
98
|
+
raise RuntimeError("mint-agent-token response missing 'token' field") from e
|
|
99
|
+
|
|
100
|
+
payload = _decode_jwt_payload(token)
|
|
101
|
+
self._token = token
|
|
102
|
+
self._deployment_id = payload["sub"]
|
|
103
|
+
self._token_exp = payload.get("exp", 0)
|
|
104
|
+
self._renew_failed = False
|
|
105
|
+
if self._persist_token_to_disk:
|
|
106
|
+
self._write_token_to_disk()
|
|
107
|
+
return token
|
|
108
|
+
|
|
109
|
+
async def _ensure_token(self):
|
|
110
|
+
if not self._token or time.time() >= self._token_exp:
|
|
111
|
+
await self._mint_token()
|
|
112
|
+
|
|
113
|
+
def _token_is_expired(self) -> bool:
|
|
114
|
+
return time.time() >= self._token_exp
|
|
115
|
+
|
|
116
|
+
def _write_token_to_disk(self):
|
|
117
|
+
path = self._token_persist_path
|
|
118
|
+
parent = os.path.dirname(path)
|
|
119
|
+
os.makedirs(parent, exist_ok=True)
|
|
120
|
+
fd, tmp = tempfile.mkstemp(dir=parent)
|
|
121
|
+
try:
|
|
122
|
+
os.write(fd, self._token.encode())
|
|
123
|
+
os.fchmod(fd, 0o600)
|
|
124
|
+
finally:
|
|
125
|
+
os.close(fd)
|
|
126
|
+
try:
|
|
127
|
+
os.rename(tmp, path)
|
|
128
|
+
except BaseException:
|
|
129
|
+
os.unlink(tmp)
|
|
130
|
+
raise
|
|
131
|
+
|
|
132
|
+
async def enable_token_persistence(self):
|
|
133
|
+
self._require_enabled()
|
|
134
|
+
path = os.environ.get("__PA_AGENT_SDK_TOKEN_PERSIST_PATH")
|
|
135
|
+
if not path:
|
|
136
|
+
raise RuntimeError(
|
|
137
|
+
"Invalid configuration, env variable __PA_AGENT_SDK_TOKEN_PERSIST_PATH is required to access the token on disk"
|
|
138
|
+
)
|
|
139
|
+
if os.path.isdir(path):
|
|
140
|
+
path = os.path.join(path, f"agent-{os.getpid()}.jwt")
|
|
141
|
+
self._token_persist_path = path
|
|
142
|
+
self._persist_token_to_disk = True
|
|
143
|
+
await self._ensure_token()
|
|
144
|
+
self._write_token_to_disk()
|
|
145
|
+
|
|
146
|
+
@property
|
|
147
|
+
def token_file_path(self) -> str:
|
|
148
|
+
if not self._persist_token_to_disk:
|
|
149
|
+
raise RuntimeError("Token persistence has not been enabled — call enable_token_persistence() first")
|
|
150
|
+
return self._token_persist_path
|
|
151
|
+
|
|
152
|
+
async def _request(self, method: str, path: str) -> dict:
|
|
153
|
+
await self._ensure_token()
|
|
154
|
+
url = f"{self._endpoint}/deployment/{self._deployment_id}{path}"
|
|
155
|
+
headers = {"Authorization": f"Bearer {self._token}"}
|
|
156
|
+
|
|
157
|
+
try:
|
|
158
|
+
async with httpx.AsyncClient(timeout=10) as client:
|
|
159
|
+
resp = await client.request(method, url, headers=headers)
|
|
160
|
+
except httpx.ConnectError as e:
|
|
161
|
+
logger.error(f"PatternSDK: cannot connect to {url}: {e}")
|
|
162
|
+
raise
|
|
163
|
+
except httpx.TimeoutException:
|
|
164
|
+
logger.error(f"PatternSDK: request to {url} timed out")
|
|
165
|
+
raise
|
|
166
|
+
|
|
167
|
+
if resp.status_code == 403 and self._token_is_expired() and not self._renew_failed:
|
|
168
|
+
logger.info("PatternSDK: got 403 with expired token, attempting renewal")
|
|
169
|
+
try:
|
|
170
|
+
await self._mint_token()
|
|
171
|
+
except Exception:
|
|
172
|
+
self._renew_failed = True
|
|
173
|
+
logger.error("PatternSDK: token renewal failed after 403, not retrying")
|
|
174
|
+
raise
|
|
175
|
+
headers = {"Authorization": f"Bearer {self._token}"}
|
|
176
|
+
try:
|
|
177
|
+
async with httpx.AsyncClient(timeout=10) as client:
|
|
178
|
+
resp = await client.request(method, url, headers=headers)
|
|
179
|
+
except httpx.ConnectError as e:
|
|
180
|
+
logger.error(f"PatternSDK: cannot connect to {url} on retry: {e}")
|
|
181
|
+
raise
|
|
182
|
+
except httpx.TimeoutException:
|
|
183
|
+
logger.error(f"PatternSDK: retry request to {url} timed out")
|
|
184
|
+
raise
|
|
185
|
+
|
|
186
|
+
if resp.status_code == 401:
|
|
187
|
+
logger.error(f"PatternSDK: 401 from {path}: {resp.text}")
|
|
188
|
+
raise PermissionError(f"Unauthorized: {resp.text}")
|
|
189
|
+
if resp.status_code == 403:
|
|
190
|
+
logger.error(f"PatternSDK: 403 from {path}: {resp.text}")
|
|
191
|
+
raise PermissionError(f"Forbidden: {resp.text}")
|
|
192
|
+
if resp.status_code >= 400:
|
|
193
|
+
logger.error(f"PatternSDK: {resp.status_code} from {path}: {resp.text}")
|
|
194
|
+
raise RuntimeError(
|
|
195
|
+
f"Error fetching {path} from pattern-agentic services: {resp.status_code} {resp.reason_phrase}"
|
|
196
|
+
)
|
|
197
|
+
|
|
198
|
+
return resp.json()
|
|
199
|
+
|
|
200
|
+
async def session_token_config(self) -> SessionTokenJWKSConfig:
|
|
201
|
+
if self._session_jwks_config:
|
|
202
|
+
return self._session_jwks_config
|
|
203
|
+
self._require_enabled()
|
|
204
|
+
url = f"{self._endpoint}/session-token-config"
|
|
205
|
+
try:
|
|
206
|
+
async with httpx.AsyncClient(timeout=10) as client:
|
|
207
|
+
resp = await client.get(url)
|
|
208
|
+
except httpx.ConnectError as e:
|
|
209
|
+
logger.error(f"PatternSDK: cannot connect to {url}: {e}")
|
|
210
|
+
raise
|
|
211
|
+
except httpx.TimeoutException:
|
|
212
|
+
logger.error(f"PatternSDK: request to {url} timed out")
|
|
213
|
+
raise
|
|
214
|
+
if resp.status_code >= 400:
|
|
215
|
+
logger.error(f"PatternSDK: {resp.status_code} from session-token-config: {resp.text}")
|
|
216
|
+
raise RuntimeError(f"Failed to fetch session token config: {resp.status_code}")
|
|
217
|
+
data = resp.json()
|
|
218
|
+
self._session_jwks_config = SessionTokenJWKSConfig(**data)
|
|
219
|
+
self._jwks_client = PyJWKClient(
|
|
220
|
+
self._session_jwks_config.jwks_url,
|
|
221
|
+
cache_keys=True,
|
|
222
|
+
)
|
|
223
|
+
return self._session_jwks_config
|
|
224
|
+
|
|
225
|
+
async def verify_session_token(self, token, *, verify_agent_peers: bool = True) -> PatternAgentSessionToken:
|
|
226
|
+
if isinstance(token, PatternAgentSessionToken):
|
|
227
|
+
raw = token.raw_token
|
|
228
|
+
else:
|
|
229
|
+
raw = token
|
|
230
|
+
if not self._jwks_client or not self._session_jwks_config:
|
|
231
|
+
await self.session_token_config()
|
|
232
|
+
signing_key = self._jwks_client.get_signing_key_from_jwt(raw)
|
|
233
|
+
decode_kwargs = {
|
|
234
|
+
"jwt": raw,
|
|
235
|
+
"key": signing_key.key,
|
|
236
|
+
"algorithms": ["RS256"],
|
|
237
|
+
"audience": self._session_jwks_config.audience,
|
|
238
|
+
}
|
|
239
|
+
if self._session_jwks_config.issuer:
|
|
240
|
+
decode_kwargs["issuer"] = self._session_jwks_config.issuer
|
|
241
|
+
payload = pyjwt.decode(**decode_kwargs)
|
|
242
|
+
verified = PatternAgentSessionToken(
|
|
243
|
+
session_id=payload["sub"],
|
|
244
|
+
tenant_id=payload["tenant_id"],
|
|
245
|
+
user_id=payload["user_id"],
|
|
246
|
+
agents=payload.get("agents", []),
|
|
247
|
+
exp=payload["exp"],
|
|
248
|
+
iat=payload["iat"],
|
|
249
|
+
raw_token=raw,
|
|
250
|
+
)
|
|
251
|
+
if verify_agent_peers:
|
|
252
|
+
if not self._deployment_id:
|
|
253
|
+
raise RuntimeError("Cannot verify agent peers: SDK has no deployment identity (token not yet minted)")
|
|
254
|
+
if self._deployment_id not in verified.agents:
|
|
255
|
+
raise PermissionError(
|
|
256
|
+
f"Session token does not authorize this agent: "
|
|
257
|
+
f"deployment={self._deployment_id} not in agents={verified.agents}"
|
|
258
|
+
)
|
|
259
|
+
return verified
|
|
260
|
+
|
|
261
|
+
async def _build_agent_app(self, *, verify_sessions: bool = True) -> PASlimApp:
|
|
262
|
+
self._require_enabled()
|
|
263
|
+
await self._ensure_token()
|
|
264
|
+
cfg = await self.config()
|
|
265
|
+
sys_params = cfg.get("__system_params", {})
|
|
266
|
+
slim_endpoint = sys_params.get("PA_SLIM_ENDPOINT", "")
|
|
267
|
+
slim_name = sys_params.get("PA_SLIM_NAME", "")
|
|
268
|
+
slim_secret = sys_params.get("PA_SLIM_SHARED_SECRET", "")
|
|
269
|
+
if not slim_endpoint or not slim_name or not slim_secret:
|
|
270
|
+
raise RuntimeError(
|
|
271
|
+
f"Missing SLIM config from agent services: "
|
|
272
|
+
f"PA_SLIM_ENDPOINT={slim_endpoint!r}, PA_SLIM_NAME={slim_name!r}, "
|
|
273
|
+
f"PA_SLIM_SHARED_SECRET={'***' if slim_secret else '(empty)'}"
|
|
274
|
+
)
|
|
275
|
+
await self.session_token_config()
|
|
276
|
+
|
|
277
|
+
config = PASlimConfig(
|
|
278
|
+
local_name=slim_name,
|
|
279
|
+
endpoint=slim_endpoint,
|
|
280
|
+
auth_secret=slim_secret,
|
|
281
|
+
message_discriminator="type",
|
|
282
|
+
)
|
|
283
|
+
app = PASlimApp(config)
|
|
284
|
+
if verify_sessions:
|
|
285
|
+
app.session_token_verifier = self.verify_session_token
|
|
286
|
+
logger.info(f"PatternSDK: agent configured (deployment={self._deployment_id} slim_name={slim_name})")
|
|
287
|
+
return app
|
|
288
|
+
|
|
289
|
+
def agent_app(self, *, verify_sessions: bool = True) -> PASlimApp:
|
|
290
|
+
return asyncio.run(self._build_agent_app(verify_sessions=verify_sessions))
|
|
291
|
+
|
|
292
|
+
async def config(self, key: Optional[str] = None):
|
|
293
|
+
self._require_enabled()
|
|
294
|
+
if self._config_cache is None:
|
|
295
|
+
self._config_cache = await self._request("GET", "/config")
|
|
296
|
+
if key is not None:
|
|
297
|
+
return self._config_cache.get(key)
|
|
298
|
+
return self._config_cache
|
|
299
|
+
|
|
300
|
+
async def prompt(self, slug: str) -> Optional[str]:
|
|
301
|
+
self._require_enabled()
|
|
302
|
+
if self._prompts_cache is None:
|
|
303
|
+
self._prompts_cache = await self._request("GET", "/prompts")
|
|
304
|
+
return self._prompts_cache.get(slug)
|
|
305
|
+
|
|
306
|
+
async def prompts(self) -> dict:
|
|
307
|
+
self._require_enabled()
|
|
308
|
+
if self._prompts_cache is None:
|
|
309
|
+
self._prompts_cache = await self._request("GET", "/prompts")
|
|
310
|
+
return self._prompts_cache
|
|
311
|
+
|
|
312
|
+
|
|
313
|
+
pa_sdk = PatternSDK(
|
|
314
|
+
os.environ.get("__PA_AGENT_SERVICES_ENDPOINT"),
|
|
315
|
+
os.environ.get("__PA_AGENT_SERVICES_TOKEN_REQUEST_HEADERS"),
|
|
316
|
+
)
|
|
@@ -0,0 +1,335 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import os
|
|
3
|
+
import stat
|
|
4
|
+
import time
|
|
5
|
+
import base64
|
|
6
|
+
from unittest.mock import AsyncMock, patch, MagicMock
|
|
7
|
+
|
|
8
|
+
import httpx
|
|
9
|
+
import jwt as pyjwt
|
|
10
|
+
import pytest
|
|
11
|
+
from cryptography.hazmat.primitives.asymmetric import rsa
|
|
12
|
+
from cryptography.hazmat.primitives import serialization
|
|
13
|
+
|
|
14
|
+
from pattern_agent_sdk import AgentSessionToken, PatternSDK, SessionTokenJWKSConfig
|
|
15
|
+
from pattern_agent_sdk.sdk import _decode_jwt_payload, _parse_extra_headers
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _make_jwt(payload: dict) -> str:
|
|
19
|
+
header = base64.urlsafe_b64encode(b'{"alg":"HS256"}').rstrip(b"=").decode()
|
|
20
|
+
body = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode()
|
|
21
|
+
return f"{header}.{body}.fakesig"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class TestDecodeJwtPayload:
|
|
25
|
+
def test_valid_token(self):
|
|
26
|
+
payload = {"sub": "dep-123", "exp": 9999999999}
|
|
27
|
+
token = _make_jwt(payload)
|
|
28
|
+
assert _decode_jwt_payload(token) == payload
|
|
29
|
+
|
|
30
|
+
def test_invalid_token(self):
|
|
31
|
+
with pytest.raises(ValueError, match="Invalid JWT format"):
|
|
32
|
+
_decode_jwt_payload("not.a.valid.jwt.token")
|
|
33
|
+
|
|
34
|
+
def test_two_parts(self):
|
|
35
|
+
with pytest.raises(ValueError, match="Invalid JWT format"):
|
|
36
|
+
_decode_jwt_payload("only.two")
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class TestParseExtraHeaders:
|
|
40
|
+
def test_valid_headers(self):
|
|
41
|
+
raw = base64.b64encode(json.dumps({"X-Forwarded-Client-Cert": "Hash=abc"}).encode()).decode()
|
|
42
|
+
assert _parse_extra_headers(raw) == {"X-Forwarded-Client-Cert": "Hash=abc"}
|
|
43
|
+
|
|
44
|
+
def test_empty_or_none(self):
|
|
45
|
+
assert _parse_extra_headers(None) == {}
|
|
46
|
+
assert _parse_extra_headers("") == {}
|
|
47
|
+
|
|
48
|
+
def test_malformed_base64(self):
|
|
49
|
+
assert _parse_extra_headers("not-valid-base64!!!") == {}
|
|
50
|
+
|
|
51
|
+
def test_non_dict_json(self):
|
|
52
|
+
raw = base64.b64encode(b'["a","b"]').decode()
|
|
53
|
+
assert _parse_extra_headers(raw) == {}
|
|
54
|
+
|
|
55
|
+
def test_non_string_values(self):
|
|
56
|
+
raw = base64.b64encode(json.dumps({"key": 123}).encode()).decode()
|
|
57
|
+
assert _parse_extra_headers(raw) == {}
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
class TestMintTokenHeaders:
|
|
61
|
+
async def test_mint_token_sends_extra_headers(self, respx_mock):
|
|
62
|
+
extra = {"X-Forwarded-Client-Cert": "Hash=abc123"}
|
|
63
|
+
raw = base64.b64encode(json.dumps(extra).encode()).decode()
|
|
64
|
+
sdk = PatternSDK("http://localhost:9999", token_request_headers=raw)
|
|
65
|
+
|
|
66
|
+
token = _make_jwt({"sub": "dep-1", "exp": 9999999999})
|
|
67
|
+
route = respx_mock.post("http://localhost:9999/mint-agent-token").mock(
|
|
68
|
+
return_value=httpx.Response(200, json={"token": token})
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
await sdk._mint_token()
|
|
72
|
+
|
|
73
|
+
assert route.called
|
|
74
|
+
req = route.calls[0].request
|
|
75
|
+
assert req.headers["X-Forwarded-Client-Cert"] == "Hash=abc123"
|
|
76
|
+
|
|
77
|
+
async def test_mint_token_no_extra_headers(self, respx_mock):
|
|
78
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
79
|
+
|
|
80
|
+
token = _make_jwt({"sub": "dep-1", "exp": 9999999999})
|
|
81
|
+
route = respx_mock.post("http://localhost:9999/mint-agent-token").mock(
|
|
82
|
+
return_value=httpx.Response(200, json={"token": token})
|
|
83
|
+
)
|
|
84
|
+
|
|
85
|
+
await sdk._mint_token()
|
|
86
|
+
|
|
87
|
+
assert route.called
|
|
88
|
+
req = route.calls[0].request
|
|
89
|
+
assert "X-Forwarded-Client-Cert" not in req.headers
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
class TestPatternSDKNoEndpoint:
|
|
93
|
+
def test_init_silent_when_no_endpoint(self):
|
|
94
|
+
sdk = PatternSDK()
|
|
95
|
+
assert not sdk._enabled
|
|
96
|
+
|
|
97
|
+
async def test_methods_raise_when_disabled(self):
|
|
98
|
+
sdk = PatternSDK()
|
|
99
|
+
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
100
|
+
await sdk.config()
|
|
101
|
+
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
102
|
+
await sdk.prompt("inference")
|
|
103
|
+
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
104
|
+
await sdk.prompts()
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
class TestConfigCaching:
|
|
108
|
+
async def test_config_cached(self):
|
|
109
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
110
|
+
mock_data = {"key1": "val1", "key2": "val2"}
|
|
111
|
+
|
|
112
|
+
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value=mock_data) as mock_req:
|
|
113
|
+
result1 = await sdk.config()
|
|
114
|
+
result2 = await sdk.config()
|
|
115
|
+
result3 = await sdk.config("key1")
|
|
116
|
+
|
|
117
|
+
assert result1 == mock_data
|
|
118
|
+
assert result2 == mock_data
|
|
119
|
+
assert result3 == "val1"
|
|
120
|
+
mock_req.assert_awaited_once_with("GET", "/config")
|
|
121
|
+
|
|
122
|
+
async def test_config_key_missing_returns_none(self):
|
|
123
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
124
|
+
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value={"a": 1}):
|
|
125
|
+
assert await sdk.config("missing") is None
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
class TestPromptsCaching:
|
|
129
|
+
async def test_prompt_shares_cache_with_prompts(self):
|
|
130
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
131
|
+
mock_data = {"greeting": "Hello!", "farewell": "Bye!"}
|
|
132
|
+
|
|
133
|
+
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value=mock_data) as mock_req:
|
|
134
|
+
single = await sdk.prompt("greeting")
|
|
135
|
+
all_prompts = await sdk.prompts()
|
|
136
|
+
|
|
137
|
+
assert single == "Hello!"
|
|
138
|
+
assert all_prompts == mock_data
|
|
139
|
+
mock_req.assert_awaited_once_with("GET", "/prompts")
|
|
140
|
+
|
|
141
|
+
async def test_prompt_missing_returns_none(self):
|
|
142
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
143
|
+
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value={}):
|
|
144
|
+
assert await sdk.prompt("nope") is None
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
class TestTokenPersistence:
|
|
148
|
+
async def test_enable_writes_token_to_disk(self, respx_mock, tmp_path):
|
|
149
|
+
token = _make_jwt({"sub": "dep-1", "exp": 9999999999})
|
|
150
|
+
respx_mock.post("http://localhost:9999/mint-agent-token").mock(
|
|
151
|
+
return_value=httpx.Response(200, json={"token": token})
|
|
152
|
+
)
|
|
153
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
154
|
+
path = str(tmp_path / "token")
|
|
155
|
+
|
|
156
|
+
with patch.dict(os.environ, {"__PA_AGENT_SDK_TOKEN_PERSIST_PATH": path}):
|
|
157
|
+
await sdk.enable_token_persistence()
|
|
158
|
+
|
|
159
|
+
assert sdk.token_file_path == path
|
|
160
|
+
assert open(path).read() == token
|
|
161
|
+
assert stat.S_IMODE(os.stat(path).st_mode) == 0o600
|
|
162
|
+
|
|
163
|
+
async def test_missing_env_var_raises(self):
|
|
164
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
165
|
+
with patch.dict(os.environ, {}, clear=True):
|
|
166
|
+
with pytest.raises(RuntimeError, match="__PA_AGENT_SDK_TOKEN_PERSIST_PATH is required"):
|
|
167
|
+
await sdk.enable_token_persistence()
|
|
168
|
+
|
|
169
|
+
def test_token_file_path_raises_when_not_enabled(self):
|
|
170
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
171
|
+
with pytest.raises(RuntimeError, match="Token persistence has not been enabled"):
|
|
172
|
+
_ = sdk.token_file_path
|
|
173
|
+
|
|
174
|
+
async def test_token_refresh_updates_file(self, respx_mock, tmp_path):
|
|
175
|
+
token1 = _make_jwt({"sub": "dep-1", "exp": 9999999999})
|
|
176
|
+
token2 = _make_jwt({"sub": "dep-1", "exp": 9999999999, "v": 2})
|
|
177
|
+
route = respx_mock.post("http://localhost:9999/mint-agent-token")
|
|
178
|
+
route.side_effect = [
|
|
179
|
+
httpx.Response(200, json={"token": token1}),
|
|
180
|
+
httpx.Response(200, json={"token": token2}),
|
|
181
|
+
]
|
|
182
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
183
|
+
path = str(tmp_path / "token")
|
|
184
|
+
|
|
185
|
+
with patch.dict(os.environ, {"__PA_AGENT_SDK_TOKEN_PERSIST_PATH": path}):
|
|
186
|
+
await sdk.enable_token_persistence()
|
|
187
|
+
|
|
188
|
+
assert open(path).read() == token1
|
|
189
|
+
|
|
190
|
+
sdk._token_exp = 0
|
|
191
|
+
await sdk._ensure_token()
|
|
192
|
+
|
|
193
|
+
assert open(path).read() == token2
|
|
194
|
+
|
|
195
|
+
async def test_disabled_sdk_raises(self):
|
|
196
|
+
sdk = PatternSDK()
|
|
197
|
+
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
198
|
+
await sdk.enable_token_persistence()
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
def _rsa_key_pair():
|
|
202
|
+
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
|
203
|
+
private_pem = private_key.private_bytes(
|
|
204
|
+
serialization.Encoding.PEM,
|
|
205
|
+
serialization.PrivateFormat.PKCS8,
|
|
206
|
+
serialization.NoEncryption(),
|
|
207
|
+
)
|
|
208
|
+
return private_key, private_pem
|
|
209
|
+
|
|
210
|
+
|
|
211
|
+
def _make_session_jwt(private_pem: bytes, payload: dict, headers: dict | None = None) -> str:
|
|
212
|
+
return pyjwt.encode(payload, private_pem, algorithm="RS256", headers=headers)
|
|
213
|
+
|
|
214
|
+
|
|
215
|
+
class TestSessionTokenConfig:
|
|
216
|
+
async def test_fetches_and_caches(self):
|
|
217
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
218
|
+
config_data = {"jwks_url": "https://example.com/.well-known/jwks.json", "audience": "my-app"}
|
|
219
|
+
|
|
220
|
+
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value=config_data) as mock_req:
|
|
221
|
+
result1 = await sdk.session_token_config()
|
|
222
|
+
result2 = await sdk.session_token_config()
|
|
223
|
+
|
|
224
|
+
assert result1 == SessionTokenJWKSConfig(jwks_url="https://example.com/.well-known/jwks.json", audience="my-app")
|
|
225
|
+
assert result2 is result1
|
|
226
|
+
mock_req.assert_awaited_once_with("GET", "/session_token_jwks_config")
|
|
227
|
+
assert sdk._jwks_client is not None
|
|
228
|
+
|
|
229
|
+
async def test_with_issuer(self):
|
|
230
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
231
|
+
config_data = {
|
|
232
|
+
"jwks_url": "https://example.com/.well-known/jwks.json",
|
|
233
|
+
"audience": "my-app",
|
|
234
|
+
"issuer": "https://auth.example.com",
|
|
235
|
+
}
|
|
236
|
+
|
|
237
|
+
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value=config_data):
|
|
238
|
+
result = await sdk.session_token_config()
|
|
239
|
+
|
|
240
|
+
assert result.issuer == "https://auth.example.com"
|
|
241
|
+
|
|
242
|
+
|
|
243
|
+
class TestVerifySessionToken:
|
|
244
|
+
def test_raises_without_config(self):
|
|
245
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
246
|
+
with pytest.raises(RuntimeError, match="Session token verification not configured"):
|
|
247
|
+
sdk.verify_session_token("some.jwt.token")
|
|
248
|
+
|
|
249
|
+
def test_valid_token(self):
|
|
250
|
+
private_key, private_pem = _rsa_key_pair()
|
|
251
|
+
now = int(time.time())
|
|
252
|
+
payload = {
|
|
253
|
+
"sub": "session-abc",
|
|
254
|
+
"tenant_id": "tenant-1",
|
|
255
|
+
"user_id": "user-42",
|
|
256
|
+
"agents": ["agent-a", "agent-b"],
|
|
257
|
+
"exp": now + 3600,
|
|
258
|
+
"iat": now,
|
|
259
|
+
"aud": "my-app",
|
|
260
|
+
}
|
|
261
|
+
token = _make_session_jwt(private_pem, payload)
|
|
262
|
+
|
|
263
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
264
|
+
sdk._session_jwks_config = SessionTokenJWKSConfig(
|
|
265
|
+
jwks_url="https://example.com/.well-known/jwks.json",
|
|
266
|
+
audience="my-app",
|
|
267
|
+
)
|
|
268
|
+
mock_signing_key = MagicMock()
|
|
269
|
+
mock_signing_key.key = private_key.public_key()
|
|
270
|
+
sdk._jwks_client = MagicMock()
|
|
271
|
+
sdk._jwks_client.get_signing_key_from_jwt.return_value = mock_signing_key
|
|
272
|
+
|
|
273
|
+
result = sdk.verify_session_token(token)
|
|
274
|
+
|
|
275
|
+
assert isinstance(result, AgentSessionToken)
|
|
276
|
+
assert result.session_id == "session-abc"
|
|
277
|
+
assert result.tenant_id == "tenant-1"
|
|
278
|
+
assert result.user_id == "user-42"
|
|
279
|
+
assert result.agents == ["agent-a", "agent-b"]
|
|
280
|
+
assert result.exp == now + 3600
|
|
281
|
+
assert result.iat == now
|
|
282
|
+
|
|
283
|
+
def test_expired_token_raises(self):
|
|
284
|
+
private_key, private_pem = _rsa_key_pair()
|
|
285
|
+
now = int(time.time())
|
|
286
|
+
payload = {
|
|
287
|
+
"sub": "session-abc",
|
|
288
|
+
"tenant_id": "tenant-1",
|
|
289
|
+
"user_id": "user-42",
|
|
290
|
+
"agents": [],
|
|
291
|
+
"exp": now - 100,
|
|
292
|
+
"iat": now - 3700,
|
|
293
|
+
"aud": "my-app",
|
|
294
|
+
}
|
|
295
|
+
token = _make_session_jwt(private_pem, payload)
|
|
296
|
+
|
|
297
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
298
|
+
sdk._session_jwks_config = SessionTokenJWKSConfig(
|
|
299
|
+
jwks_url="https://example.com/.well-known/jwks.json",
|
|
300
|
+
audience="my-app",
|
|
301
|
+
)
|
|
302
|
+
mock_signing_key = MagicMock()
|
|
303
|
+
mock_signing_key.key = private_key.public_key()
|
|
304
|
+
sdk._jwks_client = MagicMock()
|
|
305
|
+
sdk._jwks_client.get_signing_key_from_jwt.return_value = mock_signing_key
|
|
306
|
+
|
|
307
|
+
with pytest.raises(pyjwt.ExpiredSignatureError):
|
|
308
|
+
sdk.verify_session_token(token)
|
|
309
|
+
|
|
310
|
+
def test_wrong_audience_raises(self):
|
|
311
|
+
private_key, private_pem = _rsa_key_pair()
|
|
312
|
+
now = int(time.time())
|
|
313
|
+
payload = {
|
|
314
|
+
"sub": "session-abc",
|
|
315
|
+
"tenant_id": "tenant-1",
|
|
316
|
+
"user_id": "user-42",
|
|
317
|
+
"agents": [],
|
|
318
|
+
"exp": now + 3600,
|
|
319
|
+
"iat": now,
|
|
320
|
+
"aud": "wrong-audience",
|
|
321
|
+
}
|
|
322
|
+
token = _make_session_jwt(private_pem, payload)
|
|
323
|
+
|
|
324
|
+
sdk = PatternSDK("http://localhost:9999")
|
|
325
|
+
sdk._session_jwks_config = SessionTokenJWKSConfig(
|
|
326
|
+
jwks_url="https://example.com/.well-known/jwks.json",
|
|
327
|
+
audience="my-app",
|
|
328
|
+
)
|
|
329
|
+
mock_signing_key = MagicMock()
|
|
330
|
+
mock_signing_key.key = private_key.public_key()
|
|
331
|
+
sdk._jwks_client = MagicMock()
|
|
332
|
+
sdk._jwks_client.get_signing_key_from_jwt.return_value = mock_signing_key
|
|
333
|
+
|
|
334
|
+
with pytest.raises(pyjwt.InvalidAudienceError):
|
|
335
|
+
sdk.verify_session_token(token)
|
|
@@ -1,164 +0,0 @@
|
|
|
1
|
-
import os
|
|
2
|
-
import json
|
|
3
|
-
import time
|
|
4
|
-
import base64
|
|
5
|
-
import logging
|
|
6
|
-
from typing import Optional
|
|
7
|
-
|
|
8
|
-
import httpx
|
|
9
|
-
|
|
10
|
-
logger = logging.getLogger(__name__)
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
def _parse_extra_headers(raw: Optional[str]) -> dict[str, str]:
|
|
14
|
-
if not raw:
|
|
15
|
-
return {}
|
|
16
|
-
try:
|
|
17
|
-
decoded = base64.b64decode(raw)
|
|
18
|
-
headers = json.loads(decoded)
|
|
19
|
-
if isinstance(headers, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in headers.items()):
|
|
20
|
-
return headers
|
|
21
|
-
except Exception:
|
|
22
|
-
pass
|
|
23
|
-
logger.warning("PatternSDK: __PA_AGENT_SERVICES_TOKEN_REQUEST_HEADERS is malformed, ignoring")
|
|
24
|
-
return {}
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
def _decode_jwt_payload(token: str) -> dict:
|
|
28
|
-
parts = token.split(".")
|
|
29
|
-
if len(parts) != 3:
|
|
30
|
-
raise ValueError("Invalid JWT format")
|
|
31
|
-
payload_b64 = parts[1] + "=" * (-len(parts[1]) % 4)
|
|
32
|
-
return json.loads(base64.urlsafe_b64decode(payload_b64))
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
class PatternSDK:
|
|
36
|
-
def __init__(self, services_endpoint: Optional[str] = None, token_request_headers: Optional[str] = None):
|
|
37
|
-
self._services_endpoint = (services_endpoint or "").rstrip("/")
|
|
38
|
-
self._enabled = bool(self._services_endpoint)
|
|
39
|
-
if not self._enabled:
|
|
40
|
-
logger.info("PatternSDK: no services endpoint configured, SDK disabled")
|
|
41
|
-
return
|
|
42
|
-
self._mint_headers = _parse_extra_headers(token_request_headers)
|
|
43
|
-
self._token: Optional[str] = None
|
|
44
|
-
self._deployment_id: Optional[str] = None
|
|
45
|
-
self._token_exp: float = 0
|
|
46
|
-
self._renew_failed: bool = False
|
|
47
|
-
self._config_cache: Optional[dict] = None
|
|
48
|
-
self._prompts_cache: Optional[dict] = None
|
|
49
|
-
|
|
50
|
-
def _require_enabled(self):
|
|
51
|
-
if not self._enabled:
|
|
52
|
-
raise RuntimeError("PatternSDK is disabled: services endpoint not specified")
|
|
53
|
-
|
|
54
|
-
@property
|
|
55
|
-
def _endpoint(self) -> str:
|
|
56
|
-
return self._services_endpoint
|
|
57
|
-
|
|
58
|
-
async def _mint_token(self) -> str:
|
|
59
|
-
url = f"{self._endpoint}/mint-agent-token"
|
|
60
|
-
try:
|
|
61
|
-
async with httpx.AsyncClient(timeout=10) as client:
|
|
62
|
-
resp = await client.post(url, headers=self._mint_headers or None)
|
|
63
|
-
except httpx.ConnectError as e:
|
|
64
|
-
logger.error(f"PatternSDK: cannot connect to agent services at {url} — is the sidecar running? {e}")
|
|
65
|
-
raise
|
|
66
|
-
except httpx.TimeoutException:
|
|
67
|
-
logger.error(f"PatternSDK: request to {url} timed out")
|
|
68
|
-
raise
|
|
69
|
-
|
|
70
|
-
if resp.status_code != 200:
|
|
71
|
-
logger.error(f"PatternSDK: mint-agent-token returned {resp.status_code}: {resp.text}")
|
|
72
|
-
resp.raise_for_status()
|
|
73
|
-
|
|
74
|
-
try:
|
|
75
|
-
token = resp.json()["token"]
|
|
76
|
-
except (KeyError, json.JSONDecodeError) as e:
|
|
77
|
-
logger.error(f"PatternSDK: mint-agent-token response missing 'token' field: {resp.text}")
|
|
78
|
-
raise RuntimeError("mint-agent-token response missing 'token' field") from e
|
|
79
|
-
|
|
80
|
-
payload = _decode_jwt_payload(token)
|
|
81
|
-
self._token = token
|
|
82
|
-
self._deployment_id = payload["sub"]
|
|
83
|
-
self._token_exp = payload.get("exp", 0)
|
|
84
|
-
self._renew_failed = False
|
|
85
|
-
return token
|
|
86
|
-
|
|
87
|
-
async def _ensure_token(self):
|
|
88
|
-
if not self._token or time.time() >= self._token_exp:
|
|
89
|
-
await self._mint_token()
|
|
90
|
-
|
|
91
|
-
def _token_is_expired(self) -> bool:
|
|
92
|
-
return time.time() >= self._token_exp
|
|
93
|
-
|
|
94
|
-
async def _request(self, method: str, path: str) -> dict:
|
|
95
|
-
await self._ensure_token()
|
|
96
|
-
url = f"{self._endpoint}/deployment/{self._deployment_id}{path}"
|
|
97
|
-
headers = {"Authorization": f"Bearer {self._token}"}
|
|
98
|
-
|
|
99
|
-
try:
|
|
100
|
-
async with httpx.AsyncClient(timeout=10) as client:
|
|
101
|
-
resp = await client.request(method, url, headers=headers)
|
|
102
|
-
except httpx.ConnectError as e:
|
|
103
|
-
logger.error(f"PatternSDK: cannot connect to {url}: {e}")
|
|
104
|
-
raise
|
|
105
|
-
except httpx.TimeoutException:
|
|
106
|
-
logger.error(f"PatternSDK: request to {url} timed out")
|
|
107
|
-
raise
|
|
108
|
-
|
|
109
|
-
if resp.status_code == 403 and self._token_is_expired() and not self._renew_failed:
|
|
110
|
-
logger.info("PatternSDK: got 403 with expired token, attempting renewal")
|
|
111
|
-
try:
|
|
112
|
-
await self._mint_token()
|
|
113
|
-
except Exception:
|
|
114
|
-
self._renew_failed = True
|
|
115
|
-
logger.error("PatternSDK: token renewal failed after 403, not retrying")
|
|
116
|
-
raise
|
|
117
|
-
headers = {"Authorization": f"Bearer {self._token}"}
|
|
118
|
-
try:
|
|
119
|
-
async with httpx.AsyncClient(timeout=10) as client:
|
|
120
|
-
resp = await client.request(method, url, headers=headers)
|
|
121
|
-
except httpx.ConnectError as e:
|
|
122
|
-
logger.error(f"PatternSDK: cannot connect to {url} on retry: {e}")
|
|
123
|
-
raise
|
|
124
|
-
except httpx.TimeoutException:
|
|
125
|
-
logger.error(f"PatternSDK: retry request to {url} timed out")
|
|
126
|
-
raise
|
|
127
|
-
|
|
128
|
-
if resp.status_code == 401:
|
|
129
|
-
logger.error(f"PatternSDK: 401 from {path}: {resp.text}")
|
|
130
|
-
raise PermissionError(f"Unauthorized: {resp.text}")
|
|
131
|
-
if resp.status_code == 403:
|
|
132
|
-
logger.error(f"PatternSDK: 403 from {path}: {resp.text}")
|
|
133
|
-
raise PermissionError(f"Forbidden: {resp.text}")
|
|
134
|
-
if resp.status_code >= 400:
|
|
135
|
-
logger.error(f"PatternSDK: {resp.status_code} from {path}: {resp.text}")
|
|
136
|
-
resp.raise_for_status()
|
|
137
|
-
|
|
138
|
-
return resp.json()
|
|
139
|
-
|
|
140
|
-
async def config(self, key: Optional[str] = None):
|
|
141
|
-
self._require_enabled()
|
|
142
|
-
if self._config_cache is None:
|
|
143
|
-
self._config_cache = await self._request("GET", "/config")
|
|
144
|
-
if key is not None:
|
|
145
|
-
return self._config_cache.get(key)
|
|
146
|
-
return self._config_cache
|
|
147
|
-
|
|
148
|
-
async def prompt(self, slug: str) -> Optional[str]:
|
|
149
|
-
self._require_enabled()
|
|
150
|
-
if self._prompts_cache is None:
|
|
151
|
-
self._prompts_cache = await self._request("GET", "/prompts")
|
|
152
|
-
return self._prompts_cache.get(slug)
|
|
153
|
-
|
|
154
|
-
async def prompts(self) -> dict:
|
|
155
|
-
self._require_enabled()
|
|
156
|
-
if self._prompts_cache is None:
|
|
157
|
-
self._prompts_cache = await self._request("GET", "/prompts")
|
|
158
|
-
return self._prompts_cache
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
pa_sdk = PatternSDK(
|
|
162
|
-
os.environ.get("__PA_AGENT_SERVICES_ENDPOINT"),
|
|
163
|
-
os.environ.get("__PA_AGENT_SERVICES_TOKEN_REQUEST_HEADERS"),
|
|
164
|
-
)
|
|
@@ -1,138 +0,0 @@
|
|
|
1
|
-
import json
|
|
2
|
-
import base64
|
|
3
|
-
from unittest.mock import AsyncMock, patch
|
|
4
|
-
|
|
5
|
-
import httpx
|
|
6
|
-
import pytest
|
|
7
|
-
|
|
8
|
-
from pattern_agent_sdk import PatternSDK
|
|
9
|
-
from pattern_agent_sdk.sdk import _decode_jwt_payload, _parse_extra_headers
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
def _make_jwt(payload: dict) -> str:
|
|
13
|
-
header = base64.urlsafe_b64encode(b'{"alg":"HS256"}').rstrip(b"=").decode()
|
|
14
|
-
body = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode()
|
|
15
|
-
return f"{header}.{body}.fakesig"
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
class TestDecodeJwtPayload:
|
|
19
|
-
def test_valid_token(self):
|
|
20
|
-
payload = {"sub": "dep-123", "exp": 9999999999}
|
|
21
|
-
token = _make_jwt(payload)
|
|
22
|
-
assert _decode_jwt_payload(token) == payload
|
|
23
|
-
|
|
24
|
-
def test_invalid_token(self):
|
|
25
|
-
with pytest.raises(ValueError, match="Invalid JWT format"):
|
|
26
|
-
_decode_jwt_payload("not.a.valid.jwt.token")
|
|
27
|
-
|
|
28
|
-
def test_two_parts(self):
|
|
29
|
-
with pytest.raises(ValueError, match="Invalid JWT format"):
|
|
30
|
-
_decode_jwt_payload("only.two")
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
class TestParseExtraHeaders:
|
|
34
|
-
def test_valid_headers(self):
|
|
35
|
-
raw = base64.b64encode(json.dumps({"X-Forwarded-Client-Cert": "Hash=abc"}).encode()).decode()
|
|
36
|
-
assert _parse_extra_headers(raw) == {"X-Forwarded-Client-Cert": "Hash=abc"}
|
|
37
|
-
|
|
38
|
-
def test_empty_or_none(self):
|
|
39
|
-
assert _parse_extra_headers(None) == {}
|
|
40
|
-
assert _parse_extra_headers("") == {}
|
|
41
|
-
|
|
42
|
-
def test_malformed_base64(self):
|
|
43
|
-
assert _parse_extra_headers("not-valid-base64!!!") == {}
|
|
44
|
-
|
|
45
|
-
def test_non_dict_json(self):
|
|
46
|
-
raw = base64.b64encode(b'["a","b"]').decode()
|
|
47
|
-
assert _parse_extra_headers(raw) == {}
|
|
48
|
-
|
|
49
|
-
def test_non_string_values(self):
|
|
50
|
-
raw = base64.b64encode(json.dumps({"key": 123}).encode()).decode()
|
|
51
|
-
assert _parse_extra_headers(raw) == {}
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
class TestMintTokenHeaders:
|
|
55
|
-
async def test_mint_token_sends_extra_headers(self, respx_mock):
|
|
56
|
-
extra = {"X-Forwarded-Client-Cert": "Hash=abc123"}
|
|
57
|
-
raw = base64.b64encode(json.dumps(extra).encode()).decode()
|
|
58
|
-
sdk = PatternSDK("http://localhost:9999", token_request_headers=raw)
|
|
59
|
-
|
|
60
|
-
token = _make_jwt({"sub": "dep-1", "exp": 9999999999})
|
|
61
|
-
route = respx_mock.post("http://localhost:9999/mint-agent-token").mock(
|
|
62
|
-
return_value=httpx.Response(200, json={"token": token})
|
|
63
|
-
)
|
|
64
|
-
|
|
65
|
-
await sdk._mint_token()
|
|
66
|
-
|
|
67
|
-
assert route.called
|
|
68
|
-
req = route.calls[0].request
|
|
69
|
-
assert req.headers["X-Forwarded-Client-Cert"] == "Hash=abc123"
|
|
70
|
-
|
|
71
|
-
async def test_mint_token_no_extra_headers(self, respx_mock):
|
|
72
|
-
sdk = PatternSDK("http://localhost:9999")
|
|
73
|
-
|
|
74
|
-
token = _make_jwt({"sub": "dep-1", "exp": 9999999999})
|
|
75
|
-
route = respx_mock.post("http://localhost:9999/mint-agent-token").mock(
|
|
76
|
-
return_value=httpx.Response(200, json={"token": token})
|
|
77
|
-
)
|
|
78
|
-
|
|
79
|
-
await sdk._mint_token()
|
|
80
|
-
|
|
81
|
-
assert route.called
|
|
82
|
-
req = route.calls[0].request
|
|
83
|
-
assert "X-Forwarded-Client-Cert" not in req.headers
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
class TestPatternSDKNoEndpoint:
|
|
87
|
-
def test_init_silent_when_no_endpoint(self):
|
|
88
|
-
sdk = PatternSDK()
|
|
89
|
-
assert not sdk._enabled
|
|
90
|
-
|
|
91
|
-
async def test_methods_raise_when_disabled(self):
|
|
92
|
-
sdk = PatternSDK()
|
|
93
|
-
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
94
|
-
await sdk.config()
|
|
95
|
-
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
96
|
-
await sdk.prompt("inference")
|
|
97
|
-
with pytest.raises(RuntimeError, match="SDK is disabled"):
|
|
98
|
-
await sdk.prompts()
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
class TestConfigCaching:
|
|
102
|
-
async def test_config_cached(self):
|
|
103
|
-
sdk = PatternSDK("http://localhost:9999")
|
|
104
|
-
mock_data = {"key1": "val1", "key2": "val2"}
|
|
105
|
-
|
|
106
|
-
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value=mock_data) as mock_req:
|
|
107
|
-
result1 = await sdk.config()
|
|
108
|
-
result2 = await sdk.config()
|
|
109
|
-
result3 = await sdk.config("key1")
|
|
110
|
-
|
|
111
|
-
assert result1 == mock_data
|
|
112
|
-
assert result2 == mock_data
|
|
113
|
-
assert result3 == "val1"
|
|
114
|
-
mock_req.assert_awaited_once_with("GET", "/config")
|
|
115
|
-
|
|
116
|
-
async def test_config_key_missing_returns_none(self):
|
|
117
|
-
sdk = PatternSDK("http://localhost:9999")
|
|
118
|
-
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value={"a": 1}):
|
|
119
|
-
assert await sdk.config("missing") is None
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
class TestPromptsCaching:
|
|
123
|
-
async def test_prompt_shares_cache_with_prompts(self):
|
|
124
|
-
sdk = PatternSDK("http://localhost:9999")
|
|
125
|
-
mock_data = {"greeting": "Hello!", "farewell": "Bye!"}
|
|
126
|
-
|
|
127
|
-
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value=mock_data) as mock_req:
|
|
128
|
-
single = await sdk.prompt("greeting")
|
|
129
|
-
all_prompts = await sdk.prompts()
|
|
130
|
-
|
|
131
|
-
assert single == "Hello!"
|
|
132
|
-
assert all_prompts == mock_data
|
|
133
|
-
mock_req.assert_awaited_once_with("GET", "/prompts")
|
|
134
|
-
|
|
135
|
-
async def test_prompt_missing_returns_none(self):
|
|
136
|
-
sdk = PatternSDK("http://localhost:9999")
|
|
137
|
-
with patch.object(sdk, "_request", new_callable=AsyncMock, return_value={}):
|
|
138
|
-
assert await sdk.prompt("nope") is None
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|