refactor(agent): enhance agent deletion logic and add cleanup for main agents
This commit is contained in:
@@ -23,7 +23,11 @@ from app.db import crud
|
||||
from app.db.pagination import paginate
|
||||
from app.db.session import async_session_maker
|
||||
from app.integrations.openclaw_gateway import GatewayConfig as GatewayClientConfig
|
||||
from app.integrations.openclaw_gateway import OpenClawGatewayError, ensure_session, send_message
|
||||
from app.integrations.openclaw_gateway import (
|
||||
OpenClawGatewayError,
|
||||
ensure_session,
|
||||
send_message,
|
||||
)
|
||||
from app.models.activity_events import ActivityEvent
|
||||
from app.models.agents import Agent
|
||||
from app.models.boards import Board
|
||||
@@ -50,6 +54,7 @@ from app.services.openclaw.provisioning import (
|
||||
MainAgentProvisionRequest,
|
||||
ProvisionOptions,
|
||||
cleanup_agent,
|
||||
cleanup_main_agent,
|
||||
provision_agent,
|
||||
provision_main_agent,
|
||||
)
|
||||
@@ -1330,24 +1335,50 @@ class AgentLifecycleService:
|
||||
return OkResponse()
|
||||
await self.require_agent_access(agent=agent, ctx=ctx, write=True)
|
||||
|
||||
board = await self.require_board(str(agent.board_id) if agent.board_id else None)
|
||||
gateway, client_config = await self.require_gateway(board)
|
||||
try:
|
||||
workspace_path = await cleanup_agent(agent, gateway)
|
||||
except OpenClawGatewayError as exc:
|
||||
self.record_instruction_failure(self.session, agent, str(exc), "delete")
|
||||
await self.session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail=f"Gateway cleanup failed: {exc}",
|
||||
) from exc
|
||||
except (OSError, RuntimeError, ValueError) as exc: # pragma: no cover
|
||||
self.record_instruction_failure(self.session, agent, str(exc), "delete")
|
||||
await self.session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Workspace cleanup failed: {exc}",
|
||||
) from exc
|
||||
gateway: Gateway | None = None
|
||||
client_config: GatewayClientConfig | None = None
|
||||
workspace_path: str | None = None
|
||||
|
||||
if agent.board_id is None:
|
||||
# Gateway-main agents are not tied to a board; resolve via agent.gateway_id.
|
||||
gateway = await Gateway.objects.by_id(agent.gateway_id).first(self.session)
|
||||
if gateway and gateway.url:
|
||||
client_config = GatewayClientConfig(url=gateway.url, token=gateway.token)
|
||||
try:
|
||||
workspace_path = await cleanup_main_agent(agent, gateway)
|
||||
except OpenClawGatewayError as exc:
|
||||
self.record_instruction_failure(self.session, agent, str(exc), "delete")
|
||||
await self.session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail=f"Gateway cleanup failed: {exc}",
|
||||
) from exc
|
||||
except (OSError, RuntimeError, ValueError) as exc: # pragma: no cover
|
||||
self.record_instruction_failure(self.session, agent, str(exc), "delete")
|
||||
await self.session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Workspace cleanup failed: {exc}",
|
||||
) from exc
|
||||
else:
|
||||
board = await self.require_board(str(agent.board_id))
|
||||
gateway, client_config = await self.require_gateway(board)
|
||||
try:
|
||||
workspace_path = await cleanup_agent(agent, gateway)
|
||||
except OpenClawGatewayError as exc:
|
||||
self.record_instruction_failure(self.session, agent, str(exc), "delete")
|
||||
await self.session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail=f"Gateway cleanup failed: {exc}",
|
||||
) from exc
|
||||
except (OSError, RuntimeError, ValueError) as exc: # pragma: no cover
|
||||
self.record_instruction_failure(self.session, agent, str(exc), "delete")
|
||||
await self.session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=f"Workspace cleanup failed: {exc}",
|
||||
) from exc
|
||||
|
||||
record_activity(
|
||||
self.session,
|
||||
@@ -1387,6 +1418,10 @@ class AgentLifecycleService:
|
||||
await self.session.commit()
|
||||
|
||||
try:
|
||||
# Notify the gateway-main agent about cleanup for board-scoped deletes.
|
||||
# Skip when deleting the gateway-main agent itself.
|
||||
if gateway is None or client_config is None or agent.board_id is None:
|
||||
raise ValueError("skip main agent cleanup notification")
|
||||
main_session = GatewayAgentIdentity.session_key(gateway)
|
||||
if main_session and workspace_path:
|
||||
cleanup_message = (
|
||||
|
||||
@@ -941,6 +941,29 @@ async def cleanup_agent(
|
||||
return None
|
||||
|
||||
|
||||
async def cleanup_main_agent(
|
||||
agent: Agent,
|
||||
gateway: Gateway,
|
||||
) -> str | None:
|
||||
"""Remove the gateway-main agent from gateway config and delete its session."""
|
||||
if not gateway.url:
|
||||
return None
|
||||
if not gateway.workspace_root:
|
||||
msg = "gateway_workspace_root is required"
|
||||
raise ValueError(msg)
|
||||
|
||||
workspace_path = _workspace_path(agent, gateway.workspace_root)
|
||||
control_plane = _control_plane_for_gateway(gateway)
|
||||
agent_id = GatewayAgentIdentity.openclaw_agent_id(gateway)
|
||||
await control_plane.delete_agent(agent_id, delete_files=True)
|
||||
|
||||
session_key = (agent.openclaw_session_id or GatewayAgentIdentity.session_key(gateway) or "").strip()
|
||||
if session_key:
|
||||
with suppress(OpenClawGatewayError):
|
||||
await control_plane.delete_agent_session(session_key)
|
||||
return workspace_path
|
||||
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
|
||||
112
backend/tests/test_agent_delete_main_agent.py
Normal file
112
backend/tests/test_agent_delete_main_agent.py
Normal file
@@ -0,0 +1,112 @@
|
||||
# ruff: noqa: S101
|
||||
"""Unit tests for agent deletion behavior."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from types import SimpleNamespace
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
import app.services.openclaw.agent_service as agent_service
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeSession:
|
||||
committed: int = 0
|
||||
deleted: list[object] = field(default_factory=list)
|
||||
|
||||
def add(self, _value: object) -> None:
|
||||
return None
|
||||
|
||||
async def commit(self) -> None:
|
||||
self.committed += 1
|
||||
|
||||
async def delete(self, value: object) -> None:
|
||||
self.deleted.append(value)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _AgentStub:
|
||||
id: UUID
|
||||
name: str
|
||||
gateway_id: UUID
|
||||
board_id: UUID | None = None
|
||||
openclaw_session_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _GatewayStub:
|
||||
id: UUID
|
||||
url: str
|
||||
token: str | None
|
||||
workspace_root: str
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_gateway_main_agent_does_not_require_board_id(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
session = _FakeSession()
|
||||
service = agent_service.AgentLifecycleService(session) # type: ignore[arg-type]
|
||||
|
||||
gateway_id = uuid4()
|
||||
agent = _AgentStub(
|
||||
id=uuid4(),
|
||||
name="Primary Gateway Agent",
|
||||
gateway_id=gateway_id,
|
||||
board_id=None,
|
||||
openclaw_session_id="agent:gateway-x:main",
|
||||
)
|
||||
gateway = _GatewayStub(
|
||||
id=gateway_id,
|
||||
url="ws://gateway.example/ws",
|
||||
token=None,
|
||||
workspace_root="/tmp/openclaw",
|
||||
)
|
||||
ctx = SimpleNamespace(organization=SimpleNamespace(id=uuid4()), member=SimpleNamespace(id=uuid4()))
|
||||
|
||||
async def _fake_first_agent(_session: object) -> _AgentStub:
|
||||
return agent
|
||||
|
||||
async def _fake_first_gateway(_session: object) -> _GatewayStub:
|
||||
return gateway
|
||||
|
||||
monkeypatch.setattr(
|
||||
agent_service.Agent,
|
||||
"objects",
|
||||
SimpleNamespace(by_id=lambda _id: SimpleNamespace(first=_fake_first_agent)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
agent_service.Gateway,
|
||||
"objects",
|
||||
SimpleNamespace(by_id=lambda _id: SimpleNamespace(first=_fake_first_gateway)),
|
||||
)
|
||||
|
||||
async def _no_access_check(*_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
async def _should_not_be_called(*_args, **_kwargs):
|
||||
raise AssertionError("require_board/require_gateway should not be called for main agents")
|
||||
|
||||
called: dict[str, int] = {"cleanup_main": 0}
|
||||
|
||||
async def _fake_cleanup_main_agent(_agent: object, _gateway: object) -> str | None:
|
||||
called["cleanup_main"] += 1
|
||||
return "/tmp/openclaw/workspace-gateway-x"
|
||||
|
||||
async def _fake_update_where(*_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(service, "require_agent_access", _no_access_check)
|
||||
monkeypatch.setattr(service, "require_board", _should_not_be_called)
|
||||
monkeypatch.setattr(service, "require_gateway", _should_not_be_called)
|
||||
monkeypatch.setattr(agent_service, "cleanup_main_agent", _fake_cleanup_main_agent)
|
||||
monkeypatch.setattr(agent_service.crud, "update_where", _fake_update_where)
|
||||
monkeypatch.setattr(agent_service, "record_activity", lambda *_a, **_k: None)
|
||||
|
||||
result = await service.delete_agent(agent_id=str(agent.id), ctx=ctx) # type: ignore[arg-type]
|
||||
|
||||
assert result.ok is True
|
||||
assert called["cleanup_main"] == 1
|
||||
assert session.deleted and session.deleted[0] == agent
|
||||
|
||||
Reference in New Issue
Block a user