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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pattern_agent_sdk
3
- Version: 0.1.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.1.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,4 @@
1
+ from .sdk import PatternSDK, SessionTokenJWKSConfig, pa_sdk
2
+ from pattern_agentic_messaging.session_token import PatternAgentSessionToken
3
+
4
+ __all__ = ["PatternAgentSessionToken", "PatternSDK", "SessionTokenJWKSConfig", "pa_sdk"]
@@ -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,3 +0,0 @@
1
- from .sdk import PatternSDK, pa_sdk
2
-
3
- __all__ = ["PatternSDK", "pa_sdk"]
@@ -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