From 544eee7814340b098d8ca61720cdb1dd398cc04d Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 07:42:50 +0000 Subject: [PATCH 1/5] Keep the App grant matrix explicit about UC tool identity --- integrations/agentbricks/tests/e2e/README.md | 3 + .../tests/e2e/fixtures/direct_agent.toml | 3 + .../agentbricks/tests/e2e/tool_matrix.py | 4 +- .../unit_tests/tool_matrix_evidence_test.py | 82 +++++++++++-------- 4 files changed, 55 insertions(+), 37 deletions(-) diff --git a/integrations/agentbricks/tests/e2e/README.md b/integrations/agentbricks/tests/e2e/README.md index 7bd8af361..cb324d280 100644 --- a/integrations/agentbricks/tests/e2e/README.md +++ b/integrations/agentbricks/tests/e2e/README.md @@ -116,6 +116,9 @@ Direct authoring does not call `agentbricks tools add`: it replaces `agent.toml` `fixtures/direct_agent.toml`. CLI authoring invokes four managed `agentbricks tools add ...` commands. Both paths then create the same user-owned, framework-native Python tool file with no Python entry in `agent.toml`. Every exact command and code-authoring step is captured in `commands.log`. +Both paths explicitly select App identity for every managed binding, including the UC function, +so deployed grant snapshots and invocations exercise the App service principal. The CLI's default +user identity is covered separately by the request-user auth tests. The CLI path first verifies that an unavailable MCP service is rejected without changing `agent.toml`, then checks that removing the absent binding is harmless. The subsequent dev and diff --git a/integrations/agentbricks/tests/e2e/fixtures/direct_agent.toml b/integrations/agentbricks/tests/e2e/fixtures/direct_agent.toml index afd370531..77fcb86ac 100644 --- a/integrations/agentbricks/tests/e2e/fixtures/direct_agent.toml +++ b/integrations/agentbricks/tests/e2e/fixtures/direct_agent.toml @@ -6,6 +6,7 @@ server = "agentbricks" [[tools]] id = "sandbox" +auth = "app" source = { kind = "sandbox", service = "system.ai.sandbox" } policy = { downscope = [ { resource = "volume:__UC_VOLUME__", permission = "read_only" }, @@ -13,10 +14,12 @@ policy = { downscope = [ [[tools]] id = "web_search" +auth = "app" source = { kind = "mcp", service = "system.ai.web_search" } [[tools]] id = "agentbricks_uc_marker" +auth = "app" source = { kind = "uc_function", function = "__UC_FUNCTION__" } [[tools]] diff --git a/integrations/agentbricks/tests/e2e/tool_matrix.py b/integrations/agentbricks/tests/e2e/tool_matrix.py index 838a04c7a..0663563e0 100644 --- a/integrations/agentbricks/tests/e2e/tool_matrix.py +++ b/integrations/agentbricks/tests/e2e/tool_matrix.py @@ -636,6 +636,8 @@ def _author_cli(self, project: pathlib.Path) -> None: self.uc_function or "", "--name", "agentbricks_uc_marker", + "--auth", + "app", ], [ "tools", @@ -666,7 +668,7 @@ def _author_cli(self, project: pathlib.Path) -> None: f"managed bindings mismatch: expected {sorted(expected)}, got {sorted(tool_ids)}" ) tools_by_id = {tool["id"]: tool for tool in manifest.get("tools", [])} - for tool_id in ("sandbox", "web_search", "genie"): + for tool_id in sorted(expected): if tools_by_id[tool_id].get("auth") != "app": raise MatrixError(f"CLI-authored {tool_id} binding is not App-auth.") diff --git a/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py b/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py index 31956453d..4b914dee4 100644 --- a/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py +++ b/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py @@ -12,8 +12,14 @@ from unittest.mock import Mock import pytest +from click.testing import CliRunner from databricks.sdk.errors import NotFound +from databricks_agentbricks.cli.app import CliContext, agentbricks +from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.agent_project import AgentProject +from databricks_agentbricks.services.deployment.tool_access import plan_tool_access + _MATRIX_PATH = pathlib.Path(__file__).parents[1] / "e2e" / "tool_matrix.py" _SPEC = importlib.util.spec_from_file_location("agentbricks_tool_matrix", _MATRIX_PATH) assert _SPEC is not None and _SPEC.loader is not None @@ -417,7 +423,7 @@ def fail_confirmation(name): ] -def test_cli_authoring_records_app_auth_for_managed_deployed_tools(tmp_path): +def test_cli_authoring_records_app_auth_for_managed_deployed_tools(tmp_path, monkeypatch): project = tmp_path / "project" project.mkdir() manifest = project / "agent.toml" @@ -428,46 +434,44 @@ def test_cli_authoring_records_app_auth_for_managed_deployed_tools(tmp_path): runner.uc_function = "supervisor_agent.mason_agent_tools_e2e.marker" runner.uc_volume = "supervisor_agent.mason_agent_tools_e2e.matrix_volume" runner.genie_space_id = "0" * 32 - bindings = [] + client = Mock() - def fake_run(argv, **kwargs): - if any("missing_service" in str(argument) for argument in argv) and "add" in argv: - return type( - "Result", - (), - { - "returncode": 1, - "stdout": "", - "stderr": '{"error":{"code":"NOT_FOUND"}}', - }, - )() - if "add" in argv: - kind = argv[argv.index("add") + 1] - auth = argv[argv.index("--auth") + 1] if "--auth" in argv else "user" - tool_id = { - "sandbox": "sandbox", - "mcp": "web_search", - "uc-function": "agentbricks_uc_marker", - "genie-agent": "genie", - }[kind] - bindings.append((tool_id, auth)) - manifest.write_text( - 'schema_version = 1\n\n[agent]\nframework = "langgraph"\nserver = "agentbricks"\n' - + "".join( - f'\n[[tools]]\nid = "{binding_id}"\nauth = "{binding_auth}"\n' - for binding_id, binding_auth in bindings - ) - ) - return type("Result", (), {"returncode": 0, "stdout": "", "stderr": ""})() + def get_mcp_service(service): + if service == "system.ai.missing_service": + raise AgentCliError("MCP service does not exist.", error_code="NOT_FOUND") + assert service == "system.ai.web_search" + return {"name": f"mcp-services/{service}"} - runner.run = fake_run + client.get_mcp_service.side_effect = get_mcp_service + monkeypatch.setattr(CliContext, "client", lambda _self: client) + monkeypatch.setenv("AGENTBRICKS_DISABLE_TELEMETRY", "1") + + def run_cli(argv, **kwargs): + result = CliRunner().invoke(agentbricks, argv[1:]) + if kwargs.get("check", True): + assert result.exit_code == 0, result.output + return subprocess.CompletedProcess(argv, result.exit_code, result.stdout, result.stderr) + + runner.run = run_cli runner._author_cli(project) - tools = {tool["id"]: tool for tool in tool_matrix.tomli.loads(manifest.read_text())["tools"]} - assert {tools[tool_id]["auth"] for tool_id in ("sandbox", "web_search", "genie")} == {"app"} - sandbox_command = next(args for args in bindings if args[0] == "sandbox") - assert sandbox_command == ("sandbox", "app") + tools = AgentProject.load(project).tools + assert {tool.id: tool.auth for tool in tools} == { + "sandbox": "app", + "web_search": "app", + "agentbricks_uc_marker": "app", + "genie": "app", + } + assert any( + resource.get("uc_securable") + == { + "securable_full_name": runner.uc_function, + "securable_type": "FUNCTION", + "permission": "EXECUTE", + } + for resource in plan_tool_access(tools).app_resources + ) def test_cli_authoring_includes_temporary_volume_scope(tmp_path): @@ -535,6 +539,12 @@ def test_direct_authoring_includes_temporary_volume_scope(tmp_path): runner._author_direct(project, "langgraph", "abc123") manifest = tool_matrix.tomli.loads((project / "agent.toml").read_text()) + assert {tool["id"]: tool["auth"] for tool in manifest["tools"]} == { + "sandbox": "app", + "web_search": "app", + "agentbricks_uc_marker": "app", + "genie": "app", + } sandbox = next(tool for tool in manifest["tools"] if tool["id"] == "sandbox") assert sandbox["policy"]["downscope"] == [ { From c99efbf3e96300678bbf90d3a3444683be733068 Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 08:27:45 +0000 Subject: [PATCH 2/5] Clean managed Runtime Stores for every owned matrix App --- integrations/agentbricks/tests/e2e/README.md | 3 + .../agentbricks/tests/e2e/tool_matrix.py | 167 +++++---- .../unit_tests/tool_matrix_cleanup_test.py | 337 +++++++++++------- .../unit_tests/tool_matrix_evidence_test.py | 8 +- 4 files changed, 314 insertions(+), 201 deletions(-) diff --git a/integrations/agentbricks/tests/e2e/README.md b/integrations/agentbricks/tests/e2e/README.md index cb324d280..2009e8b1a 100644 --- a/integrations/agentbricks/tests/e2e/README.md +++ b/integrations/agentbricks/tests/e2e/README.md @@ -150,6 +150,9 @@ Lakebase roles, and UC resources are deleted after every run, passing or failing saved in `evidence.json`; App deletion is not considered complete until a follow-up read confirms absence. Pass `--keep-resources` to retain resources while debugging. The gated nightly test reports bounded dev/deploy log tails and captures App runtime logs for failed deployed cases. +Runtime Store cleanup covers every matrix App, including direct authoring without memory/session +bindings. It verifies the store's App owner and service principal before deletion and retains the +App if ownership checks or Runtime Store deletion fail, so cleanup can be retried safely. ## Declarative user-auth scope matrix diff --git a/integrations/agentbricks/tests/e2e/tool_matrix.py b/integrations/agentbricks/tests/e2e/tool_matrix.py index 0663563e0..0217ad658 100644 --- a/integrations/agentbricks/tests/e2e/tool_matrix.py +++ b/integrations/agentbricks/tests/e2e/tool_matrix.py @@ -121,6 +121,11 @@ class ProjectCase: session_store_name: str | None = None +@dataclasses.dataclass(frozen=True) +class RuntimeStoreCleanupTarget: + role_name: str | None + + class Transcript: def __init__(self, path: pathlib.Path): self.path = path @@ -1305,46 +1310,58 @@ def _write_evidence(self) -> None: temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") os.replace(temporary, target) - def _app_role_target(self, app: str) -> str | None: + def _runtime_store_cleanup_target(self, app: str) -> RuntimeStoreCleanupTarget | None: + """Verify the Runtime Store's owner before resolving its App's Lakebase role.""" + client = WorkspaceClient(profile=self.profile) try: - app_record = self.databricks(["apps", "get", app]) - principal = app_record.get("service_principal_client_id") - store = self.databricks(["api", "get", f"/api/2.0/agents/runtime-stores/{app}"]) - owner = store.get("owner", {}).get("app", {}) - branch = store.get("storage_backend", {}).get("lakebase", {}).get("branch") + store = client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app}") + except NotFound: + return None + if not isinstance(store, dict): + raise MatrixError(f"Runtime Store API returned an invalid resource for {app}.") + record = cast(dict[str, Any], store) + principal = client.apps.get(app).service_principal_client_id + owner = record.get("owner", {}).get("app", {}) + branch = record.get("storage_backend", {}).get("lakebase", {}).get("branch") + if ( + not isinstance(principal, str) + or not principal + or record.get("name") != f"runtime-stores/{app}" + or owner.get("name") != app + or owner.get("service_principal_id") != principal + or not isinstance(branch, str) + or not re.fullmatch(r"projects/[^/]+/branches/[^/]+", branch) + ): + raise MatrixError(f"Runtime Store ownership/backend for {app} was not verified.") + roles = cast(list[dict[str, Any]], self.databricks(["postgres", "list-roles", branch])) + targets = [] + for role in roles: + status = role.get("status", {}) + if status.get("postgres_role") != principal: + continue + name = role.get("name") if ( - not principal - or store.get("name") != f"runtime-stores/{app}" - or owner.get("name") != app - or owner.get("service_principal_id") != principal - or not isinstance(branch, str) - or not re.fullmatch(r"projects/[^/]+/branches/[^/]+", branch) + status.get("identity_type") != "SERVICE_PRINCIPAL" + or not isinstance(name, str) + or not re.fullmatch(rf"{re.escape(branch)}/roles/[^/]+", name) ): - self.transcript.write( - f"cleanup warning | Lakebase role for {app} | ownership not verified" - ) - return None - roles = cast(list[dict[str, Any]], self.databricks(["postgres", "list-roles", branch])) - for role in roles: - status = role.get("status", {}) - name = role.get("name") - if ( - status.get("postgres_role") == principal - and status.get("identity_type") == "SERVICE_PRINCIPAL" - and isinstance(name, str) - and name.startswith(f"{branch}/roles/") - ): - return name - except Exception as exc: - self.transcript.write(f"cleanup warning | Lakebase role lookup for {app} | {exc}") - return None + raise MatrixError(f"Lakebase role ownership for {app} was not verified.") + targets.append(name) + if len(targets) > 1: + raise MatrixError(f"Multiple Lakebase roles matched App {app}.") + return RuntimeStoreCleanupTarget(role_name=targets[0] if targets else None) def cleanup(self) -> None: - role_targets = { - case.app_name: self._app_role_target(case.app_name) - for case in self.cases - if case.memory_store_name or case.session_store_name - } + # Every Agent Bricks Runtime deployment gets a managed Runtime Store, including the direct + # fixture without memory/session bindings. Resolve ownership while the App still exists. + runtime_targets: dict[str, RuntimeStoreCleanupTarget | None] = {} + for app in self.apps: + try: + runtime_targets[app] = self._runtime_store_cleanup_target(app) + except Exception as exc: + self.cleanup_results.append( + {"resource": f"runtime-store:{app}", "status": "failed", "detail": str(exc)} + ) apps_with_stores = { case.app_name for case in self.cases @@ -1388,48 +1405,34 @@ def cleanup(self) -> None: if stores_ok: cleaned_stores.add(case.app_name) - apps_with_deleted_runtime_stores: set[str] = set() for app in self.apps: - if app in apps_with_stores: - # The deploy-created Runtime Store owns a dedicated Lakebase database; it must be - # deleted before the role, or the role delete fails on database ownership. - result = self.run( - [ - "databricks", - "api", - "delete", - f"/api/2.0/agents/runtime-stores/{app}", - *self._profile_args(), - ], - timeout=600, - check=False, + if app not in runtime_targets: + continue + runtime_target = runtime_targets[app] + try: + client = WorkspaceClient(profile=self.profile) + # The store owns a dedicated database; delete it before its App and role. Keep the + # App identity on failure so a later cleanup can repeat the ownership check. + if runtime_target is not None: + try: + client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{app}") + except NotFound: + pass + self.cleanup_results.append( + {"resource": f"runtime-store:{app}", "status": "deleted"} ) - output = f"{result.stdout}\n{result.stderr}".strip() - if result.returncode == 0 or "404" in output or "not_found" in output.lower(): - apps_with_deleted_runtime_stores.add(app) - self.cleanup_results.append( - {"resource": f"runtime-store:{app}", "status": "deleted"} - ) - else: - self.cleanup_results.append( - { - "resource": f"runtime-store:{app}", - "status": "failed", - "detail": output, - } - ) - result = self.run( - ["databricks", "apps", "delete", app, *self._profile_args()], - timeout=600, - check=False, - ) - if result.returncode != 0: + except Exception as exc: self.cleanup_results.append( - { - "resource": f"app:{app}", - "status": "failed", - "detail": (result.stderr or result.stdout).strip(), - } + {"resource": f"runtime-store:{app}", "status": "failed", "detail": str(exc)} + ) + continue + try: + client.apps.delete(app) + except NotFound: + pass + except Exception as exc: + self.cleanup_results.append( + {"resource": f"app:{app}", "status": "failed", "detail": str(exc)} ) continue try: @@ -1447,8 +1450,8 @@ def cleanup(self) -> None: } ) self._delete_deployment_source(app) - role_target = role_targets.get(app) - if role_target and app in cleaned_stores and app in apps_with_deleted_runtime_stores: + role_target = runtime_target.role_name if runtime_target is not None else None + if role_target and app in cleaned_stores: try: result = self.run( [ @@ -1461,11 +1464,19 @@ def cleanup(self) -> None: timeout=600, check=False, ) + deleted = result.returncode == 0 + if not deleted: + try: + client.postgres.get_role(role_target) + except NotFound: + deleted = True + except Exception: + pass role_result = { "resource": f"lakebase-role:{role_target}", - "status": "deleted" if result.returncode == 0 else "failed", + "status": "deleted" if deleted else "failed", } - if result.returncode != 0: + if not deleted: role_result["detail"] = (result.stderr or result.stdout).strip() self.cleanup_results.append(role_result) except Exception as exc: diff --git a/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py b/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py index 9330aef51..06116adea 100644 --- a/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py +++ b/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py @@ -55,46 +55,81 @@ def fake_run(argv, **kwargs): assert case.session_store_name == "test-session" -def test_cleanup_deletes_stores_runtime_app_and_matching_lakebase_role( - tmp_path: pathlib.Path, monkeypatch -) -> None: +class _RuntimeCleanupClient: + def __init__(self, app: str, commands: list): + self.app = app + self.commands = commands + self.principal = "test-app-sp" + self.store = { + "name": f"runtime-stores/{app}", + "owner": {"app": {"name": app, "service_principal_id": self.principal}}, + "storage_backend": {"lakebase": {"branch": "projects/test/branches/test-branch"}}, + } + self.get_error: Exception | None = None + self.delete_error: Exception | None = None + self.app_delete_error: Exception | None = None + self.role_get_error: Exception | None = None + self.api_client = type("ApiClient", (), {"do": self._do})() + self.apps = type("Apps", (), {"get": self._get_app, "delete": self._delete_app})() + self.postgres = type("Postgres", (), {"get_role": self._get_role})() + + def _do(self, method, path): + assert path == f"/api/2.0/agents/runtime-stores/{self.app}" + if method == "GET": + if self.get_error: + raise self.get_error + return self.store + assert method == "DELETE" + self.commands.append(["sdk", method, path]) + if self.delete_error: + raise self.delete_error + return {} + + def _get_app(self, app): + assert app == self.app + return type("App", (), {"service_principal_client_id": self.principal})() + + def _delete_app(self, app): + assert app == self.app + self.commands.append(["sdk", "apps", "delete", app]) + if self.app_delete_error: + raise self.app_delete_error + + def _get_role(self, role): + if self.role_get_error: + raise self.role_get_error + return object() + + +def _cleanup_runner(tmp_path, monkeypatch, *, authoring="direct"): wheel = tmp_path / "agentbricks.whl" wheel.write_bytes(b"wheel") runner = Runner(None, tmp_path, wheel) - app = "agent-bricks-t-la-cl-123456" + app = f"agent-bricks-t-la-{authoring[:2]}-123456" runner.apps.append(app) runner.cases.append( ProjectCase( "langgraph", - "cli", + authoring, tmp_path, app, - memory_store_name="memory-stores/test-memory-id", - session_store_name="test-session-store", + memory_store_name="memory-stores/test-memory-id" if authoring == "cli" else None, + session_store_name="test-session-store" if authoring == "cli" else None, ) ) commands = [] + client = _RuntimeCleanupClient(app, commands) + monkeypatch.setitem(Runner.cleanup.__globals__, "WorkspaceClient", lambda profile: client) + role = "projects/test/branches/test-branch/roles/test-role" def fake_databricks(args, **kwargs): - if args[:2] == ["apps", "get"]: - return {"service_principal_client_id": "test-app-sp"} - if args[:2] == ["api", "get"]: - return { - "name": f"runtime-stores/{app}", - "owner": {"app": {"name": app, "service_principal_id": "test-app-sp"}}, - "storage_backend": {"lakebase": {"branch": "projects/test/branches/test-branch"}}, + assert args == ["postgres", "list-roles", "projects/test/branches/test-branch"] + return [ + { + "name": role, + "status": {"postgres_role": client.principal, "identity_type": "SERVICE_PRINCIPAL"}, } - if args[:2] == ["postgres", "list-roles"]: - return [ - { - "name": "projects/test/branches/test-branch/roles/test-role", - "status": { - "postgres_role": "test-app-sp", - "identity_type": "SERVICE_PRINCIPAL", - }, - } - ] - raise AssertionError(args) + ] def fake_run(argv, **kwargs): commands.append(list(argv)) @@ -104,93 +139,81 @@ def fake_run(argv, **kwargs): monkeypatch.setattr(runner, "run", fake_run) monkeypatch.setattr(runner, "_wait_for_app_deleted", lambda name, **kwargs: None) monkeypatch.setattr(runner, "_delete_deployment_source", lambda name: None) + return runner, client, commands, role + + +@pytest.mark.parametrize("authoring", ["cli", "direct"]) +def test_cleanup_deletes_runtime_and_role_with_or_without_bound_stores( + tmp_path: pathlib.Path, monkeypatch, authoring +) -> None: + runner, client, commands, role = _cleanup_runner(tmp_path, monkeypatch, authoring=authoring) runner.cleanup() - assert commands == [ - [ - str(runner.agentbricks), - "memory", - "stores", - "delete", - "memory-stores/test-memory-id", - "--yes", - ], - [str(runner.agentbricks), "sessions", "stores", "delete", "test-session-store", "--yes"], - ["databricks", "api", "delete", f"/api/2.0/agents/runtime-stores/{app}"], - ["databricks", "apps", "delete", app], - [ - "databricks", - "postgres", - "delete-role", - "projects/test/branches/test-branch/roles/test-role", - ], + expected = [] + if authoring == "cli": + expected = [ + [ + str(runner.agentbricks), + "memory", + "stores", + "delete", + "memory-stores/test-memory-id", + "--yes", + ], + [ + str(runner.agentbricks), + "sessions", + "stores", + "delete", + "test-session-store", + "--yes", + ], + ] + expected += [ + ["sdk", "DELETE", f"/api/2.0/agents/runtime-stores/{client.app}"], + ["sdk", "apps", "delete", client.app], + ["databricks", "postgres", "delete-role", role], ] + assert commands == expected + assert { + "resource": f"runtime-store:{client.app}", + "status": "deleted", + } in runner.cleanup_results + assert {"resource": f"lakebase-role:{role}", "status": "deleted"} in runner.cleanup_results + assert runner.cleanup_complete is True -def test_cleanup_keeps_role_when_store_deletion_fails(tmp_path: pathlib.Path, monkeypatch) -> None: - wheel = tmp_path / "agentbricks.whl" - wheel.write_bytes(b"wheel") - runner = Runner(None, tmp_path, wheel) - app = "agent-bricks-t-la-cl-123456" - runner.apps.append(app) - runner.cases.append( - ProjectCase("langgraph", "cli", tmp_path, app, memory_store_name="memory-stores/owned") - ) - monkeypatch.setattr( - runner, "_app_role_target", lambda name: "projects/test/branches/test/roles/sp" - ) - monkeypatch.setattr(runner, "_wait_for_app_deleted", lambda name, **kwargs: None) - monkeypatch.setattr(runner, "_delete_deployment_source", lambda name: None) - commands = [] +def test_cleanup_keeps_role_when_bound_store_deletion_fails( + tmp_path: pathlib.Path, monkeypatch +) -> None: + runner, client, commands, _ = _cleanup_runner(tmp_path, monkeypatch, authoring="cli") def fake_run(argv, **kwargs): commands.append(list(argv)) return subprocess.CompletedProcess(argv, 1 if "memory" in argv else 0, "", "") monkeypatch.setattr(runner, "run", fake_run) - runner.cleanup() - assert ["databricks", "apps", "delete", app] in commands + assert ["sdk", "apps", "delete", client.app] in commands assert not any("delete-role" in command for command in commands) + assert runner.cleanup_complete is False -def test_cleanup_keeps_role_when_runtime_store_deletion_fails( +def test_cleanup_retains_app_identity_when_runtime_store_deletion_fails( tmp_path: pathlib.Path, monkeypatch ) -> None: - wheel = tmp_path / "agentbricks.whl" - wheel.write_bytes(b"wheel") - runner = Runner(None, tmp_path, wheel) - app = "agent-bricks-t-la-cl-123456" - runner.apps.append(app) - runner.cases.append( - ProjectCase("langgraph", "cli", tmp_path, app, memory_store_name="memory-stores/owned") - ) - monkeypatch.setattr( - runner, "_app_role_target", lambda name: "projects/test/branches/test/roles/sp" - ) - monkeypatch.setattr(runner, "_wait_for_app_deleted", lambda name, **kwargs: None) - monkeypatch.setattr(runner, "_delete_deployment_source", lambda name: None) - commands = [] - - def fake_run(argv, **kwargs): - commands.append(list(argv)) - failed = subprocess.CompletedProcess(argv, 1, "", "409 Conflict: resource is in use") - if argv[:2] == ["databricks", "api"]: - return failed - return subprocess.CompletedProcess(argv, 0, "", "") - - monkeypatch.setattr(runner, "run", fake_run) + runner, client, commands, _ = _cleanup_runner(tmp_path, monkeypatch) + client.delete_error = RuntimeError("resource is in use") runner.cleanup() - assert ["databricks", "apps", "delete", app] in commands - assert not any("delete-role" in command for command in commands) + assert commands == [["sdk", "DELETE", f"/api/2.0/agents/runtime-stores/{client.app}"]] assert { - "resource": f"runtime-store:{app}", + "resource": f"runtime-store:{client.app}", "status": "failed", - "detail": "409 Conflict: resource is in use", + "detail": "resource is in use", } in runner.cleanup_results assert runner.cleanup_complete is False @@ -198,53 +221,125 @@ def fake_run(argv, **kwargs): def test_cleanup_treats_missing_runtime_store_as_deleted( tmp_path: pathlib.Path, monkeypatch ) -> None: - wheel = tmp_path / "agentbricks.whl" - wheel.write_bytes(b"wheel") - runner = Runner(None, tmp_path, wheel) - app = "agent-bricks-t-la-cl-123456" - runner.apps.append(app) - runner.cases.append( - ProjectCase("langgraph", "cli", tmp_path, app, memory_store_name="memory-stores/owned") - ) - monkeypatch.setattr( - runner, "_app_role_target", lambda name: "projects/test/branches/test/roles/sp" + runner, client, commands, _ = _cleanup_runner(tmp_path, monkeypatch) + client.get_error = _NAMESPACE["NotFound"]("missing") + client.delete_error = _NAMESPACE["NotFound"]("missing") + + runner.cleanup() + + assert commands == [["sdk", "apps", "delete", client.app]] + assert { + "resource": f"runtime-store:{client.app}", + "status": "deleted", + } in runner.cleanup_results + assert runner.cleanup_complete is True + + +def test_cleanup_handles_runtime_store_removed_after_ownership_check( + tmp_path: pathlib.Path, monkeypatch +) -> None: + runner, client, commands, role = _cleanup_runner(tmp_path, monkeypatch) + client.delete_error = _NAMESPACE["NotFound"]("missing") + + runner.cleanup() + + assert commands == [ + ["sdk", "DELETE", f"/api/2.0/agents/runtime-stores/{client.app}"], + ["sdk", "apps", "delete", client.app], + ["databricks", "postgres", "delete-role", role], + ] + assert runner.cleanup_complete is True + + +def test_cleanup_handles_app_already_absent(tmp_path: pathlib.Path, monkeypatch) -> None: + runner, client, _, _ = _cleanup_runner(tmp_path, monkeypatch) + client.get_error = _NAMESPACE["NotFound"]("missing store") + client.app_delete_error = _NAMESPACE["NotFound"]("missing App") + + runner.cleanup() + + assert runner.cleanup_complete is True + assert { + "resource": f"runtime-store:{client.app}", + "status": "deleted", + } in runner.cleanup_results + assert any( + row["resource"] == f"app:{client.app}" and row["status"] == "deleted" + for row in runner.cleanup_results ) - monkeypatch.setattr(runner, "_wait_for_app_deleted", lambda name, **kwargs: None) - monkeypatch.setattr(runner, "_delete_deployment_source", lambda name: None) - commands = [] + + +def test_cleanup_handles_role_removed_after_ownership_check( + tmp_path: pathlib.Path, monkeypatch +) -> None: + runner, client, commands, role = _cleanup_runner(tmp_path, monkeypatch) + client.role_get_error = _NAMESPACE["NotFound"]("missing role") def fake_run(argv, **kwargs): commands.append(list(argv)) - if argv[:2] == ["databricks", "api"]: - return subprocess.CompletedProcess(argv, 1, "", "Error: 404 Not Found") - return subprocess.CompletedProcess(argv, 0, "", "") + return subprocess.CompletedProcess(argv, 1, "", "role delete failed") monkeypatch.setattr(runner, "run", fake_run) - runner.cleanup() - assert any("delete-role" in command for command in commands) - assert {"resource": f"runtime-store:{app}", "status": "deleted"} in runner.cleanup_results assert runner.cleanup_complete is True + assert {"resource": f"lakebase-role:{role}", "status": "deleted"} in runner.cleanup_results -def test_role_lookup_rejects_other_app_owner(tmp_path: pathlib.Path, monkeypatch) -> None: - wheel = tmp_path / "agentbricks.whl" - wheel.write_bytes(b"wheel") - runner = Runner(None, tmp_path, wheel) +@pytest.mark.parametrize("field", ["resource_name", "owner_name", "owner_principal", "branch"]) +def test_cleanup_retains_app_and_store_when_runtime_ownership_is_not_verified( + tmp_path: pathlib.Path, monkeypatch, field +) -> None: + runner, client, commands, _ = _cleanup_runner(tmp_path, monkeypatch) + if field == "resource_name": + client.store["name"] = "runtime-stores/other-app" + elif field == "owner_name": + client.store["owner"]["app"]["name"] = "other-app" + elif field == "owner_principal": + client.store["owner"]["app"]["service_principal_id"] = "other-app-sp" + else: + client.store["storage_backend"]["lakebase"]["branch"] = "projects/other/branches/test/extra" - def fake_databricks(args, **kwargs): - if args[:2] == ["apps", "get"]: - return {"service_principal_client_id": "test-app-sp"} - return { - "name": "runtime-stores/test-app", - "owner": {"app": {"name": "other-app", "service_principal_id": "test-app-sp"}}, - "storage_backend": {"lakebase": {"branch": "projects/test/branches/test"}}, - } + runner.cleanup() - monkeypatch.setattr(runner, "databricks", fake_databricks) + assert commands == [] + assert runner.cleanup_complete is False + assert runner.cleanup_results[0]["resource"] == f"runtime-store:{client.app}" + assert runner.cleanup_results[0]["status"] == "failed" - assert runner._app_role_target("test-app") is None + +def test_cleanup_does_not_treat_runtime_lookup_failure_as_absence( + tmp_path: pathlib.Path, monkeypatch +) -> None: + runner, client, commands, _ = _cleanup_runner(tmp_path, monkeypatch) + client.get_error = _NAMESPACE["DatabricksError"]("permission denied") + + runner.cleanup() + + assert commands == [] + assert runner.cleanup_complete is False + + +@pytest.mark.parametrize("bad_role", ["other_branch", "other_identity_type", "duplicate"]) +def test_cleanup_retains_app_when_role_ownership_is_ambiguous( + tmp_path: pathlib.Path, monkeypatch, bad_role +) -> None: + runner, client, commands, role = _cleanup_runner(tmp_path, monkeypatch) + record = { + "name": role if bad_role != "other_branch" else "projects/other/branches/test/roles/sp", + "status": { + "postgres_role": client.principal, + "identity_type": "USER" if bad_role == "other_identity_type" else "SERVICE_PRINCIPAL", + }, + } + monkeypatch.setattr( + runner, "databricks", lambda args: [record] * (2 if bad_role == "duplicate" else 1) + ) + + runner.cleanup() + + assert commands == [] + assert runner.cleanup_complete is False class _FakeWorkspaceClient: diff --git a/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py b/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py index 4b914dee4..f74edb92e 100644 --- a/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py +++ b/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py @@ -397,7 +397,7 @@ def get(self, name): assert calls == 2 -def test_cleanup_records_failed_when_app_absence_is_not_confirmed(tmp_path): +def test_cleanup_records_failed_when_app_absence_is_not_confirmed(tmp_path, monkeypatch): wheel = tmp_path / "wheel.whl" wheel.write_bytes(b"wheel") runner = tool_matrix.Runner("profile", tmp_path / "out", wheel) @@ -405,21 +405,25 @@ def test_cleanup_records_failed_when_app_absence_is_not_confirmed(tmp_path): runner.run = lambda *args, **kwargs: type( "Result", (), {"returncode": 0, "stdout": "", "stderr": ""} )() + runner._runtime_store_cleanup_target = lambda name: None + client = type("Client", (), {"apps": type("Apps", (), {"delete": lambda _, name: None})()})() def fail_confirmation(name): raise tool_matrix.MatrixError("App still exists") runner._wait_for_app_deleted = fail_confirmation + monkeypatch.setattr(tool_matrix, "WorkspaceClient", lambda profile: client) runner.cleanup() assert runner.cleanup_complete is False assert runner.cleanup_results == [ + {"resource": "runtime-store:agent-bricks-test", "status": "deleted"}, { "resource": "app:agent-bricks-test", "status": "failed", "detail": "App still exists", - } + }, ] From fef02d04d73a835eb4af20f7ee95d625cb4e0c4e Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 08:30:28 +0000 Subject: [PATCH 3/5] Preserve App identity and verified role targets across cleanup retries --- integrations/agentbricks/tests/e2e/README.md | 3 +- .../agentbricks/tests/e2e/tool_matrix.py | 93 ++++++++++++------- .../unit_tests/tool_matrix_cleanup_test.py | 69 +++++++++++++- 3 files changed, 125 insertions(+), 40 deletions(-) diff --git a/integrations/agentbricks/tests/e2e/README.md b/integrations/agentbricks/tests/e2e/README.md index 2009e8b1a..8988db5c1 100644 --- a/integrations/agentbricks/tests/e2e/README.md +++ b/integrations/agentbricks/tests/e2e/README.md @@ -152,7 +152,8 @@ absence. Pass `--keep-resources` to retain resources while debugging. The gated captures App runtime logs for failed deployed cases. Runtime Store cleanup covers every matrix App, including direct authoring without memory/session bindings. It verifies the store's App owner and service principal before deletion and retains the -App if ownership checks or Runtime Store deletion fail, so cleanup can be retried safely. +App if ownership checks or bound-store, Runtime Store, or role deletion fail. Verified role targets +remain in runner state and evidence for retries after the Runtime Store has been removed. ## Declarative user-auth scope matrix diff --git a/integrations/agentbricks/tests/e2e/tool_matrix.py b/integrations/agentbricks/tests/e2e/tool_matrix.py index 0217ad658..f1796efed 100644 --- a/integrations/agentbricks/tests/e2e/tool_matrix.py +++ b/integrations/agentbricks/tests/e2e/tool_matrix.py @@ -124,6 +124,8 @@ class ProjectCase: @dataclasses.dataclass(frozen=True) class RuntimeStoreCleanupTarget: role_name: str | None + app_principal: str + store_present: bool = True class Transcript: @@ -188,6 +190,7 @@ def __init__( self.agentbricks = self.runner_venv / "bin" / "agentbricks" self.rows: list[EvidenceRow] = [] self.apps: list[str] = [] + self._runtime_cleanup_targets: dict[str, RuntimeStoreCleanupTarget] = {} self.cases: list[ProjectCase] = [] self.uc_function: str | None = None self.transitive_uc_function: str | None = None @@ -1303,6 +1306,10 @@ def _write_evidence(self) -> None: "cleanup_required": self.cleanup_required, "cleanup_complete": self.cleanup_complete, "cleanup": self.cleanup_results, + "runtime_store_cleanup_targets": [ + {"app_name": app, "role_name": target.role_name} + for app, target in self._runtime_cleanup_targets.items() + ], "rows": [dataclasses.asdict(row) for row in self.rows], } target = self.output / "evidence.json" @@ -1316,7 +1323,14 @@ def _runtime_store_cleanup_target(self, app: str) -> RuntimeStoreCleanupTarget | try: store = client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app}") except NotFound: - return None + previous = self._runtime_cleanup_targets.get(app) + if previous is None: + return None + if client.apps.get(app).service_principal_client_id != previous.app_principal: + raise MatrixError( + f"App identity changed while cleanup for {app} was pending." + ) from None + return dataclasses.replace(previous, store_present=False) if not isinstance(store, dict): raise MatrixError(f"Runtime Store API returned an invalid resource for {app}.") record = cast(dict[str, Any], store) @@ -1349,7 +1363,11 @@ def _runtime_store_cleanup_target(self, app: str) -> RuntimeStoreCleanupTarget | targets.append(name) if len(targets) > 1: raise MatrixError(f"Multiple Lakebase roles matched App {app}.") - return RuntimeStoreCleanupTarget(role_name=targets[0] if targets else None) + target = RuntimeStoreCleanupTarget( + role_name=targets[0] if targets else None, app_principal=principal + ) + self._runtime_cleanup_targets[app] = target + return target def cleanup(self) -> None: # Every Agent Bricks Runtime deployment gets a managed Runtime Store, including the direct @@ -1370,10 +1388,11 @@ def cleanup(self) -> None: cleaned_stores = set(self.apps) - apps_with_stores for case in self.cases: stores_ok = True - for label, name, command in ( - ("memory store", case.memory_store_name, ["memory", "stores", "delete"]), - ("session store", case.session_store_name, ["sessions", "stores", "delete"]), + for label, attribute, command in ( + ("memory store", "memory_store_name", ["memory", "stores", "delete"]), + ("session store", "session_store_name", ["sessions", "stores", "delete"]), ): + name = getattr(case, attribute) if not name: continue try: @@ -1396,6 +1415,7 @@ def cleanup(self) -> None: self.cleanup_results.append( {"resource": resource_name, "status": "deleted"} ) + setattr(case, attribute, None) except Exception as exc: stores_ok = False self.transcript.write(f"cleanup warning | {label} {name} | {exc}") @@ -1406,14 +1426,14 @@ def cleanup(self) -> None: cleaned_stores.add(case.app_name) for app in self.apps: - if app not in runtime_targets: + if app not in runtime_targets or app not in cleaned_stores: continue runtime_target = runtime_targets[app] try: client = WorkspaceClient(profile=self.profile) # The store owns a dedicated database; delete it before its App and role. Keep the # App identity on failure so a later cleanup can repeat the ownership check. - if runtime_target is not None: + if runtime_target is not None and runtime_target.store_present: try: client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{app}") except NotFound: @@ -1426,32 +1446,8 @@ def cleanup(self) -> None: {"resource": f"runtime-store:{app}", "status": "failed", "detail": str(exc)} ) continue - try: - client.apps.delete(app) - except NotFound: - pass - except Exception as exc: - self.cleanup_results.append( - {"resource": f"app:{app}", "status": "failed", "detail": str(exc)} - ) - continue - try: - self._wait_for_app_deleted(app) - except MatrixError as exc: - self.cleanup_results.append( - {"resource": f"app:{app}", "status": "failed", "detail": str(exc)} - ) - continue - self.cleanup_results.append( - { - "resource": f"app:{app}", - "status": "deleted", - "confirmed_absent_at": dt.datetime.now(dt.timezone.utc).isoformat(), - } - ) - self._delete_deployment_source(app) role_target = runtime_target.role_name if runtime_target is not None else None - if role_target and app in cleaned_stores: + if role_target: try: result = self.run( [ @@ -1479,6 +1475,8 @@ def cleanup(self) -> None: if not deleted: role_result["detail"] = (result.stderr or result.stdout).strip() self.cleanup_results.append(role_result) + if not deleted: + continue except Exception as exc: self.transcript.write(f"cleanup warning | Lakebase role {role_target} | {exc}") self.cleanup_results.append( @@ -1488,6 +1486,32 @@ def cleanup(self) -> None: "detail": str(exc), } ) + continue + try: + client.apps.delete(app) + except NotFound: + pass + except Exception as exc: + self.cleanup_results.append( + {"resource": f"app:{app}", "status": "failed", "detail": str(exc)} + ) + continue + try: + self._wait_for_app_deleted(app) + except MatrixError as exc: + self.cleanup_results.append( + {"resource": f"app:{app}", "status": "failed", "detail": str(exc)} + ) + continue + self.cleanup_results.append( + { + "resource": f"app:{app}", + "status": "deleted", + "confirmed_absent_at": dt.datetime.now(dt.timezone.utc).isoformat(), + } + ) + self._delete_deployment_source(app) + self._runtime_cleanup_targets.pop(app, None) if self.uc_function: catalog, schema, function_name = self.uc_function.split(".") try: @@ -1558,9 +1582,8 @@ def cleanup(self) -> None: "detail": str(exc), } ) - self.cleanup_complete = not any( - result.get("status") == "failed" for result in self.cleanup_results - ) + latest_status = {result["resource"]: result["status"] for result in self.cleanup_results} + self.cleanup_complete = "failed" not in latest_status.values() self._write_evidence() def _delete_deployment_source(self, app: str) -> None: diff --git a/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py b/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py index 06116adea..d8338e243 100644 --- a/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py +++ b/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py @@ -172,8 +172,8 @@ def test_cleanup_deletes_runtime_and_role_with_or_without_bound_stores( ] expected += [ ["sdk", "DELETE", f"/api/2.0/agents/runtime-stores/{client.app}"], - ["sdk", "apps", "delete", client.app], ["databricks", "postgres", "delete-role", role], + ["sdk", "apps", "delete", client.app], ] assert commands == expected assert { @@ -196,8 +196,7 @@ def fake_run(argv, **kwargs): monkeypatch.setattr(runner, "run", fake_run) runner.cleanup() - assert ["sdk", "apps", "delete", client.app] in commands - assert not any("delete-role" in command for command in commands) + assert all(command[0] == str(runner.agentbricks) for command in commands) assert runner.cleanup_complete is False @@ -218,6 +217,68 @@ def test_cleanup_retains_app_identity_when_runtime_store_deletion_fails( assert runner.cleanup_complete is False +def test_cleanup_retries_only_the_failed_bound_store_before_app_teardown( + tmp_path: pathlib.Path, monkeypatch +) -> None: + runner, client, commands, _ = _cleanup_runner(tmp_path, monkeypatch, authoring="cli") + memory_attempts = 0 + + def fake_run(argv, **kwargs): + nonlocal memory_attempts + commands.append(list(argv)) + if "memory" in argv: + memory_attempts += 1 + return subprocess.CompletedProcess(argv, 1 if memory_attempts == 1 else 0, "", "") + return subprocess.CompletedProcess(argv, 0, "", "") + + monkeypatch.setattr(runner, "run", fake_run) + runner.cleanup() + assert runner.cleanup_complete is False + assert not any(command[0] == "sdk" for command in commands) + + runner.cleanup() + + assert memory_attempts == 2 + assert sum("sessions" in command for command in commands) == 1 + assert ["sdk", "apps", "delete", client.app] in commands + assert runner.cleanup_complete is True + + +@pytest.mark.parametrize("identity_changed", [False, True]) +def test_cleanup_preserves_verified_role_target_for_retry_after_runtime_store_deletion( + tmp_path: pathlib.Path, monkeypatch, identity_changed +) -> None: + runner, client, commands, role = _cleanup_runner(tmp_path, monkeypatch) + role_attempts = 0 + + def fake_run(argv, **kwargs): + nonlocal role_attempts + commands.append(list(argv)) + role_attempts += 1 + return subprocess.CompletedProcess(argv, 1 if role_attempts == 1 else 0, "", "role busy") + + monkeypatch.setattr(runner, "run", fake_run) + runner.cleanup() + assert runner.cleanup_complete is False + assert not any("apps" in command for command in commands) + assert runner._runtime_cleanup_targets[client.app].role_name == role + client.get_error = _NAMESPACE["NotFound"]("runtime store already removed") + if identity_changed: + client.principal = "replacement-app-sp" + + runner.cleanup() + + if identity_changed: + assert role_attempts == 1 + assert runner.cleanup_complete is False + assert not any("apps" in command for command in commands) + else: + assert role_attempts == 2 + assert sum(command[:2] == ["sdk", "DELETE"] for command in commands) == 1 + assert ["sdk", "apps", "delete", client.app] in commands + assert runner.cleanup_complete is True + + def test_cleanup_treats_missing_runtime_store_as_deleted( tmp_path: pathlib.Path, monkeypatch ) -> None: @@ -245,8 +306,8 @@ def test_cleanup_handles_runtime_store_removed_after_ownership_check( assert commands == [ ["sdk", "DELETE", f"/api/2.0/agents/runtime-stores/{client.app}"], - ["sdk", "apps", "delete", client.app], ["databricks", "postgres", "delete-role", role], + ["sdk", "apps", "delete", client.app], ] assert runner.cleanup_complete is True From 1a766fd74eb0836119305f0e10d966ca24f36062 Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 08:32:38 +0000 Subject: [PATCH 4/5] Verify the latest cleanup attempt while preserving retry history --- .../agentbricks/tests/e2e/tool_matrix.py | 15 ++++++++---- .../unit_tests/tool_matrix_evidence_test.py | 23 +++++++++++++++++++ 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/integrations/agentbricks/tests/e2e/tool_matrix.py b/integrations/agentbricks/tests/e2e/tool_matrix.py index f1796efed..bf90e212f 100644 --- a/integrations/agentbricks/tests/e2e/tool_matrix.py +++ b/integrations/agentbricks/tests/e2e/tool_matrix.py @@ -1973,24 +1973,29 @@ def verify_evidence(path: pathlib.Path, *, require_cleanup: bool = True) -> int: sys.stdout.write("cleanup evidence: cleanup_required was not recorded\n") return 1 cleanup = document.get("cleanup", []) + cleanup_valid = isinstance(cleanup, list) and all( + isinstance(result, dict) and isinstance(result.get("resource"), str) for result in cleanup + ) + # Keep earlier attempts in evidence, but evaluate each resource's latest observation. + latest_cleanup = {result["resource"]: result for result in cleanup} if cleanup_valid else {} if ( require_cleanup and cleanup_required and ( document.get("cleanup_complete") is not True - or not isinstance(cleanup, list) - or any(result.get("status") == "failed" for result in cleanup) + or not cleanup_valid + or any(result.get("status") == "failed" for result in latest_cleanup.values()) or any( result.get("resource", "").startswith("app:") and result.get("status") == "deleted" and not result.get("confirmed_absent_at") - for result in cleanup + for result in latest_cleanup.values() ) ) ): sys.stdout.write("cleanup evidence: required cleanup is incomplete or failed\n") - if isinstance(cleanup, list): - for result in cleanup: + if cleanup_valid: + for result in latest_cleanup.values(): if result.get("status") == "failed": sys.stdout.write( f"cleanup failed | {result.get('resource')} | {result.get('detail')}\n" diff --git a/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py b/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py index f74edb92e..87e9795a9 100644 --- a/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py +++ b/integrations/agentbricks/tests/unit_tests/tool_matrix_evidence_test.py @@ -142,6 +142,29 @@ def test_verify_evidence_rejects_failed_required_cleanup(tmp_path): assert tool_matrix.verify_evidence(path) == 1 +@pytest.mark.parametrize( + ("latest_status", "confirmed", "expected"), + [("deleted", True, 0), ("failed", False, 1), ("deleted", False, 1)], +) +def test_verify_evidence_uses_latest_cleanup_attempt_per_resource( + tmp_path, latest_status, confirmed, expected +): + path = tmp_path / "evidence.json" + evidence = _evidence(cleanup_required=True) + evidence["cleanup"] = [ + {"resource": "app:test", "status": "failed", "detail": "temporary failure"}, + { + "resource": "app:test", + "status": latest_status, + **({"confirmed_absent_at": evidence["ended_at"]} if confirmed else {}), + }, + ] + path.write_text(json.dumps(evidence)) + + assert tool_matrix.verify_evidence(path) == expected + assert json.loads(path.read_text())["cleanup"][0]["status"] == "failed" + + def test_verify_evidence_allows_explicit_keep_resources(tmp_path): path = tmp_path / "evidence.json" evidence = _evidence(cleanup_required=False) From 95051bd5d70a9b2553d0470c693c1eddaec6f0ac Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 15:05:48 +0000 Subject: [PATCH 5/5] Allow safe cleanup retry after App deletion --- .../agentbricks/tests/e2e/tool_matrix.py | 17 +++- .../unit_tests/tool_matrix_cleanup_test.py | 89 +++++++++++++++++++ 2 files changed, 105 insertions(+), 1 deletion(-) diff --git a/integrations/agentbricks/tests/e2e/tool_matrix.py b/integrations/agentbricks/tests/e2e/tool_matrix.py index bf90e212f..3982e99c4 100644 --- a/integrations/agentbricks/tests/e2e/tool_matrix.py +++ b/integrations/agentbricks/tests/e2e/tool_matrix.py @@ -1326,7 +1326,22 @@ def _runtime_store_cleanup_target(self, app: str) -> RuntimeStoreCleanupTarget | previous = self._runtime_cleanup_targets.get(app) if previous is None: return None - if client.apps.get(app).service_principal_client_id != previous.app_principal: + try: + principal = client.apps.get(app).service_principal_client_id + except NotFound: + # App deletion can finish before its absence check fails transiently. Without + # that identity, only an already-absent role is safe to finish cleaning up. + if previous.role_name is not None: + try: + client.postgres.get_role(previous.role_name) + except NotFound: + pass + else: + raise MatrixError( + f"App {app} is absent but its verified Lakebase role still exists." + ) from None + return dataclasses.replace(previous, role_name=None, store_present=False) + if principal != previous.app_principal: raise MatrixError( f"App identity changed while cleanup for {app} was pending." ) from None diff --git a/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py b/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py index d8338e243..a8de2d171 100644 --- a/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py +++ b/integrations/agentbricks/tests/unit_tests/tool_matrix_cleanup_test.py @@ -279,6 +279,95 @@ def fake_run(argv, **kwargs): assert runner.cleanup_complete is True +@pytest.mark.parametrize("role_state", ["absent", "present", "unreadable"]) +def test_cleanup_retry_after_app_deletion_requires_cached_role_absence( + tmp_path: pathlib.Path, monkeypatch, role_state +) -> None: + runner, client, commands, role = _cleanup_runner(tmp_path, monkeypatch) + app_deleted = False + transient_read_pending = True + role_reads = [] + source_deletions = [] + original_api = client.api_client.do + + def api(method, path): + result = original_api(method, path) + if method == "DELETE": + client.get_error = _NAMESPACE["NotFound"]("Runtime Store removed") + return result + + def get_app(app): + nonlocal transient_read_pending + assert app == client.app + if not app_deleted: + return type("App", (), {"service_principal_client_id": client.principal})() + if transient_read_pending: + transient_read_pending = False + raise _NAMESPACE["DatabricksError"]("temporary App read failure") + raise _NAMESPACE["NotFound"]("App removed") + + def delete_app(app): + nonlocal app_deleted + assert app == client.app + commands.append(["sdk", "apps", "delete", app]) + if app_deleted: + raise _NAMESPACE["NotFound"]("App removed") + app_deleted = True + + def get_role(name): + assert name == role + role_reads.append(name) + if role_state == "absent": + raise _NAMESPACE["NotFound"]("role removed") + if role_state == "unreadable": + raise _NAMESPACE["DatabricksError"]("role read denied") + return object() + + monkeypatch.setattr(client.api_client, "do", api) + monkeypatch.setattr(client.apps, "get", get_app) + monkeypatch.setattr(client.apps, "delete", delete_app) + monkeypatch.setattr(client.postgres, "get_role", get_role) + monkeypatch.setattr( + runner, + "_wait_for_app_deleted", + lambda app: Runner._wait_for_app_deleted(runner, app), + ) + monkeypatch.setattr(runner, "_delete_deployment_source", source_deletions.append) + + # The first attempt deletes the resources, but its App-absence read fails transiently. + runner.cleanup() + + assert app_deleted is True + assert runner.cleanup_complete is False + assert source_deletions == [] + assert runner._runtime_cleanup_targets[client.app].role_name == role + first_attempt_commands = list(commands) + assert first_attempt_commands == [ + ["sdk", "DELETE", f"/api/2.0/agents/runtime-stores/{client.app}"], + ["databricks", "postgres", "delete-role", role], + ["sdk", "apps", "delete", client.app], + ] + + # Both App and Runtime Store are now absent. Never delete a role without its App identity. + runner.cleanup() + + assert role_reads == [role] + if role_state == "absent": + assert commands == [*first_attempt_commands, ["sdk", "apps", "delete", client.app]] + assert source_deletions == [client.app] + assert client.app not in runner._runtime_cleanup_targets + assert runner.cleanup_complete is True + assert any( + row["resource"] == f"app:{client.app}" and row.get("confirmed_absent_at") + for row in runner.cleanup_results + ) + else: + assert commands == first_attempt_commands + assert source_deletions == [] + assert runner._runtime_cleanup_targets[client.app].role_name == role + assert runner.cleanup_complete is False + + def test_cleanup_treats_missing_runtime_store_as_deleted( tmp_path: pathlib.Path, monkeypatch ) -> None: