whiteroom-bedrock 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.
- whiteroom_bedrock-0.1.0/PKG-INFO +100 -0
- whiteroom_bedrock-0.1.0/README.md +88 -0
- whiteroom_bedrock-0.1.0/pyproject.toml +23 -0
- whiteroom_bedrock-0.1.0/setup.cfg +4 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock/__init__.py +10 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock/attest.py +146 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock/hook.py +92 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock.egg-info/PKG-INFO +100 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock.egg-info/SOURCES.txt +12 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock.egg-info/dependency_links.txt +1 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock.egg-info/requires.txt +5 -0
- whiteroom_bedrock-0.1.0/src/whiteroom_bedrock.egg-info/top_level.txt +1 -0
- whiteroom_bedrock-0.1.0/tests/test_attest.py +190 -0
- whiteroom_bedrock-0.1.0/tests/test_hook.py +139 -0
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: whiteroom-bedrock
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Attaches WhiteRoom session attestation to a boto3 bedrock-runtime client
|
|
5
|
+
License: BSL-1.1
|
|
6
|
+
Requires-Python: >=3.9
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
Requires-Dist: botocore>=1.34
|
|
9
|
+
Provides-Extra: test
|
|
10
|
+
Requires-Dist: pytest>=7; extra == "test"
|
|
11
|
+
Requires-Dist: boto3>=1.34; extra == "test"
|
|
12
|
+
|
|
13
|
+
# whiteroom-bedrock
|
|
14
|
+
|
|
15
|
+
Client-side half of WhiteRoom's Bedrock adapter (server side lives in
|
|
16
|
+
`apps/engine/src/bedrock/`).
|
|
17
|
+
|
|
18
|
+
An unmodified AWS SDK has no way to prove its identity to WhiteRoom the way
|
|
19
|
+
`ANTHROPIC_BASE_URL`/`OPENAI_BASE_URL` callers do with a bearer key — every
|
|
20
|
+
Bedrock call is SigV4-signed by the caller's own credentials, and WhiteRoom
|
|
21
|
+
discards that signature (it re-signs outbound with its own assumed-role
|
|
22
|
+
credentials). Proving who called instead happens once, out of band, by
|
|
23
|
+
replaying a signed `sts:GetCallerIdentity` probe to WhiteRoom's `/attest`
|
|
24
|
+
endpoint. This package does that handshake and attaches the resulting
|
|
25
|
+
session token to every subsequent call, so the rest of your agent code never
|
|
26
|
+
has to know WhiteRoom is in the loop.
|
|
27
|
+
|
|
28
|
+
## Install
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
pip install -e sdks/bedrock-python
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
## Test
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install -e "sdks/bedrock-python[test]"
|
|
38
|
+
pytest sdks/bedrock-python/tests
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
No network access or AWS credentials needed — everything's mocked at the
|
|
42
|
+
botocore/urllib layer. `tests/test_attest.py` is the one worth reading
|
|
43
|
+
closely: it directly checks the signed probe has `host`,
|
|
44
|
+
`x-whiteroom-server-id` and `x-amz-date` in `SignedHeaders`, which is
|
|
45
|
+
exactly what the server (`attest.ts`) requires — a change here that broke
|
|
46
|
+
that would otherwise only surface as every real attestation failing
|
|
47
|
+
against a live proxy.
|
|
48
|
+
|
|
49
|
+
## Use
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
import boto3
|
|
53
|
+
import whiteroom_bedrock
|
|
54
|
+
|
|
55
|
+
client = boto3.client(
|
|
56
|
+
"bedrock-runtime",
|
|
57
|
+
region_name="us-east-1",
|
|
58
|
+
endpoint_url="https://proxy.whiteroom.tech",
|
|
59
|
+
)
|
|
60
|
+
whiteroom_bedrock.install(
|
|
61
|
+
client,
|
|
62
|
+
proxy_base_url="https://proxy.whiteroom.tech",
|
|
63
|
+
server_id="proxy.whiteroom.tech", # must match the proxy's configured serverId
|
|
64
|
+
region="us-east-1",
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
# Everything from here down is a normal bedrock-runtime client.
|
|
68
|
+
response = client.converse(
|
|
69
|
+
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
|
|
70
|
+
messages=[{"role": "user", "content": [{"text": "hi"}]}],
|
|
71
|
+
)
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## What `install()` actually does
|
|
75
|
+
|
|
76
|
+
- On first call (and again whenever the cached session goes stale), signs a
|
|
77
|
+
throwaway `GetCallerIdentity` request with whatever credentials this
|
|
78
|
+
process already resolves for AWS (env vars, shared config, an instance or
|
|
79
|
+
task role — the standard botocore chain) and POSTs it to `/attest`.
|
|
80
|
+
- Attaches the session token WhiteRoom hands back as `x-whiteroom-session`
|
|
81
|
+
on every `invoke_model`/`converse` call, via a `before-send` hook — after
|
|
82
|
+
signing, so it never interferes with the client's own (ultimately
|
|
83
|
+
discarded) signature.
|
|
84
|
+
- If the proxy ever answers 403 — the session expired, or WhiteRoom
|
|
85
|
+
restarted and lost its in-memory session table — invalidates the cached
|
|
86
|
+
token and forces exactly one retry, so callers don't see a transient
|
|
87
|
+
session death as a failure.
|
|
88
|
+
|
|
89
|
+
## Not handled here
|
|
90
|
+
|
|
91
|
+
- **Streaming.** `invoke_model_with_response_stream` and
|
|
92
|
+
`converse_stream` reach the proxy the same as any other call, but the
|
|
93
|
+
server actively refuses to proxy streaming responses today (see
|
|
94
|
+
`routes.ts`) — you'll get a `ValidationException` back, not a stream.
|
|
95
|
+
- **Credential source mismatch.** `region` and credentials are resolved
|
|
96
|
+
from the default botocore chain unless you pass `region=`/
|
|
97
|
+
`botocore_session=` explicitly. If the client you're hooking was built
|
|
98
|
+
with credentials that *don't* come from that chain (e.g. explicit
|
|
99
|
+
`aws_access_key_id=` on `boto3.client(...)`), attest with the same
|
|
100
|
+
credentials by passing a `botocore.session.Session` that resolves them.
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
# whiteroom-bedrock
|
|
2
|
+
|
|
3
|
+
Client-side half of WhiteRoom's Bedrock adapter (server side lives in
|
|
4
|
+
`apps/engine/src/bedrock/`).
|
|
5
|
+
|
|
6
|
+
An unmodified AWS SDK has no way to prove its identity to WhiteRoom the way
|
|
7
|
+
`ANTHROPIC_BASE_URL`/`OPENAI_BASE_URL` callers do with a bearer key — every
|
|
8
|
+
Bedrock call is SigV4-signed by the caller's own credentials, and WhiteRoom
|
|
9
|
+
discards that signature (it re-signs outbound with its own assumed-role
|
|
10
|
+
credentials). Proving who called instead happens once, out of band, by
|
|
11
|
+
replaying a signed `sts:GetCallerIdentity` probe to WhiteRoom's `/attest`
|
|
12
|
+
endpoint. This package does that handshake and attaches the resulting
|
|
13
|
+
session token to every subsequent call, so the rest of your agent code never
|
|
14
|
+
has to know WhiteRoom is in the loop.
|
|
15
|
+
|
|
16
|
+
## Install
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
pip install -e sdks/bedrock-python
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
## Test
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install -e "sdks/bedrock-python[test]"
|
|
26
|
+
pytest sdks/bedrock-python/tests
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
No network access or AWS credentials needed — everything's mocked at the
|
|
30
|
+
botocore/urllib layer. `tests/test_attest.py` is the one worth reading
|
|
31
|
+
closely: it directly checks the signed probe has `host`,
|
|
32
|
+
`x-whiteroom-server-id` and `x-amz-date` in `SignedHeaders`, which is
|
|
33
|
+
exactly what the server (`attest.ts`) requires — a change here that broke
|
|
34
|
+
that would otherwise only surface as every real attestation failing
|
|
35
|
+
against a live proxy.
|
|
36
|
+
|
|
37
|
+
## Use
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
import boto3
|
|
41
|
+
import whiteroom_bedrock
|
|
42
|
+
|
|
43
|
+
client = boto3.client(
|
|
44
|
+
"bedrock-runtime",
|
|
45
|
+
region_name="us-east-1",
|
|
46
|
+
endpoint_url="https://proxy.whiteroom.tech",
|
|
47
|
+
)
|
|
48
|
+
whiteroom_bedrock.install(
|
|
49
|
+
client,
|
|
50
|
+
proxy_base_url="https://proxy.whiteroom.tech",
|
|
51
|
+
server_id="proxy.whiteroom.tech", # must match the proxy's configured serverId
|
|
52
|
+
region="us-east-1",
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
# Everything from here down is a normal bedrock-runtime client.
|
|
56
|
+
response = client.converse(
|
|
57
|
+
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
|
|
58
|
+
messages=[{"role": "user", "content": [{"text": "hi"}]}],
|
|
59
|
+
)
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
## What `install()` actually does
|
|
63
|
+
|
|
64
|
+
- On first call (and again whenever the cached session goes stale), signs a
|
|
65
|
+
throwaway `GetCallerIdentity` request with whatever credentials this
|
|
66
|
+
process already resolves for AWS (env vars, shared config, an instance or
|
|
67
|
+
task role — the standard botocore chain) and POSTs it to `/attest`.
|
|
68
|
+
- Attaches the session token WhiteRoom hands back as `x-whiteroom-session`
|
|
69
|
+
on every `invoke_model`/`converse` call, via a `before-send` hook — after
|
|
70
|
+
signing, so it never interferes with the client's own (ultimately
|
|
71
|
+
discarded) signature.
|
|
72
|
+
- If the proxy ever answers 403 — the session expired, or WhiteRoom
|
|
73
|
+
restarted and lost its in-memory session table — invalidates the cached
|
|
74
|
+
token and forces exactly one retry, so callers don't see a transient
|
|
75
|
+
session death as a failure.
|
|
76
|
+
|
|
77
|
+
## Not handled here
|
|
78
|
+
|
|
79
|
+
- **Streaming.** `invoke_model_with_response_stream` and
|
|
80
|
+
`converse_stream` reach the proxy the same as any other call, but the
|
|
81
|
+
server actively refuses to proxy streaming responses today (see
|
|
82
|
+
`routes.ts`) — you'll get a `ValidationException` back, not a stream.
|
|
83
|
+
- **Credential source mismatch.** `region` and credentials are resolved
|
|
84
|
+
from the default botocore chain unless you pass `region=`/
|
|
85
|
+
`botocore_session=` explicitly. If the client you're hooking was built
|
|
86
|
+
with credentials that *don't* come from that chain (e.g. explicit
|
|
87
|
+
`aws_access_key_id=` on `boto3.client(...)`), attest with the same
|
|
88
|
+
credentials by passing a `botocore.session.Session` that resolves them.
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "whiteroom-bedrock"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Attaches WhiteRoom session attestation to a boto3 bedrock-runtime client"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "BSL-1.1" }
|
|
12
|
+
dependencies = [
|
|
13
|
+
"botocore>=1.34",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[project.optional-dependencies]
|
|
17
|
+
test = [
|
|
18
|
+
"pytest>=7",
|
|
19
|
+
"boto3>=1.34",
|
|
20
|
+
]
|
|
21
|
+
|
|
22
|
+
[tool.setuptools.packages.find]
|
|
23
|
+
where = ["src"]
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Client side of the attestation handshake described in
|
|
3
|
+
apps/engine/src/bedrock/attest.ts.
|
|
4
|
+
|
|
5
|
+
WhiteRoom cannot verify a request's SigV4 signature itself (that would need
|
|
6
|
+
the caller's secret key). Instead we sign a throwaway sts:GetCallerIdentity
|
|
7
|
+
request with whatever AWS credentials this process already has, and send it
|
|
8
|
+
to WhiteRoom's /attest endpoint. WhiteRoom replays it to STS, which can
|
|
9
|
+
verify it, and answers with our ARN plus a short-lived session token. That
|
|
10
|
+
token — not the probe — is what rides on every subsequent Bedrock call.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import json
|
|
16
|
+
import threading
|
|
17
|
+
import time
|
|
18
|
+
import urllib.error
|
|
19
|
+
import urllib.request
|
|
20
|
+
from dataclasses import dataclass
|
|
21
|
+
|
|
22
|
+
from botocore.auth import SigV4Auth
|
|
23
|
+
from botocore.awsrequest import AWSRequest
|
|
24
|
+
from botocore.session import Session as BotocoreSession
|
|
25
|
+
|
|
26
|
+
# The exact, minimal GetCallerIdentity body attest.ts's assertGetCallerIdentity
|
|
27
|
+
# expects: Action=GetCallerIdentity and nothing else but Version.
|
|
28
|
+
_GET_CALLER_IDENTITY_BODY = "Action=GetCallerIdentity&Version=2011-06-15"
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class AttestationError(Exception):
|
|
32
|
+
"""Raised when the probe can't be built, or WhiteRoom/STS rejects it."""
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@dataclass(frozen=True)
|
|
36
|
+
class AttestedIdentity:
|
|
37
|
+
session_token: str
|
|
38
|
+
arn: str
|
|
39
|
+
issued_at: float
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class SessionManager:
|
|
43
|
+
"""
|
|
44
|
+
Holds one WhiteRoom session token per (proxy, server_id, region) and
|
|
45
|
+
refreshes it as needed.
|
|
46
|
+
|
|
47
|
+
/attest's response carries no expiry (see routes.ts — deliberately: the
|
|
48
|
+
server owns TTL policy, not the client), so freshness here is a mix of:
|
|
49
|
+
- proactive: refresh after SESSION_TTL_HINT locally, so a long-idle
|
|
50
|
+
process doesn't eat the round trip lazily on its very next call
|
|
51
|
+
- reactive: hook.py invalidates and forces one retry on a 403 from
|
|
52
|
+
the proxy, which covers a token dying earlier than the hint (server
|
|
53
|
+
restart, explicit revocation) without the caller doing anything
|
|
54
|
+
|
|
55
|
+
Thread-safe: a single lock guards refresh, so concurrent callers on a
|
|
56
|
+
cold cache queue behind one /attest call instead of each firing their
|
|
57
|
+
own (the same thundering-herd concern credentials.ts guards against
|
|
58
|
+
server-side for AssumeRole).
|
|
59
|
+
"""
|
|
60
|
+
|
|
61
|
+
SESSION_TTL_HINT_SECONDS = 10 * 60
|
|
62
|
+
|
|
63
|
+
def __init__(
|
|
64
|
+
self,
|
|
65
|
+
proxy_base_url: str,
|
|
66
|
+
server_id: str,
|
|
67
|
+
region: str,
|
|
68
|
+
botocore_session: BotocoreSession | None = None,
|
|
69
|
+
timeout_seconds: float = 10.0,
|
|
70
|
+
) -> None:
|
|
71
|
+
self._proxy_base_url = proxy_base_url.rstrip("/")
|
|
72
|
+
self._server_id = server_id
|
|
73
|
+
self._region = region
|
|
74
|
+
self._botocore_session = botocore_session or BotocoreSession()
|
|
75
|
+
self._timeout_seconds = timeout_seconds
|
|
76
|
+
self._lock = threading.Lock()
|
|
77
|
+
self._identity: AttestedIdentity | None = None
|
|
78
|
+
|
|
79
|
+
def get_token(self, force: bool = False) -> str:
|
|
80
|
+
with self._lock:
|
|
81
|
+
if not force and self._identity is not None:
|
|
82
|
+
age = time.time() - self._identity.issued_at
|
|
83
|
+
if age < self.SESSION_TTL_HINT_SECONDS:
|
|
84
|
+
return self._identity.session_token
|
|
85
|
+
self._identity = self._attest()
|
|
86
|
+
return self._identity.session_token
|
|
87
|
+
|
|
88
|
+
def invalidate(self) -> None:
|
|
89
|
+
with self._lock:
|
|
90
|
+
self._identity = None
|
|
91
|
+
|
|
92
|
+
def _build_probe(self) -> dict:
|
|
93
|
+
creds = self._botocore_session.get_credentials()
|
|
94
|
+
if creds is None:
|
|
95
|
+
raise AttestationError(
|
|
96
|
+
"no AWS credentials available to attest with — "
|
|
97
|
+
"whiteroom_bedrock needs the same credential chain your "
|
|
98
|
+
"bedrock-runtime client uses"
|
|
99
|
+
)
|
|
100
|
+
frozen = creds.get_frozen_credentials()
|
|
101
|
+
|
|
102
|
+
host = f"sts.{self._region}.amazonaws.com"
|
|
103
|
+
request = AWSRequest(
|
|
104
|
+
method="POST",
|
|
105
|
+
url=f"https://{host}/",
|
|
106
|
+
data=_GET_CALLER_IDENTITY_BODY,
|
|
107
|
+
headers={
|
|
108
|
+
"Host": host,
|
|
109
|
+
"Content-Type": "application/x-www-form-urlencoded",
|
|
110
|
+
# Must be present BEFORE signing — SigV4Auth signs whatever
|
|
111
|
+
# headers exist on the request at add_auth() time, and
|
|
112
|
+
# attest.ts refuses any probe where this (and host) aren't
|
|
113
|
+
# in SignedHeaders.
|
|
114
|
+
"x-whiteroom-server-id": self._server_id,
|
|
115
|
+
},
|
|
116
|
+
)
|
|
117
|
+
SigV4Auth(frozen, "sts", self._region).add_auth(request)
|
|
118
|
+
|
|
119
|
+
return {
|
|
120
|
+
"method": "POST",
|
|
121
|
+
"headers": dict(request.headers.items()),
|
|
122
|
+
"body": _GET_CALLER_IDENTITY_BODY,
|
|
123
|
+
}
|
|
124
|
+
|
|
125
|
+
def _attest(self) -> AttestedIdentity:
|
|
126
|
+
probe = self._build_probe()
|
|
127
|
+
req = urllib.request.Request(
|
|
128
|
+
f"{self._proxy_base_url}/attest",
|
|
129
|
+
data=json.dumps(probe).encode("utf-8"),
|
|
130
|
+
headers={"Content-Type": "application/json"},
|
|
131
|
+
method="POST",
|
|
132
|
+
)
|
|
133
|
+
try:
|
|
134
|
+
with urllib.request.urlopen(req, timeout=self._timeout_seconds) as resp:
|
|
135
|
+
payload = json.loads(resp.read())
|
|
136
|
+
except urllib.error.HTTPError as e:
|
|
137
|
+
detail = e.read().decode("utf-8", "replace")
|
|
138
|
+
raise AttestationError(f"WhiteRoom rejected attestation ({e.code}): {detail}") from e
|
|
139
|
+
except urllib.error.URLError as e:
|
|
140
|
+
raise AttestationError(f"could not reach WhiteRoom attest endpoint: {e}") from e
|
|
141
|
+
|
|
142
|
+
token = payload.get("sessionToken")
|
|
143
|
+
arn = payload.get("arn", "")
|
|
144
|
+
if not token:
|
|
145
|
+
raise AttestationError("WhiteRoom attest response missing sessionToken")
|
|
146
|
+
return AttestedIdentity(session_token=token, arn=arn, issued_at=time.time())
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Wires session attachment into a boto3/botocore bedrock-runtime client.
|
|
3
|
+
|
|
4
|
+
An unmodified AWS SDK has no concept of WhiteRoom's session token — every
|
|
5
|
+
Bedrock call it makes is SigV4-signed with the caller's own credentials and
|
|
6
|
+
sent as-is. WhiteRoom's proxy discards that signature outright (it re-signs
|
|
7
|
+
outbound with its own assumed-role credentials — see sigv4.ts's STRIP list)
|
|
8
|
+
and instead identifies the caller from an `x-whiteroom-session` header. This
|
|
9
|
+
module is what attaches that header, transparently, to every call a hooked
|
|
10
|
+
client makes.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
from botocore.client import BaseClient
|
|
16
|
+
|
|
17
|
+
from .attest import SessionManager
|
|
18
|
+
|
|
19
|
+
SESSION_HEADER = "x-whiteroom-session"
|
|
20
|
+
|
|
21
|
+
# Botocore's event-name slug for the Bedrock Runtime service. Fixed — this
|
|
22
|
+
# module only ever hooks that one service.
|
|
23
|
+
_SERVICE_EVENT_ID = "bedrock-runtime"
|
|
24
|
+
|
|
25
|
+
# Force at most one re-attest-and-retry per call. Anything past that is a
|
|
26
|
+
# real, persistent rejection (bad server_id, revoked tenant, ...) and should
|
|
27
|
+
# surface to the caller as a normal ClientError rather than loop.
|
|
28
|
+
_MAX_FORCED_RETRIES = 1
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def install(
|
|
32
|
+
client: BaseClient,
|
|
33
|
+
*,
|
|
34
|
+
proxy_base_url: str,
|
|
35
|
+
server_id: str,
|
|
36
|
+
region: str | None = None,
|
|
37
|
+
botocore_session=None,
|
|
38
|
+
) -> SessionManager:
|
|
39
|
+
"""
|
|
40
|
+
Call once, right after constructing a bedrock-runtime client pointed at
|
|
41
|
+
a WhiteRoom proxy:
|
|
42
|
+
|
|
43
|
+
client = boto3.client(
|
|
44
|
+
"bedrock-runtime",
|
|
45
|
+
region_name="us-east-1",
|
|
46
|
+
endpoint_url="https://proxy.whiteroom.tech",
|
|
47
|
+
)
|
|
48
|
+
whiteroom_bedrock.install(
|
|
49
|
+
client,
|
|
50
|
+
proxy_base_url="https://proxy.whiteroom.tech",
|
|
51
|
+
server_id="proxy.whiteroom.tech",
|
|
52
|
+
region="us-east-1",
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
Everything past this point — invoke_model, converse, retries, streaming
|
|
56
|
+
refusals — flows through the client exactly as if it were talking to
|
|
57
|
+
real Bedrock. `region` defaults to the client's own configured region if
|
|
58
|
+
omitted; pass it explicitly if the client's credentials weren't resolved
|
|
59
|
+
from the default chain (e.g. an assumed role built by hand).
|
|
60
|
+
|
|
61
|
+
Returns the SessionManager backing this client, in case the caller wants
|
|
62
|
+
manual control (e.g. `manager.invalidate()` after rotating credentials).
|
|
63
|
+
"""
|
|
64
|
+
resolved_region = region or client.meta.region_name
|
|
65
|
+
if not resolved_region:
|
|
66
|
+
raise ValueError("region could not be determined from the client; pass region= explicitly")
|
|
67
|
+
|
|
68
|
+
manager = SessionManager(
|
|
69
|
+
proxy_base_url=proxy_base_url,
|
|
70
|
+
server_id=server_id,
|
|
71
|
+
region=resolved_region,
|
|
72
|
+
botocore_session=botocore_session,
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
events = client.meta.events
|
|
76
|
+
|
|
77
|
+
def _attach_session(request, **kwargs):
|
|
78
|
+
request.headers[SESSION_HEADER] = manager.get_token()
|
|
79
|
+
|
|
80
|
+
def _retry_on_session_rejection(response, attempts, **kwargs):
|
|
81
|
+
if response is None or attempts > _MAX_FORCED_RETRIES:
|
|
82
|
+
return None
|
|
83
|
+
http_response, _parsed = response
|
|
84
|
+
if http_response is not None and http_response.status_code == 403:
|
|
85
|
+
manager.invalidate()
|
|
86
|
+
return 0 # retry immediately; _attach_session will fetch a fresh token
|
|
87
|
+
return None
|
|
88
|
+
|
|
89
|
+
events.register(f"before-send.{_SERVICE_EVENT_ID}", _attach_session)
|
|
90
|
+
events.register(f"needs-retry.{_SERVICE_EVENT_ID}", _retry_on_session_rejection)
|
|
91
|
+
|
|
92
|
+
return manager
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: whiteroom-bedrock
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Attaches WhiteRoom session attestation to a boto3 bedrock-runtime client
|
|
5
|
+
License: BSL-1.1
|
|
6
|
+
Requires-Python: >=3.9
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
Requires-Dist: botocore>=1.34
|
|
9
|
+
Provides-Extra: test
|
|
10
|
+
Requires-Dist: pytest>=7; extra == "test"
|
|
11
|
+
Requires-Dist: boto3>=1.34; extra == "test"
|
|
12
|
+
|
|
13
|
+
# whiteroom-bedrock
|
|
14
|
+
|
|
15
|
+
Client-side half of WhiteRoom's Bedrock adapter (server side lives in
|
|
16
|
+
`apps/engine/src/bedrock/`).
|
|
17
|
+
|
|
18
|
+
An unmodified AWS SDK has no way to prove its identity to WhiteRoom the way
|
|
19
|
+
`ANTHROPIC_BASE_URL`/`OPENAI_BASE_URL` callers do with a bearer key — every
|
|
20
|
+
Bedrock call is SigV4-signed by the caller's own credentials, and WhiteRoom
|
|
21
|
+
discards that signature (it re-signs outbound with its own assumed-role
|
|
22
|
+
credentials). Proving who called instead happens once, out of band, by
|
|
23
|
+
replaying a signed `sts:GetCallerIdentity` probe to WhiteRoom's `/attest`
|
|
24
|
+
endpoint. This package does that handshake and attaches the resulting
|
|
25
|
+
session token to every subsequent call, so the rest of your agent code never
|
|
26
|
+
has to know WhiteRoom is in the loop.
|
|
27
|
+
|
|
28
|
+
## Install
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
pip install -e sdks/bedrock-python
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
## Test
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install -e "sdks/bedrock-python[test]"
|
|
38
|
+
pytest sdks/bedrock-python/tests
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
No network access or AWS credentials needed — everything's mocked at the
|
|
42
|
+
botocore/urllib layer. `tests/test_attest.py` is the one worth reading
|
|
43
|
+
closely: it directly checks the signed probe has `host`,
|
|
44
|
+
`x-whiteroom-server-id` and `x-amz-date` in `SignedHeaders`, which is
|
|
45
|
+
exactly what the server (`attest.ts`) requires — a change here that broke
|
|
46
|
+
that would otherwise only surface as every real attestation failing
|
|
47
|
+
against a live proxy.
|
|
48
|
+
|
|
49
|
+
## Use
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
import boto3
|
|
53
|
+
import whiteroom_bedrock
|
|
54
|
+
|
|
55
|
+
client = boto3.client(
|
|
56
|
+
"bedrock-runtime",
|
|
57
|
+
region_name="us-east-1",
|
|
58
|
+
endpoint_url="https://proxy.whiteroom.tech",
|
|
59
|
+
)
|
|
60
|
+
whiteroom_bedrock.install(
|
|
61
|
+
client,
|
|
62
|
+
proxy_base_url="https://proxy.whiteroom.tech",
|
|
63
|
+
server_id="proxy.whiteroom.tech", # must match the proxy's configured serverId
|
|
64
|
+
region="us-east-1",
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
# Everything from here down is a normal bedrock-runtime client.
|
|
68
|
+
response = client.converse(
|
|
69
|
+
modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
|
|
70
|
+
messages=[{"role": "user", "content": [{"text": "hi"}]}],
|
|
71
|
+
)
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## What `install()` actually does
|
|
75
|
+
|
|
76
|
+
- On first call (and again whenever the cached session goes stale), signs a
|
|
77
|
+
throwaway `GetCallerIdentity` request with whatever credentials this
|
|
78
|
+
process already resolves for AWS (env vars, shared config, an instance or
|
|
79
|
+
task role — the standard botocore chain) and POSTs it to `/attest`.
|
|
80
|
+
- Attaches the session token WhiteRoom hands back as `x-whiteroom-session`
|
|
81
|
+
on every `invoke_model`/`converse` call, via a `before-send` hook — after
|
|
82
|
+
signing, so it never interferes with the client's own (ultimately
|
|
83
|
+
discarded) signature.
|
|
84
|
+
- If the proxy ever answers 403 — the session expired, or WhiteRoom
|
|
85
|
+
restarted and lost its in-memory session table — invalidates the cached
|
|
86
|
+
token and forces exactly one retry, so callers don't see a transient
|
|
87
|
+
session death as a failure.
|
|
88
|
+
|
|
89
|
+
## Not handled here
|
|
90
|
+
|
|
91
|
+
- **Streaming.** `invoke_model_with_response_stream` and
|
|
92
|
+
`converse_stream` reach the proxy the same as any other call, but the
|
|
93
|
+
server actively refuses to proxy streaming responses today (see
|
|
94
|
+
`routes.ts`) — you'll get a `ValidationException` back, not a stream.
|
|
95
|
+
- **Credential source mismatch.** `region` and credentials are resolved
|
|
96
|
+
from the default botocore chain unless you pass `region=`/
|
|
97
|
+
`botocore_session=` explicitly. If the client you're hooking was built
|
|
98
|
+
with credentials that *don't* come from that chain (e.g. explicit
|
|
99
|
+
`aws_access_key_id=` on `boto3.client(...)`), attest with the same
|
|
100
|
+
credentials by passing a `botocore.session.Session` that resolves them.
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
src/whiteroom_bedrock/__init__.py
|
|
4
|
+
src/whiteroom_bedrock/attest.py
|
|
5
|
+
src/whiteroom_bedrock/hook.py
|
|
6
|
+
src/whiteroom_bedrock.egg-info/PKG-INFO
|
|
7
|
+
src/whiteroom_bedrock.egg-info/SOURCES.txt
|
|
8
|
+
src/whiteroom_bedrock.egg-info/dependency_links.txt
|
|
9
|
+
src/whiteroom_bedrock.egg-info/requires.txt
|
|
10
|
+
src/whiteroom_bedrock.egg-info/top_level.txt
|
|
11
|
+
tests/test_attest.py
|
|
12
|
+
tests/test_hook.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
whiteroom_bedrock
|
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Tests the probe construction and session caching in attest.py.
|
|
3
|
+
|
|
4
|
+
The probe-shape tests are the important ones: attest.ts on the server side
|
|
5
|
+
specifically requires host, x-whiteroom-server-id and x-amz-date to all be
|
|
6
|
+
in SignedHeaders (see apps/engine/src/bedrock/attest.ts and its own test
|
|
7
|
+
suite, bedrock-attest.test.ts). A change here that stops signing one of
|
|
8
|
+
those would only ever surface as every real attestation failing against a
|
|
9
|
+
live proxy -- this is the one place that's caught locally instead.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
import time
|
|
14
|
+
from unittest.mock import patch
|
|
15
|
+
from urllib.error import HTTPError, URLError
|
|
16
|
+
|
|
17
|
+
import pytest
|
|
18
|
+
from botocore.credentials import Credentials
|
|
19
|
+
|
|
20
|
+
from whiteroom_bedrock.attest import AttestationError, SessionManager
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class FakeCredentialSession:
|
|
24
|
+
"""Stands in for botocore.session.Session -- just get_credentials()."""
|
|
25
|
+
|
|
26
|
+
def __init__(self, creds=None):
|
|
27
|
+
self._creds = creds if creds is not None else Credentials(
|
|
28
|
+
"AKIAFAKEEXAMPLE", "fakesecret1234567890", token="fake-session-token"
|
|
29
|
+
)
|
|
30
|
+
|
|
31
|
+
def get_credentials(self):
|
|
32
|
+
return self._creds
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def make_manager(**overrides):
|
|
36
|
+
kwargs = dict(
|
|
37
|
+
proxy_base_url="https://proxy.whiteroom.tech",
|
|
38
|
+
server_id="proxy.whiteroom.tech",
|
|
39
|
+
region="us-east-1",
|
|
40
|
+
botocore_session=FakeCredentialSession(),
|
|
41
|
+
)
|
|
42
|
+
kwargs.update(overrides)
|
|
43
|
+
return SessionManager(**kwargs)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class TestBuildProbe:
|
|
47
|
+
def test_method_and_body_are_exactly_get_caller_identity(self):
|
|
48
|
+
probe = make_manager()._build_probe()
|
|
49
|
+
assert probe["method"] == "POST"
|
|
50
|
+
assert probe["body"] == "Action=GetCallerIdentity&Version=2011-06-15"
|
|
51
|
+
|
|
52
|
+
def test_signed_headers_include_everything_the_server_requires(self):
|
|
53
|
+
# Direct regression test for the server-side fix: attest.ts rejects
|
|
54
|
+
# a probe unless host, x-whiteroom-server-id AND x-amz-date are all
|
|
55
|
+
# in SignedHeaders (not just present as headers).
|
|
56
|
+
probe = make_manager()._build_probe()
|
|
57
|
+
auth = probe["headers"]["Authorization"]
|
|
58
|
+
signed_headers = auth.split("SignedHeaders=")[1].split(",")[0].split(";")
|
|
59
|
+
for required in ("host", "x-whiteroom-server-id", "x-amz-date"):
|
|
60
|
+
assert required in signed_headers, f"{required} missing from SignedHeaders: {signed_headers}"
|
|
61
|
+
|
|
62
|
+
def test_server_id_header_matches_what_was_configured(self):
|
|
63
|
+
probe = make_manager(server_id="my-proxy.example.com")._build_probe()
|
|
64
|
+
assert probe["headers"]["x-whiteroom-server-id"] == "my-proxy.example.com"
|
|
65
|
+
|
|
66
|
+
def test_host_targets_sts_in_the_configured_region(self):
|
|
67
|
+
probe = make_manager(region="eu-west-1")._build_probe()
|
|
68
|
+
assert probe["headers"]["Host"] == "sts.eu-west-1.amazonaws.com"
|
|
69
|
+
|
|
70
|
+
def test_includes_security_token_for_temporary_credentials(self):
|
|
71
|
+
probe = make_manager()._build_probe()
|
|
72
|
+
assert probe["headers"]["X-Amz-Security-Token"] == "fake-session-token"
|
|
73
|
+
|
|
74
|
+
def test_omits_security_token_for_static_credentials(self):
|
|
75
|
+
static_creds = Credentials("AKIASTATIC", "staticsecret")
|
|
76
|
+
probe = make_manager(botocore_session=FakeCredentialSession(static_creds))._build_probe()
|
|
77
|
+
assert "X-Amz-Security-Token" not in probe["headers"]
|
|
78
|
+
|
|
79
|
+
def test_raises_when_no_credentials_are_available(self):
|
|
80
|
+
class NoCreds:
|
|
81
|
+
def get_credentials(self):
|
|
82
|
+
return None
|
|
83
|
+
|
|
84
|
+
with pytest.raises(AttestationError, match="no AWS credentials"):
|
|
85
|
+
make_manager(botocore_session=NoCreds())._build_probe()
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
class TestGetToken:
|
|
89
|
+
def test_caches_within_the_ttl_window(self):
|
|
90
|
+
manager = make_manager()
|
|
91
|
+
calls = []
|
|
92
|
+
|
|
93
|
+
def fake_attest():
|
|
94
|
+
calls.append(1)
|
|
95
|
+
from whiteroom_bedrock.attest import AttestedIdentity
|
|
96
|
+
return AttestedIdentity(session_token=f"tok-{len(calls)}", arn="arn:aws:iam::123:user/a", issued_at=time.time())
|
|
97
|
+
|
|
98
|
+
with patch.object(manager, "_attest", side_effect=fake_attest):
|
|
99
|
+
first = manager.get_token()
|
|
100
|
+
second = manager.get_token()
|
|
101
|
+
|
|
102
|
+
assert first == second == "tok-1"
|
|
103
|
+
assert len(calls) == 1
|
|
104
|
+
|
|
105
|
+
def test_force_refreshes_even_within_the_ttl_window(self):
|
|
106
|
+
manager = make_manager()
|
|
107
|
+
calls = []
|
|
108
|
+
|
|
109
|
+
def fake_attest():
|
|
110
|
+
calls.append(1)
|
|
111
|
+
from whiteroom_bedrock.attest import AttestedIdentity
|
|
112
|
+
return AttestedIdentity(session_token=f"tok-{len(calls)}", arn="arn:aws:iam::123:user/a", issued_at=time.time())
|
|
113
|
+
|
|
114
|
+
with patch.object(manager, "_attest", side_effect=fake_attest):
|
|
115
|
+
manager.get_token()
|
|
116
|
+
second = manager.get_token(force=True)
|
|
117
|
+
|
|
118
|
+
assert second == "tok-2"
|
|
119
|
+
assert len(calls) == 2
|
|
120
|
+
|
|
121
|
+
def test_refreshes_once_the_ttl_hint_has_passed(self):
|
|
122
|
+
manager = make_manager()
|
|
123
|
+
calls = []
|
|
124
|
+
|
|
125
|
+
def fake_attest():
|
|
126
|
+
calls.append(1)
|
|
127
|
+
from whiteroom_bedrock.attest import AttestedIdentity
|
|
128
|
+
# issued_at backdated past the TTL hint, as if time had passed.
|
|
129
|
+
return AttestedIdentity(
|
|
130
|
+
session_token=f"tok-{len(calls)}",
|
|
131
|
+
arn="arn:aws:iam::123:user/a",
|
|
132
|
+
issued_at=time.time() - SessionManager.SESSION_TTL_HINT_SECONDS - 1,
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
with patch.object(manager, "_attest", side_effect=fake_attest):
|
|
136
|
+
manager.get_token()
|
|
137
|
+
second = manager.get_token()
|
|
138
|
+
|
|
139
|
+
assert second == "tok-2"
|
|
140
|
+
assert len(calls) == 2
|
|
141
|
+
|
|
142
|
+
def test_invalidate_forces_a_fresh_attest_on_the_next_call(self):
|
|
143
|
+
manager = make_manager()
|
|
144
|
+
calls = []
|
|
145
|
+
|
|
146
|
+
def fake_attest():
|
|
147
|
+
calls.append(1)
|
|
148
|
+
from whiteroom_bedrock.attest import AttestedIdentity
|
|
149
|
+
return AttestedIdentity(session_token=f"tok-{len(calls)}", arn="arn:aws:iam::123:user/a", issued_at=time.time())
|
|
150
|
+
|
|
151
|
+
with patch.object(manager, "_attest", side_effect=fake_attest):
|
|
152
|
+
manager.get_token()
|
|
153
|
+
manager.invalidate()
|
|
154
|
+
second = manager.get_token()
|
|
155
|
+
|
|
156
|
+
assert second == "tok-2"
|
|
157
|
+
assert len(calls) == 2
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
class TestAttestErrorHandling:
|
|
161
|
+
def test_wraps_an_http_error_from_the_proxy(self):
|
|
162
|
+
manager = make_manager()
|
|
163
|
+
err = HTTPError("https://proxy.whiteroom.tech/attest", 401, "Unauthorized", {}, None)
|
|
164
|
+
err.read = lambda: b'{"message":"server-ID mismatch"}'
|
|
165
|
+
with patch("urllib.request.urlopen", side_effect=err):
|
|
166
|
+
with pytest.raises(AttestationError, match="401"):
|
|
167
|
+
manager.get_token()
|
|
168
|
+
|
|
169
|
+
def test_wraps_a_network_error_reaching_the_proxy(self):
|
|
170
|
+
manager = make_manager()
|
|
171
|
+
with patch("urllib.request.urlopen", side_effect=URLError("connection refused")):
|
|
172
|
+
with pytest.raises(AttestationError, match="could not reach"):
|
|
173
|
+
manager.get_token()
|
|
174
|
+
|
|
175
|
+
def test_raises_when_the_proxy_response_has_no_session_token(self):
|
|
176
|
+
manager = make_manager()
|
|
177
|
+
|
|
178
|
+
class FakeResponse:
|
|
179
|
+
def read(self):
|
|
180
|
+
return json.dumps({"arn": "arn:aws:iam::123:user/a"}).encode()
|
|
181
|
+
|
|
182
|
+
def __enter__(self):
|
|
183
|
+
return self
|
|
184
|
+
|
|
185
|
+
def __exit__(self, *a):
|
|
186
|
+
return False
|
|
187
|
+
|
|
188
|
+
with patch("urllib.request.urlopen", return_value=FakeResponse()):
|
|
189
|
+
with pytest.raises(AttestationError, match="missing sessionToken"):
|
|
190
|
+
manager.get_token()
|
|
@@ -0,0 +1,139 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Tests install()'s session-attachment and 403-retry logic in hook.py.
|
|
3
|
+
|
|
4
|
+
Uses a fake botocore client (just enough of .meta.events.register to capture
|
|
5
|
+
the handlers install() wires up) rather than a real boto3 client, so these
|
|
6
|
+
run without network access or a real AWS SDK event bus -- the actual
|
|
7
|
+
end-to-end path (a real boto3 client, a real HTTP server) was validated
|
|
8
|
+
manually when this adapter was built; these pin the handler logic itself so
|
|
9
|
+
a future change can't silently break it.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
from unittest.mock import MagicMock, patch
|
|
13
|
+
|
|
14
|
+
import pytest
|
|
15
|
+
|
|
16
|
+
from whiteroom_bedrock.hook import SESSION_HEADER, install
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class FakeEvents:
|
|
20
|
+
"""Records registrations instead of actually dispatching anything."""
|
|
21
|
+
|
|
22
|
+
def __init__(self):
|
|
23
|
+
self.registered = {}
|
|
24
|
+
|
|
25
|
+
def register(self, event_name, handler):
|
|
26
|
+
self.registered[event_name] = handler
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class FakeMeta:
|
|
30
|
+
def __init__(self, region_name="us-east-1"):
|
|
31
|
+
self.region_name = region_name
|
|
32
|
+
self.events = FakeEvents()
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class FakeClient:
|
|
36
|
+
def __init__(self, region_name="us-east-1"):
|
|
37
|
+
self.meta = FakeMeta(region_name)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class FakeRequest:
|
|
41
|
+
def __init__(self):
|
|
42
|
+
self.headers = {}
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def install_with_fake_manager(client=None, **overrides):
|
|
46
|
+
"""install() always builds its own real SessionManager; tests patch its
|
|
47
|
+
get_token/invalidate afterward via the manager it returns."""
|
|
48
|
+
client = client or FakeClient()
|
|
49
|
+
kwargs = dict(proxy_base_url="https://proxy.whiteroom.tech", server_id="proxy.whiteroom.tech")
|
|
50
|
+
kwargs.update(overrides)
|
|
51
|
+
manager = install(client, **kwargs)
|
|
52
|
+
return client, manager
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class TestInstall:
|
|
56
|
+
def test_registers_before_send_and_needs_retry_for_bedrock_runtime(self):
|
|
57
|
+
client, _ = install_with_fake_manager()
|
|
58
|
+
assert "before-send.bedrock-runtime" in client.meta.events.registered
|
|
59
|
+
assert "needs-retry.bedrock-runtime" in client.meta.events.registered
|
|
60
|
+
|
|
61
|
+
def test_resolves_region_from_the_client_when_not_given_explicitly(self):
|
|
62
|
+
client = FakeClient(region_name="eu-west-1")
|
|
63
|
+
_, manager = install_with_fake_manager(client)
|
|
64
|
+
assert manager._region == "eu-west-1"
|
|
65
|
+
|
|
66
|
+
def test_explicit_region_overrides_the_client_default(self):
|
|
67
|
+
client = FakeClient(region_name="eu-west-1")
|
|
68
|
+
_, manager = install_with_fake_manager(client, region="ap-southeast-2")
|
|
69
|
+
assert manager._region == "ap-southeast-2"
|
|
70
|
+
|
|
71
|
+
def test_raises_when_no_region_can_be_determined(self):
|
|
72
|
+
client = FakeClient(region_name=None)
|
|
73
|
+
with pytest.raises(ValueError, match="region"):
|
|
74
|
+
install(client, proxy_base_url="https://proxy.whiteroom.tech", server_id="proxy.whiteroom.tech")
|
|
75
|
+
|
|
76
|
+
def test_returns_the_session_manager_for_manual_control(self):
|
|
77
|
+
_, manager = install_with_fake_manager()
|
|
78
|
+
assert hasattr(manager, "get_token")
|
|
79
|
+
assert hasattr(manager, "invalidate")
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class TestAttachSession:
|
|
83
|
+
def test_sets_the_session_header_from_get_token(self):
|
|
84
|
+
client, manager = install_with_fake_manager()
|
|
85
|
+
with patch.object(manager, "get_token", return_value="wrbs_testtoken123"):
|
|
86
|
+
handler = client.meta.events.registered["before-send.bedrock-runtime"]
|
|
87
|
+
request = FakeRequest()
|
|
88
|
+
handler(request)
|
|
89
|
+
assert request.headers[SESSION_HEADER] == "wrbs_testtoken123"
|
|
90
|
+
|
|
91
|
+
def test_does_not_force_a_refresh_on_every_call(self):
|
|
92
|
+
# get_token() is called with no force argument -- the hook relies on
|
|
93
|
+
# SessionManager's own TTL caching, it doesn't re-attest every call.
|
|
94
|
+
client, manager = install_with_fake_manager()
|
|
95
|
+
with patch.object(manager, "get_token", return_value="tok") as mock_get_token:
|
|
96
|
+
handler = client.meta.events.registered["before-send.bedrock-runtime"]
|
|
97
|
+
handler(FakeRequest())
|
|
98
|
+
mock_get_token.assert_called_once_with()
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
class TestRetryOnSessionRejection:
|
|
102
|
+
def _invoke(self, client, response, attempts):
|
|
103
|
+
handler = client.meta.events.registered["needs-retry.bedrock-runtime"]
|
|
104
|
+
return handler(response=response, attempts=attempts)
|
|
105
|
+
|
|
106
|
+
def test_forces_an_immediate_retry_on_the_first_403(self):
|
|
107
|
+
client, manager = install_with_fake_manager()
|
|
108
|
+
http_response = MagicMock(status_code=403)
|
|
109
|
+
with patch.object(manager, "invalidate") as mock_invalidate:
|
|
110
|
+
result = self._invoke(client, (http_response, None), attempts=1)
|
|
111
|
+
assert result == 0
|
|
112
|
+
mock_invalidate.assert_called_once()
|
|
113
|
+
|
|
114
|
+
def test_does_not_force_a_second_retry(self):
|
|
115
|
+
# Two failures in a row is treated as a real, persistent rejection
|
|
116
|
+
# (bad server_id, revoked tenant) -- not something re-attesting a
|
|
117
|
+
# second time will fix.
|
|
118
|
+
client, manager = install_with_fake_manager()
|
|
119
|
+
http_response = MagicMock(status_code=403)
|
|
120
|
+
with patch.object(manager, "invalidate") as mock_invalidate:
|
|
121
|
+
result = self._invoke(client, (http_response, None), attempts=2)
|
|
122
|
+
assert result is None
|
|
123
|
+
mock_invalidate.assert_not_called()
|
|
124
|
+
|
|
125
|
+
def test_ignores_non_403_statuses(self):
|
|
126
|
+
client, manager = install_with_fake_manager()
|
|
127
|
+
http_response = MagicMock(status_code=500)
|
|
128
|
+
with patch.object(manager, "invalidate") as mock_invalidate:
|
|
129
|
+
result = self._invoke(client, (http_response, None), attempts=1)
|
|
130
|
+
assert result is None
|
|
131
|
+
mock_invalidate.assert_not_called()
|
|
132
|
+
|
|
133
|
+
def test_ignores_a_missing_response(self):
|
|
134
|
+
client, _ = install_with_fake_manager()
|
|
135
|
+
assert self._invoke(client, None, attempts=1) is None
|
|
136
|
+
|
|
137
|
+
def test_ignores_a_missing_http_response(self):
|
|
138
|
+
client, _ = install_with_fake_manager()
|
|
139
|
+
assert self._invoke(client, (None, None), attempts=1) is None
|