From 01133082a429f0fc19e007d7e7f993efd42f5ae5 Mon Sep 17 00:00:00 2001 From: evanlowe <62918515+evanlowe@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:17:21 +0800 Subject: [PATCH 1/2] feat(studio): reuse IAM roles for new agent runtimes --- frontend/README.md | 10 ++ frontend/server/runtime_iam.py | 122 ++++++++++++++ tests/cli/conftest.py | 9 ++ tests/cli/test_agentkit_runtime_iam.py | 90 ----------- tests/cli/test_studio_rbac.py | 86 +++++++--- tests/frontend/server/test_runtime_iam.py | 188 ++++++++++++++++++++++ veadk/cli/agentkit_runtime_iam.py | 91 ----------- veadk/cli/cli_frontend.py | 58 ++----- 8 files changed, 407 insertions(+), 247 deletions(-) create mode 100644 frontend/server/runtime_iam.py delete mode 100644 tests/cli/test_agentkit_runtime_iam.py create mode 100644 tests/frontend/server/test_runtime_iam.py delete mode 100644 veadk/cli/agentkit_runtime_iam.py diff --git a/frontend/README.md b/frontend/README.md index b056e41e5..c56dcc7dc 100644 --- a/frontend/README.md +++ b/frontend/README.md @@ -37,6 +37,16 @@ See [deployment and operation](service/studio_release_notifier/README.md). Studio's ADK chat protocol. Sessions require persistent storage to survive container replacement, and multiple replicas require appropriate session routing. +- **Runtime IAM role reuse**: ordinary and quick Agent creation reuse the first + role with the `AgentKitDefaultRuntimeAccess` system policy in the selected + cloud account, including roles on later IAM result pages. If none matches, + Studio creates `AgentKit_Runtime_Default_ServiceRole_<7 random characters>` + with only that policy. Existing roles keep all their current permissions; + quick creation no longer adds `AgentKitFullAccess`. Lookup errors stop the + deployment instead of triggering role creation. This applies to Volcengine + and BytePlus; existing Runtime updates and Sidecar deployments keep their + existing role behavior + - **Sandbox updates** in System Information compare each Tool's current image with `ListToolTypes` for its cloud provider and actual region. Volcengine and BytePlus use their own credentials and API hosts; catalogs are cached for diff --git a/frontend/server/runtime_iam.py b/frontend/server/runtime_iam.py new file mode 100644 index 000000000..e59fb0b39 --- /dev/null +++ b/frontend/server/runtime_iam.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 + +"""Select IAM roles for new Studio Agent Runtimes""" + +from __future__ import annotations + +import json +import os +from typing import Any + +from veadk.utils.cloud_provider import ( + DEFAULT_CLOUD_PROVIDER, + CloudProvider, + iam_openapi_host, +) + +DEFAULT_RUNTIME_POLICY = "AgentKitDefaultRuntimeAccess" +_ROLE_PAGE_SIZE = 100 + + +def _result(response: dict[str, Any]) -> dict[str, Any]: + error = (response.get("ResponseMetadata") or {}).get("Error") + if error: + raise RuntimeError(error.get("Message") or str(error)) + result = response.get("Result", {}) + if not isinstance(result, dict): + raise RuntimeError("IAM response is missing Result") + return result + + +def _find_reusable_role(iam: Any) -> str | None: + offset = 0 + while True: + page = _result(iam.list_roles({"Limit": _ROLE_PAGE_SIZE, "Offset": offset})) + roles = page.get("RoleMetadata") + total = page.get("Total") + if not isinstance(roles, list) or not isinstance(total, int) or total < 0: + raise RuntimeError("IAM returned an invalid role list") + for role in roles: + name = role.get("RoleName") + if not isinstance(name, str) or not name.strip(): + raise RuntimeError("IAM role is missing RoleName") + policies = _result(iam.list_attached_role_policies({"RoleName": name})).get( + "AttachedPolicyMetadata" + ) + if not isinstance(policies, list): + raise RuntimeError("IAM returned an invalid role policy list") + if any( + policy.get("PolicyName") == DEFAULT_RUNTIME_POLICY + and policy.get("PolicyType") == "System" + for policy in policies + ): + return name + offset += len(roles) + if offset >= total: + return None + if not roles: + raise RuntimeError("IAM returned an incomplete role list") + + +def ensure_runtime_role( + *, + access_key: str, + secret_key: str, + session_token: str | None = None, + provider: CloudProvider = DEFAULT_CLOUD_PROVIDER, +) -> str: + """Reuse a matching role, or create one with only the default runtime policy + + Called under Studio's deployment lock so concurrent local deployments can + reuse the role created by the previous deployment + """ + from volcengine.iam.IamService import IamService + + iam = IamService() + iam.set_ak(access_key) + iam.set_sk(secret_key) + iam.set_host(iam_openapi_host(provider)) + iam.set_scheme("https") + if session_token: + iam.set_session_token(session_token) + + existing = _find_reusable_role(iam) + if existing is not None: + return existing + + from agentkit.utils.misc import generate_runtime_role_name + + name = generate_runtime_role_name() + service_code = ( + os.getenv("VOLCENGINE_AGENTKIT_SERVICE") + or os.getenv("VOLC_AGENTKIT_SERVICE") + or os.getenv("BYTEPLUS_AGENTKIT_SERVICE") + or "" + ).lower() + trust_policy = { + "Statement": [ + { + "Effect": "Allow", + "Action": ["sts:AssumeRole"], + "Principal": { + "Service": ["vefaas_dev" if "stg" in service_code else "vefaas"] + }, + } + ] + } + _result( + iam.create_role( + {"RoleName": name, "TrustPolicyDocument": json.dumps(trust_policy)} + ) + ) + _result( + iam.attach_role_policy( + { + "RoleName": name, + "PolicyName": DEFAULT_RUNTIME_POLICY, + "PolicyType": "System", + } + ) + ) + return name diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py index 54559f6c5..3b930c7b7 100644 --- a/tests/cli/conftest.py +++ b/tests/cli/conftest.py @@ -15,6 +15,15 @@ from __future__ import annotations import pytest +from unittest.mock import MagicMock + + +@pytest.fixture(autouse=True) +def _stub_studio_runtime_role(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Keep deployment tests isolated from live Runtime IAM operations""" + resolver = MagicMock(return_value="shared-runtime-role") + monkeypatch.setattr("frontend.server.runtime_iam.ensure_runtime_role", resolver) + return resolver @pytest.fixture(autouse=True) diff --git a/tests/cli/test_agentkit_runtime_iam.py b/tests/cli/test_agentkit_runtime_iam.py deleted file mode 100644 index b7e6d48ea..000000000 --- a/tests/cli/test_agentkit_runtime_iam.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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. - -import importlib -from unittest.mock import MagicMock - -import pytest - -from veadk.cli.agentkit_runtime_iam import ( - AGENTKIT_RUNTIME_FULL_ACCESS_POLICY, - ensure_quick_runtime_full_access, - is_agentkit_default_runtime_role, -) - - -def _install_iam_service(monkeypatch: pytest.MonkeyPatch, service: MagicMock) -> None: - iam_module = importlib.import_module("volcengine.iam.IamService") - monkeypatch.setattr(iam_module, "IamService", lambda: service) - - -def test_recognizes_current_and_legacy_agentkit_default_roles() -> None: - assert is_agentkit_default_runtime_role( - "AgentKit_Runtime_Default_ServiceRole_abcd123" - ) - assert is_agentkit_default_runtime_role( - "trn:iam::123:role/AgentKit-Runtime-Default-ServiceRole-abcd123" - ) - assert not is_agentkit_default_runtime_role("CustomerRuntimeRole") - - -def test_attaches_full_access_to_generated_runtime_role( - monkeypatch: pytest.MonkeyPatch, -) -> None: - service = MagicMock() - service.list_attached_role_policies.return_value = { - "Result": {"AttachedPolicyMetadata": [{"PolicyName": "AgentKitRuntimeAccess"}]} - } - service.attach_role_policy.return_value = {"Result": {}} - _install_iam_service(monkeypatch, service) - - assert ensure_quick_runtime_full_access( - "AgentKit_Runtime_Default_ServiceRole_abcd123", - access_key="ak", - secret_key="sk", - ) - service.attach_role_policy.assert_called_once_with( - { - "RoleName": "AgentKit_Runtime_Default_ServiceRole_abcd123", - "PolicyName": AGENTKIT_RUNTIME_FULL_ACCESS_POLICY, - "PolicyType": "System", - } - ) - - -def test_keeps_existing_full_access_and_customer_roles_unchanged( - monkeypatch: pytest.MonkeyPatch, -) -> None: - service = MagicMock() - service.list_attached_role_policies.return_value = { - "Result": { - "AttachedPolicyMetadata": [ - {"PolicyName": AGENTKIT_RUNTIME_FULL_ACCESS_POLICY} - ] - } - } - _install_iam_service(monkeypatch, service) - - assert ensure_quick_runtime_full_access( - "AgentKit_Runtime_Default_ServiceRole_abcd123", - access_key="ak", - secret_key="sk", - ) - assert not ensure_quick_runtime_full_access( - "CustomerRuntimeRole", - access_key="ak", - secret_key="sk", - ) - service.attach_role_policy.assert_not_called() - service.list_attached_role_policies.assert_called_once() diff --git a/tests/cli/test_studio_rbac.py b/tests/cli/test_studio_rbac.py index e9ccd8511..ca70145a3 100644 --- a/tests/cli/test_studio_rbac.py +++ b/tests/cli/test_studio_rbac.py @@ -58,6 +58,49 @@ ) +@pytest.mark.parametrize("provider", ["volcengine", "byteplus"]) +def test_runtime_role_lookup_failure_finishes_deployment_and_allows_retry( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + _stub_studio_runtime_role, + provider: str, +) -> None: + monkeypatch.setenv("BYTEPLUS_ACCESS_KEY", "test-ak") + monkeypatch.setenv("BYTEPLUS_SECRET_KEY", "test-sk") + _stub_studio_runtime_role.side_effect = RuntimeError("IAM role lookup denied") + monkeypatch.setattr( + "agentkit.toolkit.sdk.launch", + lambda **_kwargs: pytest.fail("IAM failure must stop deployment before launch"), + ) + app = _create_studio_app( + monkeypatch, tmp_path, developers="developer", provider=provider + ) + with TestClient(app) as client: + for _ in range(2): + response = client.post( + "/web/deploy-agentkit", + headers={"X-VeADK-Local-User": "developer"}, + json={ + "name": "role-lookup-test", + "taskId": "role-lookup-task", + "createEvaluationSets": False, + "envs": [{"key": "MODEL_AGENT_API_KEY", "value": "test-only-key"}], + "files": [{"path": "app.py", "content": "app = object()\n"}], + "config": {"region": "cn-beijing"}, + }, + ) + assert response.status_code == 200 + frames = [ + json.loads(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: ") + ] + assert frames[-1]["done"] is True + assert frames[-1]["success"] is False + assert "IAM role lookup denied" in frames[-1]["error"] + assert _stub_studio_runtime_role.call_count == 2 + + @pytest.mark.parametrize( ("code", "expected"), [ @@ -5452,6 +5495,7 @@ def test_update_deployment_reuses_owned_runtime_and_returns_new_version( has_resource_tags: bool, provider: str, region: str, + _stub_studio_runtime_role, ) -> None: from agentkit.sdk.runtime.client import AgentkitRuntimeClient @@ -5783,6 +5827,7 @@ async def _mark_validated_oauth_token(request: Request, call_next): assert cloud["runtime_id"] == runtime.runtime_id assert cloud["runtime_name"] == runtime.name assert cloud["runtime_role_name"] == "runtime-role" + _stub_studio_runtime_role.assert_not_called() assert cloud["image_tag"] == "veadk-v4" if provider == "volcengine": tencent = "https://mirrors.cloud.tencent.com/pypi/simple" @@ -6135,6 +6180,7 @@ def get_completed_capability(client: TestClient) -> httpx.Response: ("persistent", 1, 5, False, True), ], ) +@pytest.mark.parametrize("provider", ["volcengine", "byteplus"]) def test_new_deployment_only_updates_non_default_instance_range( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -6143,6 +6189,8 @@ def test_new_deployment_only_updates_non_default_instance_range( max_instance: int, expects_update: bool, quick_mode: bool, + provider: str, + _stub_studio_runtime_role, ) -> None: from agentkit.sdk.runtime.client import AgentkitRuntimeClient @@ -6150,7 +6198,6 @@ def test_new_deployment_only_updates_non_default_instance_range( update_requests: list[Any] = [] create_requests: list[Any] = [] captured_config: dict[str, Any] = {} - full_access_calls: list[dict[str, Any]] = [] def create_runtime(_self: Any, request: Any) -> SimpleNamespace: create_requests.append(request) @@ -6188,14 +6235,11 @@ def launch(*, config_file: str, **_kwargs: Any) -> SimpleNamespace: monkeypatch.setattr(AgentkitRuntimeClient, "update_runtime", update_runtime) monkeypatch.setattr(AgentkitRuntimeClient, "get_runtime", get_runtime) monkeypatch.setattr("agentkit.toolkit.sdk.launch", launch) - monkeypatch.setattr( - "veadk.cli.agentkit_runtime_iam.ensure_quick_runtime_full_access", - lambda role_name, **kwargs: full_access_calls.append( - {"role_name": role_name, **kwargs} - ) - or True, + monkeypatch.setenv("BYTEPLUS_ACCESS_KEY", "test-ak") + monkeypatch.setenv("BYTEPLUS_SECRET_KEY", "test-sk") + app = _create_studio_app( + monkeypatch, tmp_path, developers="developer", provider=provider ) - app = _create_studio_app(monkeypatch, tmp_path, developers="developer") with TestClient(app) as client: with client.stream( @@ -6275,21 +6319,15 @@ def launch(*, config_file: str, **_kwargs: Any) -> SimpleNamespace: assert bool(update_requests) is expects_update assert all(request.apmplus_enable is True for request in update_requests) assert any(frame.get("phase") == "update" for frame in frames) is expects_update - assert bool(full_access_calls) is quick_mode - if quick_mode: - assert full_access_calls == [ - { - "role_name": "AgentKit_Runtime_Default_ServiceRole_test", - "access_key": "test-ak", - "secret_key": "test-sk", - "session_token": None, - "provider": "volcengine", - } - ] - assert any( - frame.get("message") == "快速模式 Runtime 已具备 AgentKit 资源访问权限" - for frame in frames - ) + assert captured_config["launch_types"]["cloud"]["runtime_role_name"] == ( + "shared-runtime-role" + ) + _stub_studio_runtime_role.assert_called_once_with( + access_key="test-ak", + secret_key="test-sk", + session_token=None, + provider=provider, + ) if expects_update: request = update_requests[0] assert request.runtime_id == runtime_id @@ -6383,6 +6421,7 @@ def test_new_deployment_rejects_invalid_instance_range( def test_sidecar_deployment_uses_agentkit_cli_structured_release( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, + _stub_studio_runtime_role, ) -> None: from agentkit.sdk.runtime.client import AgentkitRuntimeClient from veadk.cli.studio_sidecar_prerequisites import DEFAULT_SIDECAR_BASE_IMAGE @@ -6616,6 +6655,7 @@ def kill(self) -> None: assert frames[-1]["agentName"] == agent_name assert frames[-1]["runtimeName"] == runtime_name assert captured["command"] == ["/fake/agentkit", "release", "--json"] + _stub_studio_runtime_role.assert_not_called() assert captured["managed_base_in_env"] is True assert captured["create_only"] is True assert captured["cli_env"]["AGENTKIT_RUNTIME_READY_TIMEOUT_MS"] == "900000" diff --git a/tests/frontend/server/test_runtime_iam.py b/tests/frontend/server/test_runtime_iam.py new file mode 100644 index 000000000..9dbb56dec --- /dev/null +++ b/tests/frontend/server/test_runtime_iam.py @@ -0,0 +1,188 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 + +import importlib +import json +import re +from unittest.mock import MagicMock, call + +import pytest + +from frontend.server.runtime_iam import ( + DEFAULT_RUNTIME_POLICY, + ensure_runtime_role, +) + + +@pytest.fixture +def iam(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + service = MagicMock() + service.list_roles.return_value = {"Result": {"RoleMetadata": [], "Total": 0}} + service.list_attached_role_policies.return_value = { + "Result": {"AttachedPolicyMetadata": []} + } + service.create_role.return_value = {"ResponseMetadata": {"Action": "CreateRole"}} + service.attach_role_policy.return_value = { + "ResponseMetadata": {"Action": "AttachRolePolicy"} + } + module = importlib.import_module("volcengine.iam.IamService") + monkeypatch.setattr(module, "IamService", lambda: service) + monkeypatch.delenv("IAM_OPENAPI_HOST", raising=False) + for key in ( + "VOLCENGINE_AGENTKIT_SERVICE", + "VOLC_AGENTKIT_SERVICE", + "BYTEPLUS_AGENTKIT_SERVICE", + ): + monkeypatch.delenv(key, raising=False) + return service + + +@pytest.mark.parametrize( + ("provider", "host"), + [("volcengine", "iam.volcengineapi.com"), ("byteplus", "iam.byteplusapi.com")], +) +def test_reuses_matching_role_on_later_page_without_changing_permissions( + iam: MagicMock, provider, host: str +) -> None: + iam.list_roles.side_effect = [ + {"Result": {"RoleMetadata": [{"RoleName": "old-role"}], "Total": 3}}, + { + "Result": { + "RoleMetadata": [ + {"RoleName": "shared-role"}, + {"RoleName": "another-role"}, + ], + "Total": 3, + } + }, + ] + iam.list_attached_role_policies.side_effect = [ + {"Result": {"AttachedPolicyMetadata": [{"PolicyName": "AgentKitFullAccess"}]}}, + { + "Result": { + "AttachedPolicyMetadata": [ + {"PolicyName": "OtherAccess", "PolicyType": "System"}, + {"PolicyName": DEFAULT_RUNTIME_POLICY, "PolicyType": "System"}, + ] + } + }, + ] + + assert ( + ensure_runtime_role( + access_key="test-ak", + secret_key="test-sk", + session_token="test-token", + provider=provider, + ) + == "shared-role" + ) + + assert iam.list_roles.call_args_list == [ + call({"Limit": 100, "Offset": 0}), + call({"Limit": 100, "Offset": 1}), + ] + assert iam.list_attached_role_policies.call_args_list == [ + call({"RoleName": "old-role"}), + call({"RoleName": "shared-role"}), + ] + iam.set_host.assert_called_once_with(host) + iam.set_ak.assert_called_once_with("test-ak") + iam.set_sk.assert_called_once_with("test-sk") + iam.set_session_token.assert_called_once_with("test-token") + iam.set_scheme.assert_called_once_with("https") + iam.create_role.assert_not_called() + iam.attach_role_policy.assert_not_called() + iam.update_role.assert_not_called() + + +@pytest.mark.parametrize("provider", ["volcengine", "byteplus"]) +@pytest.mark.parametrize("existing_roles", [False, True]) +def test_creates_only_default_policy_when_no_role_matches( + iam: MagicMock, provider, existing_roles: bool +) -> None: + if existing_roles: + iam.list_roles.return_value = { + "Result": {"RoleMetadata": [{"RoleName": "unrelated"}], "Total": 1} + } + iam.list_attached_role_policies.return_value = { + "Result": { + "AttachedPolicyMetadata": [ + {"PolicyName": DEFAULT_RUNTIME_POLICY, "PolicyType": "Custom"}, + {"PolicyName": "AgentKitFullAccess", "PolicyType": "System"}, + ] + } + } + + name = ensure_runtime_role(access_key="ak", secret_key="sk", provider=provider) + + assert re.fullmatch(r"AgentKit_Runtime_Default_ServiceRole_[a-zA-Z0-9]{7}", name) + created = iam.create_role.call_args.args[0] + assert created["RoleName"] == name + assert json.loads(created["TrustPolicyDocument"]) == { + "Statement": [ + { + "Effect": "Allow", + "Action": ["sts:AssumeRole"], + "Principal": {"Service": ["vefaas"]}, + } + ] + } + iam.attach_role_policy.assert_called_once_with( + {"RoleName": name, "PolicyName": DEFAULT_RUNTIME_POLICY, "PolicyType": "System"} + ) + iam.set_session_token.assert_not_called() + + +def test_keeps_staging_trust_service(iam: MagicMock, monkeypatch) -> None: + monkeypatch.setenv("VOLCENGINE_AGENTKIT_SERVICE", "agentkit_stg") + ensure_runtime_role(access_key="ak", secret_key="sk") + document = json.loads(iam.create_role.call_args.args[0]["TrustPolicyDocument"]) + assert document["Statement"][0]["Principal"]["Service"] == ["vefaas_dev"] + + +@pytest.mark.parametrize("operation", ["list_roles", "list_attached_role_policies"]) +def test_lookup_errors_never_create_a_fallback_role( + iam: MagicMock, operation: str +) -> None: + iam.list_roles.return_value = { + "Result": {"RoleMetadata": [{"RoleName": "existing"}], "Total": 1} + } + getattr(iam, operation).return_value = { + "ResponseMetadata": {"Error": {"Code": "AccessDenied", "Message": "denied"}} + } + with pytest.raises(RuntimeError, match="denied"): + ensure_runtime_role(access_key="ak", secret_key="sk") + iam.create_role.assert_not_called() + iam.attach_role_policy.assert_not_called() + + +def test_network_failure_never_creates_a_fallback_role(iam: MagicMock) -> None: + iam.list_roles.side_effect = TimeoutError("IAM timeout") + with pytest.raises(TimeoutError): + ensure_runtime_role(access_key="ak", secret_key="sk") + iam.create_role.assert_not_called() + + +@pytest.mark.parametrize( + "result", + [{}, {"RoleMetadata": [], "Total": 1}, {"RoleMetadata": "invalid", "Total": 0}], +) +def test_incomplete_role_list_never_creates_a_role(iam: MagicMock, result) -> None: + iam.list_roles.return_value = {"Result": result} + with pytest.raises(RuntimeError): + ensure_runtime_role(access_key="ak", secret_key="sk") + iam.create_role.assert_not_called() + + +@pytest.mark.parametrize("operation", ["create_role", "attach_role_policy"]) +def test_creation_or_attachment_failure_is_reported( + iam: MagicMock, operation: str +) -> None: + getattr(iam, operation).return_value = { + "ResponseMetadata": {"Error": {"Message": "IAM operation failed"}} + } + with pytest.raises(RuntimeError, match="IAM operation failed"): + ensure_runtime_role(access_key="ak", secret_key="sk") + if operation == "create_role": + iam.attach_role_policy.assert_not_called() diff --git a/veadk/cli/agentkit_runtime_iam.py b/veadk/cli/agentkit_runtime_iam.py deleted file mode 100644 index dfe8f1586..000000000 --- a/veadk/cli/agentkit_runtime_iam.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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. - -"""IAM policy synchronization for AgentKit-created Runtime roles.""" - -from typing import Any - -from veadk.utils.cloud_provider import ( - DEFAULT_CLOUD_PROVIDER, - CloudProvider, - iam_openapi_host, -) - -AGENTKIT_RUNTIME_FULL_ACCESS_POLICY = "AgentKitFullAccess" -_AGENTKIT_DEFAULT_RUNTIME_ROLE_PREFIXES = ( - "AgentKit_Runtime_Default_ServiceRole", - "AgentKit-Runtime-Default-ServiceRole", -) - - -def _result(response: dict[str, Any]) -> dict[str, Any]: - metadata = response.get("ResponseMetadata", {}) or {} - if metadata.get("Error"): - error = metadata["Error"] - raise RuntimeError(error.get("Message") or str(error)) - return response.get("Result", {}) or {} - - -def is_agentkit_default_runtime_role(role_name: str) -> bool: - """Return whether AgentKit generated the Runtime role automatically.""" - normalized = str(role_name or "").strip().rsplit("/", 1)[-1] - return normalized.startswith(_AGENTKIT_DEFAULT_RUNTIME_ROLE_PREFIXES) - - -def ensure_quick_runtime_full_access( - role_name: str, - *, - access_key: str, - secret_key: str, - session_token: str = "", - provider: CloudProvider = DEFAULT_CLOUD_PROVIDER, -) -> bool: - """Attach AgentKitFullAccess to a quick Runtime's generated default role. - - Customer-managed Runtime roles are never modified. The return value tells - the caller whether the supplied role is an AgentKit-generated default role. - """ - normalized = str(role_name or "").strip().rsplit("/", 1)[-1] - if not is_agentkit_default_runtime_role(normalized): - return False - - from volcengine.iam.IamService import IamService - - iam = IamService() - iam.set_ak(access_key) - iam.set_sk(secret_key) - iam.set_host(iam_openapi_host(provider)) - if provider == "byteplus": - iam.set_scheme("https") - if session_token: - iam.set_session_token(session_token) - - attached = _result(iam.list_attached_role_policies({"RoleName": normalized})).get( - "AttachedPolicyMetadata", [] - ) - if any( - str(policy.get("PolicyName") or "") == AGENTKIT_RUNTIME_FULL_ACCESS_POLICY - for policy in attached - ): - return True - _result( - iam.attach_role_policy( - { - "RoleName": normalized, - "PolicyName": AGENTKIT_RUNTIME_FULL_ACCESS_POLICY, - "PolicyType": "System", - } - ) - ) - return True diff --git a/veadk/cli/cli_frontend.py b/veadk/cli/cli_frontend.py index 91eb97683..8b20cb68a 100644 --- a/veadk/cli/cli_frontend.py +++ b/veadk/cli/cli_frontend.py @@ -6507,14 +6507,6 @@ def _draft_has_model_fallbacks(value: Any) -> bool: ) requested_runtime_name = (data.get("runtimeName") or agent_name).strip() files = data.get("files", []) - quick_mode_requested = ( - isinstance(requested_draft, Mapping) - and requested_draft.get("dynamicAgentDelegation") is True - ) or any( - isinstance(item, Mapping) - and str(item.get("path") or "").endswith("/quick_mode_compat.py") - for item in files - ) migration_task_id = str(data.get("migrationTaskId") or "").strip() source = data.get("source") or ( {"kind": "migration", "migrationId": migration_task_id} @@ -8777,6 +8769,21 @@ def _tagged_update(self, req, _orig=orig_update): try: import copy + if existing_runtime is None and not sidecar_enabled: + from frontend.server.runtime_iam import ensure_runtime_role + + access_key, secret_key, session_token = ( + _resolve_ve_credentials() + ) + sdk_agentkit_config["launch_types"]["cloud"][ + "runtime_role_name" + ] = ensure_runtime_role( + access_key=access_key, + secret_key=secret_key, + session_token=session_token, + provider=provider, + ) + def _launch_config(config: dict[str, Any]): config_path = base / "agentkit.yaml" persisted_config = config @@ -8941,41 +8948,6 @@ def _launch_config(config: dict[str, Any]): 100, ) if result is not None and getattr(result, "success", False): - if quick_mode_requested: - created_runtime_id = str( - task_state.get("runtime_id") or runtime_id - ) - if not created_runtime_id: - raise RuntimeError( - "快速模式 Runtime 创建成功，但未返回 Runtime ID" - ) - runtime_detail = _get_runtime(created_runtime_id, region) - runtime_role_name = str( - getattr(runtime_detail, "role_name", "") or "" - ).strip() - from veadk.cli.agentkit_runtime_iam import ( - ensure_quick_runtime_full_access, - ) - - access_key, secret_key, session_token = ( - _resolve_ve_credentials() - ) - if not ensure_quick_runtime_full_access( - runtime_role_name, - access_key=access_key, - secret_key=secret_key, - session_token=session_token, - provider=provider, - ): - raise RuntimeError( - "快速模式 Runtime 使用了自定义运行角色；" - "请为该角色授予 AgentKitFullAccess。" - ) - _emit( - "success", - "快速模式 Runtime 已具备 AgentKit 资源访问权限", - 100, - ) _verify_sdk_sidecar_release(result) if existing_runtime is not None and provider != "byteplus": try: From a2b049cc8b9ea1081edd4cd8a10553eb12c04ba1 Mon Sep 17 00:00:00 2001 From: evanlowe <62918515+evanlowe@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:19:47 +0800 Subject: [PATCH 2/2] chore(studio): add complete license headers for runtime IAM --- frontend/server/runtime_iam.py | 13 ++++++++++++- tests/frontend/server/test_runtime_iam.py | 13 ++++++++++++- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/frontend/server/runtime_iam.py b/frontend/server/runtime_iam.py index e59fb0b39..af6a30fd2 100644 --- a/frontend/server/runtime_iam.py +++ b/frontend/server/runtime_iam.py @@ -1,5 +1,16 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# SPDX-License-Identifier: Apache-2.0 +# +# 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. """Select IAM roles for new Studio Agent Runtimes""" diff --git a/tests/frontend/server/test_runtime_iam.py b/tests/frontend/server/test_runtime_iam.py index 9dbb56dec..d8f6529cc 100644 --- a/tests/frontend/server/test_runtime_iam.py +++ b/tests/frontend/server/test_runtime_iam.py @@ -1,5 +1,16 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. -# SPDX-License-Identifier: Apache-2.0 +# +# 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 importlib import json