File size: 13,129 Bytes
a42a6a0 44b1f9c a42a6a0 44b1f9c a42a6a0 44b1f9c a42a6a0 44b1f9c 96175fe be19956 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 | import httpx
import pytest
from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
import app.hub as hub_module
from app.deps import get_registration_limiter
from app.main import app as fastapi_app
from app.rate_limit import TokenBucket
from fakes import seed_agent
def test_list_and_filters(env):
seed_agent(env.hub, "agent-1", model="opus-4.7", hf_user="user-one")
seed_agent(env.hub, "agent-2", model="gemma-3", hf_user="user-two", bio="byte-level ideas")
data = env.client.get("/v1/agents").json()
assert data["count"] == 2
assert data["items"] == ["agent-1.md", "agent-2.md"]
assert env.client.get("/v1/agents?model=gemma-3").json()["matched"] == 1
assert env.client.get("/v1/agents?hf_user=user-one").json()["matched"] == 1
assert env.client.get("/v1/agents?q=byte-level").json()["matched"] == 1
def test_expand_returns_agent_info(env):
seed_agent(env.hub, "agent-1", bio="hello world")
items = env.client.get("/v1/agents?expand=true").json()["items"]
assert items[0]["agent_id"] == "agent-1"
assert items[0]["hf_user"] == "test-user"
assert items[0]["bio"] == "hello world"
def test_single_agent_lookup(env):
seed_agent(env.hub, "agent-1")
r = env.client.get("/v1/agents/agent-1")
assert r.status_code == 200
assert r.json()["agent_bucket"] == "test-org/test-agent-1"
assert env.client.get("/v1/agents/ghost").status_code == 404
def test_human_namespace_is_reserved_at_registration(env):
for reserved in ("human-cmpatino", "human"):
r = env.client.post(
"/v1/agents/register",
json={"agent_id": reserved, "model": "m", "harness": "h", "tools": []},
)
assert r.status_code == 400, reserved
assert "reserved" in r.json()["error"]["message"]
def test_normal_registration_still_works_and_is_immediately_listed(env):
bucket = "test-org/test-agent-9"
env.hub.buckets[bucket] = {}
env.hub.seed(".bucket-sync-handshake", "test-user", bucket=bucket)
r = env.client.post(
"/v1/agents/register",
json={"agent_id": "agent-9", "model": "m", "harness": "h", "tools": ["bash"]},
headers={"authorization": "Bearer hf_dummy"},
)
assert r.status_code == 201, r.json()
assert r.json()["hf_user"] == "test-user"
# write-through: visible to the read model without waiting for a listing
assert env.client.get("/v1/agents").json()["count"] == 1
# and the new agent is immediately mentionable
msg = env.client.post(
"/v1/messages", json={"agent_id": "agent-9", "body": "I have arrived"}
)
assert msg.status_code == 201
def _register(env, agent_id="agent-9", **extra):
return env.client.post(
"/v1/agents/register",
json={"agent_id": agent_id, "model": "m", "harness": "h", "tools": [], **extra},
headers={"authorization": "Bearer hf_caller"},
)
def test_fresh_registration_provisions_bucket_with_caller_token(env):
r = _register(env)
assert r.status_code == 201, r.json()
bucket = "test-org/test-agent-9"
assert env.hub.created_buckets == [(bucket, "hf_caller")]
assert env.hub.caller_writes == [(bucket, ".bucket-sync-handshake", "hf_caller")]
assert env.hub.buckets[bucket][".bucket-sync-handshake"] == b"test-user"
BUCKET = "test-org/test-agent-9"
HANDSHAKE = ".bucket-sync-handshake"
def test_matching_handshake_is_still_proved_by_a_caller_write(env):
env.hub.seed(HANDSHAKE, "test-user", bucket=BUCKET)
env.hub.bucket_owners[BUCKET] = "test-user"
assert _register(env).status_code == 201
assert env.hub.created_buckets == []
assert env.hub.caller_writes == [(BUCKET, HANDSHAKE, "hf_caller")]
def test_handshake_planted_in_someone_elses_bucket_is_refused(env):
# An attacker pre-creates the bucket a victim will register and writes the
# victim's public username into it. Matching text is not proof: the
# victim's token cannot write there, so registration must refuse.
env.hub.seed(HANDSHAKE, "test-user", bucket=BUCKET)
env.hub.bucket_owners[BUCKET] = "attacker"
r = _register(env)
assert r.status_code == 403
assert r.json()["error"]["code"] == "BUCKET_NOT_YOURS"
assert env.hub.caller_writes == []
assert env.client.get("/v1/agents").json()["count"] == 0
def _victim_state(env):
"""Victim's registration of agent-9 and the handshake in its bucket."""
seed_agent(env.hub, "agent-9", hf_user="victim")
env.hub.seed(HANDSHAKE, "victim", bucket=BUCKET)
env.hub.bucket_owners[BUCKET] = "victim"
env.hub.caller_is_org_admin = True # the caller could write anywhere
return dict(env.hub.buckets[env.settings.central_bucket]), dict(env.hub.buckets[BUCKET])
def _assert_untouched(env, before):
central, bucket = before
assert env.hub.buckets[env.settings.central_bucket] == central
assert env.hub.buckets[BUCKET] == bucket
assert env.hub.caller_writes == [] and env.hub.created_buckets == []
def test_registration_read_outage_aborts_before_provisioning(env):
before = _victim_state(env)
env.hub.failing_reads.add("agents/agent-9.md")
r = _register(env, force=True)
assert r.status_code == 503
assert r.json()["error"]["code"] == "HUB_UNAVAILABLE"
_assert_untouched(env, before)
def test_handshake_read_outage_aborts_before_writing(env):
before = _victim_state(env)
env.hub.buckets[env.settings.central_bucket].pop("agents/agent-9.md")
before = (dict(env.hub.buckets[env.settings.central_bucket]), before[1])
env.hub.failing_reads.add(f"{BUCKET}/{HANDSHAKE}")
r = _register(env)
assert r.status_code == 503
_assert_untouched(env, before)
def test_foreign_handshake_blocks_even_a_caller_who_could_write(env):
# An unregistered bucket whose handshake names another user is theirs; an
# org admin's token could overwrite it, but registration must not.
before = _victim_state(env)
env.hub.buckets[env.settings.central_bucket].pop("agents/agent-9.md")
before = (dict(env.hub.buckets[env.settings.central_bucket]), before[1])
r = _register(env)
assert r.status_code == 403
assert r.json()["error"]["code"] == "BUCKET_NOT_YOURS"
assert "victim" in r.json()["error"]["message"]
_assert_untouched(env, before)
SECRET = "hf_SECRETtokenDoNotLeak123"
def test_caller_token_never_reaches_responses_or_logs(env, monkeypatch, caplog):
caplog.set_level("DEBUG")
headers = {"authorization": f"Bearer {SECRET}"}
body = {"agent_id": "agent-9", "model": "m", "harness": "h", "tools": []}
post = lambda: env.client.post("/v1/agents/register", json=body, headers=headers)
texts = []
env.hub.whoami_fails = True # 401
texts.append(post()); env.hub.whoami_fails = False
env.hub.whoami_orgs = set() # 403 NOT_ORG_MEMBER
texts.append(post()); env.hub.whoami_orgs = {env.settings.org}
env.hub.failing_reads.add("agents/agent-9.md") # 503 lookup outage
texts.append(post()); env.hub.failing_reads.clear()
def forbid(bucket, token):
raise PermissionError(f"403 for token {token}")
monkeypatch.setattr(env.hub, "create_bucket_as", forbid) # 403 BUCKET_CREATE_FORBIDDEN
texts.append(post()); monkeypatch.undo()
env.hub.seed(HANDSHAKE, "test-user", bucket=BUCKET) # 403 BUCKET_NOT_YOURS
env.hub.bucket_owners[BUCKET] = "attacker"
texts.append(post())
env.hub.bucket_owners[BUCKET] = "test-user" # 201, audited
texts.append(post())
assert [r.status_code for r in texts] == [401, 403, 503, 403, 403, 201]
for r in texts:
assert SECRET not in r.text
assert SECRET not in caplog.text
audit = b"".join(env.hub.buckets[env.settings.audit_bucket].values())
assert SECRET.encode() not in audit
def test_own_bucket_without_handshake_gets_one(env):
bucket = "test-org/test-agent-9"
env.hub.buckets[bucket] = {}
env.hub.bucket_owners[bucket] = "test-user"
assert _register(env).status_code == 201
assert env.hub.caller_writes == [(bucket, ".bucket-sync-handshake", "hf_caller")]
def test_non_member_gets_invite_url(make_env):
env = make_env(INVITE_URL="https://hf.co/invite/abc")
env.hub.whoami_orgs = {"other-org"}
r = _register(env)
assert r.status_code == 403
err = r.json()["error"]
assert err["code"] == "NOT_ORG_MEMBER"
assert "https://hf.co/invite/abc" in err["message"]
assert env.hub.created_buckets == []
def test_non_member_without_invite_url_asks_organizer(env):
env.hub.whoami_orgs = set()
r = _register(env)
assert "ask the organizer for the invite link" in r.json()["error"]["message"]
def test_foreign_bucket_is_not_yours(env):
bucket = "test-org/test-agent-9"
env.hub.seed(".bucket-sync-handshake", "someone-else", bucket=bucket)
env.hub.bucket_owners[bucket] = "someone-else"
r = _register(env)
assert r.status_code == 403
err = r.json()["error"]
assert err["code"] == "BUCKET_NOT_YOURS"
assert "pick another agent_id" in err["hint"]
assert env.hub.buckets[bucket][".bucket-sync-handshake"] == b"someone-else"
assert env.client.get("/v1/agents").json()["count"] == 0
def test_transient_whoami_error_is_503_not_401(env):
env.hub.whoami_unreachable = True
r = _register(env)
assert r.status_code == 503
assert "retry" in r.json()["error"]["message"]
assert env.hub.created_buckets == []
def test_bad_token_is_401(env):
env.hub.whoami_fails = True
assert _register(env).status_code == 401
def test_duplicate_without_force_says_it_is_yours(env):
assert _register(env).status_code == 201
r = _register(env)
assert r.status_code == 409
assert "already registered to you" in r.json()["error"]["message"]
assert _register(env, force=True).status_code == 201
def test_rate_limit_is_keyed_by_user_after_auth(env):
limiter = TokenBucket(capacity=1, refill_per_minute=1)
fastapi_app.dependency_overrides[get_registration_limiter] = lambda: limiter
# A failed auth burns nothing.
env.hub.whoami_fails = True
assert _register(env).status_code == 401
env.hub.whoami_fails = False
assert _register(env, agent_id="agent-a").status_code == 201
# Same user, another agent_id: still limited.
assert _register(env, agent_id="agent-b").status_code == 429
# A different user is not throttled by the first user's attempts.
env.hub.whoami_user = "other-user"
assert _register(env, agent_id="agent-b").status_code == 201
def _hub_error(status):
resp = httpx.Response(status, request=httpx.Request("GET", "https://hf.co/api/x"))
return HfHubHTTPError(f"{status}", response=resp)
def test_bucket_exists_maps_only_404_to_false(env, monkeypatch):
client = hub_module.HubClient(env.settings)
def raise_(status):
def f(*a, **k):
raise _hub_error(status)
return f
monkeypatch.setattr(hub_module, "bucket_info", raise_(404))
assert client.bucket_exists("test-org/x") is False
monkeypatch.setattr(hub_module, "bucket_info", raise_(503))
with pytest.raises(hub_module.HubUnreachable):
client.bucket_exists("test-org/x")
def test_transient_bucket_check_is_503(env, monkeypatch):
def unreachable(bucket):
raise hub_module.HubUnreachable("503")
monkeypatch.setattr(env.hub, "bucket_exists", unreachable)
r = _register(env)
assert r.status_code == 503
assert r.json()["error"]["code"] == "HUB_UNAVAILABLE"
def test_optional_reads_map_only_missing_entries_to_none(env, monkeypatch):
client = hub_module.HubClient(env.settings)
def missing(**k):
raise EntryNotFoundError("no such file")
def outage(**k):
raise _hub_error(500)
monkeypatch.setattr(hub_module, "download_bucket_files", missing)
assert client.read_text_optional("hf://buckets/test-org/test-agent-9/.bucket-sync-handshake") is None
assert client.read_central_bytes_optional("agents/agent-9.md") is None
monkeypatch.setattr(hub_module, "download_bucket_files", outage)
with pytest.raises(HfHubHTTPError):
client.read_text_optional("hf://buckets/test-org/test-agent-9/.bucket-sync-handshake")
with pytest.raises(HfHubHTTPError):
client.read_central_bytes_optional("agents/agent-9.md")
def test_rejected_signup_leaves_registration_working(env):
# After a refused handshake write, the next signups still go all the way
# through the handshake write (a 409 alone would never reach it).
env.hub.seed(HANDSHAKE, "test-user", bucket=BUCKET)
env.hub.bucket_owners[BUCKET] = "attacker"
assert _register(env).status_code == 403
fresh = _register(env, agent_id="agent-10")
assert fresh.status_code == 201, fresh.json()
assert _register(env, agent_id="agent-10", force=True).status_code == 201
assert env.hub.caller_writes == [("test-org/test-agent-10", HANDSHAKE, "hf_caller")] * 2
post = env.client.post("/v1/messages", json={"agent_id": "agent-10", "body": "still here"})
assert post.status_code == 201
|