diff --git a/agentkit/auth/admin.py b/agentkit/auth/admin.py index 33c70e45..aa62d7ab 100644 --- a/agentkit/auth/admin.py +++ b/agentkit/auth/admin.py @@ -383,27 +383,30 @@ def publish_discovery( sk = secret_key or _os.getenv("VOLCENGINE_SECRET_KEY") token = session_token or _os.getenv("VOLCENGINE_SESSION_TOKEN") endpoint = f"tos-{coords.region}.volces.com" - client = tos.TosClientV2(ak, sk, endpoint, coords.region, security_token=token) + client = tos.TosClientV2(ak, sk, endpoint, coords.region, security_token=token, dns_cache_time=0) try: - client.create_bucket(bucket, acl=tos.ACLType.ACL_Public_Read) - except Exception as exc: # noqa: BLE001 - msg = str(exc) - if not any(k in msg for k in ("Exist", "exist", "Owned", "owned", "Conflict", "conflict")): - 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.", - ) from exc - client.put_object( - bucket, WELL_KNOWN_KEY, content=json.dumps(coords.discovery_doc()).encode(), - acl=tos.ACLType.ACL_Public_Read, content_type="application/json", - ) - if custom_domain: - with __import__("contextlib").suppress(Exception): - client.put_bucket_custom_domain(bucket, domain=custom_domain) # best-effort bind - return f"https://{custom_domain}" - url = f"https://{bucket}.{endpoint}" - _verify_public_discovery(url) # fail loudly if the doc isn't anonymously readable - return url + try: + client.create_bucket(bucket, acl=tos.ACLType.ACL_Public_Read) + except Exception as exc: # noqa: BLE001 + msg = str(exc) + if not any(k in msg for k in ("Exist", "exist", "Owned", "owned", "Conflict", "conflict")): + 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.", + ) from exc + client.put_object( + bucket, WELL_KNOWN_KEY, content=json.dumps(coords.discovery_doc()).encode(), + acl=tos.ACLType.ACL_Public_Read, content_type="application/json", + ) + if custom_domain: + with __import__("contextlib").suppress(Exception): + client.put_bucket_custom_domain(bucket, domain=custom_domain) # best-effort bind + return f"https://{custom_domain}" + url = f"https://{bucket}.{endpoint}" + _verify_public_discovery(url) # fail loudly if the doc isn't anonymously readable + return url + finally: + client.close() def _verify_public_discovery(base_url: str) -> None: diff --git a/agentkit/toolkit/builders/ve_pipeline.py b/agentkit/toolkit/builders/ve_pipeline.py index eaf9aa1f..63ce8df8 100644 --- a/agentkit/toolkit/builders/ve_pipeline.py +++ b/agentkit/toolkit/builders/ve_pipeline.py @@ -421,8 +421,13 @@ def remove_artifact(self, config: VeCPCRBuilderConfig) -> bool: None, ) tos_service = TOSService(tos_config, provider=provider) - tos_service.delete_file(builder_config.tos_object_key) - logger.info(f"Deleted TOS archive: {builder_config.tos_object_key}") + try: + tos_service.delete_file(builder_config.tos_object_key) + logger.info( + f"Deleted TOS archive: {builder_config.tos_object_key}" + ) + finally: + tos_service.close() except Exception as e: logger.warning(f"Failed to delete TOS archive: {str(e)}") @@ -945,148 +950,155 @@ def _upload_to_tos( getattr(config, "common_config", None), "cloud_provider", None ) tos_service = TOSService(tos_config, provider=provider) + try: + # Two-step safety check: + # 1) Ensure the bucket exists and is accessible. + # 2) Verify the bucket is owned by the current account via ListBuckets before uploading. + import time + + created_in_this_run = False + name_conflict_not_created = False + + # Step 1: ensure bucket exists / accessible + if auto_created_bucket: + self.reporter.info( + f"Creating auto-generated TOS bucket in current account: {bucket_name}" + ) - # Two-step safety check: - # 1) Ensure the bucket exists and is accessible. - # 2) Verify the bucket is owned by the current account via ListBuckets before uploading. - import time - - created_in_this_run = False - name_conflict_not_created = False - - # Step 1: ensure bucket exists / accessible - if auto_created_bucket: - self.reporter.info( - f"Creating auto-generated TOS bucket in current account: {bucket_name}" - ) - - # Very low probability: name collision. Retry with a new generated name. - max_attempts = 3 - for attempt in range(1, max_attempts + 1): - tos_service.config.bucket = bucket_name - try: - tos_service.create_bucket() - created_in_this_run = True - break - except tos.exceptions.TosServerError as e: - if e.status_code == 409 and attempt < max_attempts: - bucket_name = TOSService.generate_bucket_name() - self.reporter.warning( - "Auto-generated bucket name already taken, retrying with a new name " - f"(attempt {attempt + 1}/{max_attempts}): {bucket_name}" - ) - continue - raise - else: - # User-specified bucket: if not accessible/existing, attempt to create. - self.reporter.info(f"Checking TOS bucket accessibility: {bucket_name}") - if not tos_service.bucket_exists(): - self.reporter.warning( - f"TOS bucket '{bucket_name}' is not accessible or does not exist, attempting to create it..." + # Very low probability: name collision. Retry with a new generated name. + max_attempts = 3 + for attempt in range(1, max_attempts + 1): + tos_service.config.bucket = bucket_name + try: + tos_service.create_bucket() + created_in_this_run = True + break + except tos.exceptions.TosServerError as e: + if e.status_code == 409 and attempt < max_attempts: + bucket_name = TOSService.generate_bucket_name() + self.reporter.warning( + "Auto-generated bucket name already taken, retrying with a new name " + f"(attempt {attempt + 1}/{max_attempts}): {bucket_name}" + ) + continue + raise + else: + # User-specified bucket: if not accessible/existing, attempt to create. + self.reporter.info( + f"Checking TOS bucket accessibility: {bucket_name}" ) + if not tos_service.bucket_exists(): + self.reporter.warning( + f"TOS bucket '{bucket_name}' is not accessible or does not exist, attempting to create it..." + ) + try: + tos_service.create_bucket() + created_in_this_run = True + except tos.exceptions.TosServerError as e: + if e.status_code == 409: + name_conflict_not_created = True + else: + raise + + # Step 2: verify bucket ownership via ListBuckets + self.reporter.info(f"Verifying TOS bucket ownership: {bucket_name}") + + def check_owned() -> bool: try: - tos_service.create_bucket() - created_in_this_run = True - except tos.exceptions.TosServerError as e: - if e.status_code == 409: - name_conflict_not_created = True - else: - raise - - # Step 2: verify bucket ownership via ListBuckets - self.reporter.info(f"Verifying TOS bucket ownership: {bucket_name}") + return tos_service.bucket_is_owned(bucket_name) + except Exception as e: + error_msg = ( + "Failed to determine TOS bucket ownership via ListBuckets. " + "Upload has been blocked for security reasons. " + "Please ensure your credentials have TOS ListBuckets permission, or set 'tos_bucket: Auto'." + ) + self.reporter.error(error_msg) + logger.error(f"Bucket ownership check failed: {str(e)}") + raise Exception(error_msg) + + if created_in_this_run: + # ListBuckets may be eventually consistent shortly after creation. + timeout_s = 10 + interval_s = 2 + deadline = time.time() + timeout_s + owned = False + while time.time() < deadline: + owned = check_owned() + if owned: + break + time.sleep(interval_s) + else: + owned = check_owned() - def check_owned() -> bool: - try: - return tos_service.bucket_is_owned(bucket_name) - except Exception as e: + if not owned: error_msg = ( - "Failed to determine TOS bucket ownership via ListBuckets. " - "Upload has been blocked for security reasons. " - "Please ensure your credentials have TOS ListBuckets permission, or set 'tos_bucket: Auto'." + f"Security notice: The configured TOS bucket '{bucket_name}' is not owned by the current account. " + "To prevent uploading your source code to a bucket you do not own (which could leak secrets), this upload has been blocked. " + "Please choose a bucket owned by your account, use 'agentkit config --tos_bucket ' to set it." ) - self.reporter.error(error_msg) - logger.error(f"Bucket ownership check failed: {str(e)}") raise Exception(error_msg) - if created_in_this_run: - # ListBuckets may be eventually consistent shortly after creation. - timeout_s = 10 - interval_s = 2 - deadline = time.time() + timeout_s - owned = False - while time.time() < deadline: - owned = check_owned() - if owned: - break - time.sleep(interval_s) - else: - owned = check_owned() - - if not owned: - error_msg = ( - f"Security notice: The configured TOS bucket '{bucket_name}' is not owned by the current account. " - "To prevent uploading your source code to a bucket you do not own (which could leak secrets), this upload has been blocked. " - "Please choose a bucket owned by your account, use 'agentkit config --tos_bucket ' to set it." - ) - raise Exception(error_msg) - - if name_conflict_not_created and owned: - actual_location = tos_service.get_bucket_location(bucket_name) - current_region = getattr( - tos_service, "actual_region", config.tos_region - ) - if actual_location: - raise Exception( - f"TOS bucket '{bucket_name}' exists in region '{actual_location}', " - f"but the current configuration targets region '{current_region}'. " - "TOS buckets are region-scoped and cannot be accessed across regions. " - "Please either switch to the bucket's region or use a different bucket name " - f"(e.g. '{bucket_name}-{current_region.split('-')[-1]}')." + if name_conflict_not_created and owned: + actual_location = tos_service.get_bucket_location(bucket_name) + current_region = getattr( + tos_service, "actual_region", config.tos_region ) - else: - raise Exception( - f"TOS bucket '{bucket_name}' is owned by your account but does not exist " - f"in the current region '{current_region}'. It may reside in another region. " - "TOS buckets are region-scoped and cannot be accessed across regions. " - "Please either switch to the bucket's region or use a different bucket name." - ) - - self.reporter.success( - f"TOS bucket ownership verified for current account: {bucket_name}" - ) - - if created_in_this_run: - data_plane_timeout_s = 30 - data_plane_interval_s = 3 - data_plane_deadline = time.time() + data_plane_timeout_s - while not tos_service.bucket_exists(): - if time.time() >= data_plane_deadline: + if actual_location: + raise Exception( + f"TOS bucket '{bucket_name}' exists in region '{actual_location}', " + f"but the current configuration targets region '{current_region}'. " + "TOS buckets are region-scoped and cannot be accessed across regions. " + "Please either switch to the bucket's region or use a different bucket name " + f"(e.g. '{bucket_name}-{current_region.split('-')[-1]}')." + ) + else: raise Exception( - f"TOS bucket '{bucket_name}' was created but is not yet " - "available for uploads. Please retry in a few seconds." + f"TOS bucket '{bucket_name}' is owned by your account but does not exist " + f"in the current region '{current_region}'. It may reside in another region. " + "TOS buckets are region-scoped and cannot be accessed across regions. " + "Please either switch to the bucket's region or use a different bucket name." ) - time.sleep(data_plane_interval_s) - # Update config with auto-generated bucket name if applicable - if auto_created_bucket: - config.tos_bucket = bucket_name + self.reporter.success( + f"TOS bucket ownership verified for current account: {bucket_name}" + ) - # Generate object key for the archive - archive_name = os.path.basename(archive_path) - object_key = f"{config.tos_prefix}/{archive_name}" + if created_in_this_run: + data_plane_timeout_s = 30 + data_plane_interval_s = 3 + data_plane_deadline = time.time() + data_plane_timeout_s + while not tos_service.bucket_exists(): + if time.time() >= data_plane_deadline: + raise Exception( + f"TOS bucket '{bucket_name}' was created but is not yet " + "available for uploads. Please retry in a few seconds." + ) + time.sleep(data_plane_interval_s) - # Upload file to TOS - tos_url = tos_service.upload_file(archive_path, object_key) + # Update config with auto-generated bucket name if applicable + if auto_created_bucket: + config.tos_bucket = bucket_name - # Get the actual region resolved by the service, or fallback to config - actual_region = getattr(tos_service, "actual_region", config.tos_region) + # Generate object key for the archive + archive_name = os.path.basename(archive_path) + object_key = f"{config.tos_prefix}/{archive_name}" - # Save object key to config for later reference - config.tos_object_key = object_key + # Upload file to TOS + tos_url = tos_service.upload_file(archive_path, object_key) - logger.info(f"File uploaded to TOS: {tos_url} (Region: {actual_region})") - return tos_url, actual_region + # Get the actual region resolved by the service, or fallback to config + actual_region = getattr(tos_service, "actual_region", config.tos_region) + + # Save object key to config for later reference + config.tos_object_key = object_key + + logger.info( + f"File uploaded to TOS: {tos_url} (Region: {actual_region})" + ) + return tos_url, actual_region + finally: + if hasattr(tos_service, "close"): + tos_service.close() except Exception as e: if "AccountDisable" in str(e): diff --git a/agentkit/toolkit/cli/cli_skills_workflow.py b/agentkit/toolkit/cli/cli_skills_workflow.py index f39be40c..b51ceada 100644 --- a/agentkit/toolkit/cli/cli_skills_workflow.py +++ b/agentkit/toolkit/cli/cli_skills_workflow.py @@ -248,56 +248,59 @@ def _ensure_bucket_ready( prefix=(prefix or "").strip(), ) ) - - exists = service.bucket_exists() - created_in_this_run = False - if not exists: - if assume_no: - raise typer.BadParameter(f"TOS bucket not found: {bucket_name}") - if auto_bucket or assume_yes: - service.create_bucket() - created_in_this_run = True - else: - if _is_interactive(): - typer.confirm( - f"TOS bucket '{bucket_name}' not found. Create it in current account?", - abort=True, - ) + try: + exists = service.bucket_exists() + created_in_this_run = False + if not exists: + if assume_no: + raise typer.BadParameter(f"TOS bucket not found: {bucket_name}") + if auto_bucket or assume_yes: service.create_bucket() created_in_this_run = True else: + if _is_interactive(): + typer.confirm( + f"TOS bucket '{bucket_name}' not found. Create it in current account?", + abort=True, + ) + service.create_bucket() + created_in_this_run = True + else: + raise typer.BadParameter( + f"TOS bucket '{bucket_name}' not found. Use -y/--yes to create it automatically." + ) + + def check_owned() -> bool: + try: + return service.bucket_is_owned(bucket_name) + except Exception as e: raise typer.BadParameter( - f"TOS bucket '{bucket_name}' not found. Use -y/--yes to create it automatically." - ) + "Failed to determine TOS bucket ownership via ListBuckets. " + "Upload has been blocked for security reasons. " + "Please ensure your credentials have TOS ListBuckets permission, or configure a bucket you own." + ) from e - def check_owned() -> bool: - try: - return service.bucket_is_owned(bucket_name) - except Exception as e: - raise typer.BadParameter( - "Failed to determine TOS bucket ownership via ListBuckets. " - "Upload has been blocked for security reasons. " - "Please ensure your credentials have TOS ListBuckets permission, or configure a bucket you own." - ) from e - - if created_in_this_run: - import time - - deadline = time.time() + 10 - while time.time() < deadline: - if check_owned(): - break - time.sleep(2) + if created_in_this_run: + import time + + deadline = time.time() + 10 + while time.time() < deadline: + if check_owned(): + break + time.sleep(2) + else: + raise typer.BadParameter( + f"Failed to verify ownership for newly created TOS bucket: {bucket_name}" + ) else: - raise typer.BadParameter( - f"Failed to verify ownership for newly created TOS bucket: {bucket_name}" - ) - else: - if not check_owned(): - raise typer.BadParameter( - f"Security notice: The configured TOS bucket '{bucket_name}' is not owned by the current account. " - "To prevent uploading your code to a bucket you do not own, this upload has been blocked." - ) + if not check_owned(): + raise typer.BadParameter( + f"Security notice: The configured TOS bucket '{bucket_name}' is not owned by the current account. " + "To prevent uploading your code to a bucket you do not own, this upload has been blocked." + ) + finally: + if hasattr(service, "close"): + service.close() def _tos_upload( @@ -327,14 +330,18 @@ def _tos_upload( prefix=effective_prefix, ) ) - if verify_bucket: - if not service.bucket_exists(): - raise typer.BadParameter(f"Bucket not found: {bucket}") - if not service.bucket_is_owned(bucket): - raise typer.BadParameter( - f"Bucket is not owned by current credentials: {bucket}" - ) - return service.upload_file(zip_abs, key) + try: + if verify_bucket: + if not service.bucket_exists(): + raise typer.BadParameter(f"Bucket not found: {bucket}") + if not service.bucket_is_owned(bucket): + raise typer.BadParameter( + f"Bucket is not owned by current credentials: {bucket}" + ) + return service.upload_file(zip_abs, key) + finally: + if hasattr(service, "close"): + service.close() def _pick_latest_version( diff --git a/agentkit/toolkit/cli/sandbox/tos_config.py b/agentkit/toolkit/cli/sandbox/tos_config.py index c046e8ac..a2625abd 100644 --- a/agentkit/toolkit/cli/sandbox/tos_config.py +++ b/agentkit/toolkit/cli/sandbox/tos_config.py @@ -64,11 +64,15 @@ def build_create_tool_tos_mount_config( region=region, ) ) - mount_config = service.build_mount_config( - bucket_path=DEFAULT_TOS_BUCKET_PATH, - local_mount_path=resolved_local_mount_path, - ) - return to_create_tool_tos_mount_config(mount_config) + try: + mount_config = service.build_mount_config( + bucket_path=DEFAULT_TOS_BUCKET_PATH, + local_mount_path=resolved_local_mount_path, + ) + return to_create_tool_tos_mount_config(mount_config) + finally: + if hasattr(service, "close"): + service.close() def to_create_tool_tos_mount_config( diff --git a/agentkit/toolkit/volcengine/services/tos_service.py b/agentkit/toolkit/volcengine/services/tos_service.py index b644ecf5..7ec45e22 100644 --- a/agentkit/toolkit/volcengine/services/tos_service.py +++ b/agentkit/toolkit/volcengine/services/tos_service.py @@ -111,6 +111,7 @@ def _init_client(self) -> None: ep.host, ep.region, security_token=getattr(creds, "session_token", None) or "", + dns_cache_time=0, ) self.credentials = creds # Expose the actual region resolved by VolcConfiguration @@ -125,6 +126,23 @@ def _init_client(self) -> None: logger.error(f"Failed to initialize TOS client: {str(e)}") raise + def close(self) -> None: + """Close the TOS client and release underlying connections and resources.""" + client, self.client = self.client, None + if client is None: + return + + try: + client.close() + except Exception: + logger.debug("Error closing TOS client", exc_info=True) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + def upload_file(self, local_path: str, object_key: str) -> str: """Upload a file to TOS. diff --git a/tests/toolkit/builders/test_ve_pipeline_upload_tos.py b/tests/toolkit/builders/test_ve_pipeline_upload_tos.py index 20342f1a..4a97073e 100644 --- a/tests/toolkit/builders/test_ve_pipeline_upload_tos.py +++ b/tests/toolkit/builders/test_ve_pipeline_upload_tos.py @@ -27,6 +27,7 @@ - ``bucket_is_owned(name)`` -> bool (ListBuckets ownership gate) - ``get_bucket_location(name)``-> Optional[str] - ``upload_file(path, key)`` -> str (URL) + - ``close()`` -> releases client resources - ``generate_bucket_name()`` -> staticmethod - ``.config`` (has a mutable ``.bucket`` attr) and ``.actual_region`` @@ -135,6 +136,15 @@ def upload_file(self, local_path: str, object_key: str) -> str: self.calls.append(("upload_file", local_path, object_key)) return self.upload_file_returns + def close(self) -> None: + self.calls.append(("close",)) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + @pytest.fixture(autouse=True) def _reset_fake_state(monkeypatch): @@ -211,6 +221,7 @@ def test_upload_blocked_when_bucket_not_owned_by_current_account(monkeypatch, tm called_methods = [c[0] for c in svc.calls] assert "bucket_is_owned" in called_methods assert "upload_file" not in called_methods + assert called_methods[-1] == "close" # --------------------------------------------------------------------------- @@ -474,3 +485,4 @@ def test_happy_path_uploads_and_returns_url_and_actual_region(monkeypatch, tmp_p # object key persisted back onto the config for later reference (L1086). assert config.tos_object_key == "agentkit-builds/app.tar.gz" + assert svc.calls[-1] == ("close",) diff --git a/tests/toolkit/volcengine/test_tos_service.py b/tests/toolkit/volcengine/test_tos_service.py new file mode 100644 index 00000000..42094c17 --- /dev/null +++ b/tests/toolkit/volcengine/test_tos_service.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 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. + +from agentkit.toolkit.volcengine.services.tos_service import TOSService + + +class _FakeClient: + def __init__(self) -> None: + self.close_calls = 0 + + def close(self) -> None: + self.close_calls += 1 + + +def test_close_releases_client_once() -> None: + service = TOSService.__new__(TOSService) + client = _FakeClient() + service.client = client + + service.close() + service.close() + + assert client.close_calls == 1 + assert service.client is None + + +def test_context_manager_closes_client() -> None: + service = TOSService.__new__(TOSService) + client = _FakeClient() + service.client = client + + with service as entered: + assert entered is service + + assert client.close_calls == 1 + assert service.client is None