diff --git a/agentkit/auth/admin.py b/agentkit/auth/admin.py index 33c70e45..79401e5c 100644 --- a/agentkit/auth/admin.py +++ b/agentkit/auth/admin.py @@ -46,6 +46,12 @@ # create sandbox sessions (the session/tool-read actions come from the custom policy). SANDBOX_ACCESS_POLICY = "AgentKitSandboxAccess" WELL_KNOWN_KEY = ".well-known/agentkit-cli" +_TOS_BUCKET_CONFLICT_CODES = frozenset( + { + "BucketAlreadyExists", + "BucketAlreadyOwnedByYou", + } +) ROLE_ACTIONS = ( "agentkit:CreateSession", "agentkit:GetSession", "agentkit:DeleteSession", "agentkit:GetSessionLogs", "agentkit:SetSessionTtl", @@ -386,9 +392,10 @@ def publish_discovery( client = tos.TosClientV2(ak, sk, endpoint, coords.region, security_token=token) try: client.create_bucket(bucket, acl=tos.ACLType.ACL_Public_Read) - except Exception as exc: # noqa: BLE001 + except Exception as exc: msg = str(exc) - if not any(k in msg for k in ("Exist", "exist", "Owned", "owned", "Conflict", "conflict")): + error_code = str(getattr(exc, "code", "")) + if error_code not in _TOS_BUCKET_CONFLICT_CODES: raise AuthError( f"could not create the TOS bucket for the discovery doc: {msg[:120]}", hint="enable TOS on this account and allow public-read buckets, or pass an existing bucket.", diff --git a/tests/auth/test_admin.py b/tests/auth/test_admin.py new file mode 100644 index 00000000..4f7d6c39 --- /dev/null +++ b/tests/auth/test_admin.py @@ -0,0 +1,94 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from agentkit.auth.admin import CliAccessCoords, publish_discovery +from agentkit.auth.errors import AuthError + + +class TosError(Exception): + def __init__(self, code: str, message: str): + super().__init__(message) + self.code = code + + +def install_fake_tos(monkeypatch, create_bucket_error: Exception): + client = MagicMock() + client.create_bucket.side_effect = create_bucket_error + tos_module = SimpleNamespace( + ACLType=SimpleNamespace(ACL_Public_Read="public-read"), + TosClientV2=MagicMock(return_value=client), + ) + monkeypatch.setitem(sys.modules, "tos", tos_module) + return client + + +def make_coords() -> CliAccessCoords: + return CliAccessCoords( + account_id="account", + region="cn-beijing", + user_pool_uid="pool", + issuer="https://issuer.example", + client_id="client", + role_trn="role", + provider_trn="provider", + ) + + +@pytest.mark.parametrize( + "error_code", + ["BucketAlreadyExists", "BucketAlreadyOwnedByYou"], +) +def test_publish_discovery_ignores_explicit_bucket_conflicts(monkeypatch, error_code): + client = install_fake_tos( + monkeypatch, + TosError(error_code, "bucket is already owned"), + ) + + url = publish_discovery( + make_coords(), + bucket="existing-bucket", + custom_domain="agent.example", + access_key="ak", + secret_key="sk", + ) + + assert url == "https://agent.example" + client.put_object.assert_called_once() + + +def test_publish_discovery_does_not_swallow_authentication_errors(monkeypatch): + client = install_fake_tos( + monkeypatch, + TosError( + "InvalidAccessKeyId", + "the specified access key does not exist", + ), + ) + + with pytest.raises(AuthError, match="could not create the TOS bucket"): + publish_discovery( + make_coords(), + bucket="new-bucket", + custom_domain="agent.example", + access_key="invalid-ak", + secret_key="invalid-sk", + ) + + client.put_object.assert_not_called()