feat(agents): Add agent provisioning confirmation mechanism and update message structure
This commit is contained in:
@@ -0,0 +1,41 @@
|
||||
"""add agent provision confirmation
|
||||
|
||||
Revision ID: 6df47d330227
|
||||
Revises: e0f28e965fa5
|
||||
Create Date: 2026-02-04 17:16:44.472239
|
||||
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from alembic import op
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '6df47d330227'
|
||||
down_revision = 'e0f28e965fa5'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(
|
||||
"ALTER TABLE agents ADD COLUMN IF NOT EXISTS provision_requested_at TIMESTAMP"
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE agents ADD COLUMN IF NOT EXISTS provision_confirm_token_hash VARCHAR"
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE agents ADD COLUMN IF NOT EXISTS provision_action VARCHAR"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute(
|
||||
"ALTER TABLE agents DROP COLUMN IF EXISTS provision_action"
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE agents DROP COLUMN IF EXISTS provision_confirm_token_hash"
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE agents DROP COLUMN IF EXISTS provision_requested_at"
|
||||
)
|
||||
@@ -29,6 +29,7 @@ from app.schemas.agents import (
|
||||
AgentHeartbeatCreate,
|
||||
AgentRead,
|
||||
AgentUpdate,
|
||||
AgentProvisionConfirm,
|
||||
)
|
||||
from app.services.activity_log import record_activity
|
||||
from app.services.agent_provisioning import (
|
||||
@@ -150,6 +151,7 @@ async def _send_wakeup_message(
|
||||
agent: Agent, config: GatewayConfig, verb: str = "provisioned"
|
||||
) -> None:
|
||||
session_key = agent.openclaw_session_id or _build_session_key(agent.name)
|
||||
await ensure_session(session_key, config=config, label=agent.name)
|
||||
message = (
|
||||
f"Hello {agent.name}. Your workspace has been {verb}.\n\n"
|
||||
"Start the agent, run BOOT.md, and if BOOTSTRAP.md exists run it once "
|
||||
@@ -181,6 +183,10 @@ async def create_agent(
|
||||
agent.agent_token_hash = hash_agent_token(raw_token)
|
||||
if agent.heartbeat_config is None:
|
||||
agent.heartbeat_config = DEFAULT_HEARTBEAT_CONFIG.copy()
|
||||
provision_token = generate_agent_token()
|
||||
agent.provision_confirm_token_hash = hash_agent_token(provision_token)
|
||||
agent.provision_requested_at = datetime.utcnow()
|
||||
agent.provision_action = "provision"
|
||||
session_key, session_error = await _ensure_gateway_session(agent.name, config)
|
||||
agent.openclaw_session_id = session_key
|
||||
session.add(agent)
|
||||
@@ -202,21 +208,18 @@ async def create_agent(
|
||||
)
|
||||
session.commit()
|
||||
try:
|
||||
await send_provisioning_message(agent, board, raw_token)
|
||||
await _send_wakeup_message(agent, config)
|
||||
await send_provisioning_message(agent, board, raw_token, provision_token)
|
||||
record_activity(
|
||||
session,
|
||||
event_type="agent.wakeup.sent",
|
||||
message=f"Wakeup message sent to {agent.name}.",
|
||||
event_type="agent.provision.requested",
|
||||
message=f"Provisioning requested for {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
except OpenClawGatewayError as exc:
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
except Exception as exc: # pragma: no cover - unexpected provisioning errors
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
return agent
|
||||
|
||||
@@ -275,33 +278,28 @@ async def update_agent(
|
||||
_record_instruction_failure(session, agent, str(exc), "update")
|
||||
session.commit()
|
||||
raw_token = generate_agent_token()
|
||||
provision_token = generate_agent_token()
|
||||
agent.agent_token_hash = hash_agent_token(raw_token)
|
||||
agent.provision_confirm_token_hash = hash_agent_token(provision_token)
|
||||
agent.provision_requested_at = datetime.utcnow()
|
||||
agent.provision_action = "update"
|
||||
session.add(agent)
|
||||
session.commit()
|
||||
session.refresh(agent)
|
||||
try:
|
||||
await send_update_message(agent, board, raw_token)
|
||||
await _send_wakeup_message(agent, config, verb="updated")
|
||||
await send_update_message(agent, board, raw_token, provision_token)
|
||||
record_activity(
|
||||
session,
|
||||
event_type="agent.updated",
|
||||
message=f"Updated agent {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
record_activity(
|
||||
session,
|
||||
event_type="agent.wakeup.sent",
|
||||
message=f"Wakeup message sent to {agent.name}.",
|
||||
event_type="agent.update.requested",
|
||||
message=f"Update requested for {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
session.commit()
|
||||
except OpenClawGatewayError as exc:
|
||||
_record_instruction_failure(session, agent, str(exc), "update")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
except Exception as exc: # pragma: no cover - unexpected provisioning errors
|
||||
_record_instruction_failure(session, agent, str(exc), "update")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
return _with_computed_status(agent)
|
||||
|
||||
@@ -351,6 +349,10 @@ async def heartbeat_or_create_agent(
|
||||
)
|
||||
raw_token = generate_agent_token()
|
||||
agent.agent_token_hash = hash_agent_token(raw_token)
|
||||
provision_token = generate_agent_token()
|
||||
agent.provision_confirm_token_hash = hash_agent_token(provision_token)
|
||||
agent.provision_requested_at = datetime.utcnow()
|
||||
agent.provision_action = "provision"
|
||||
session_key, session_error = await _ensure_gateway_session(agent.name, config)
|
||||
agent.openclaw_session_id = session_key
|
||||
session.add(agent)
|
||||
@@ -372,21 +374,18 @@ async def heartbeat_or_create_agent(
|
||||
)
|
||||
session.commit()
|
||||
try:
|
||||
await send_provisioning_message(agent, board, raw_token)
|
||||
await _send_wakeup_message(agent, config)
|
||||
await send_provisioning_message(agent, board, raw_token, provision_token)
|
||||
record_activity(
|
||||
session,
|
||||
event_type="agent.wakeup.sent",
|
||||
message=f"Wakeup message sent to {agent.name}.",
|
||||
event_type="agent.provision.requested",
|
||||
message=f"Provisioning requested for {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
except OpenClawGatewayError as exc:
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
except Exception as exc: # pragma: no cover - unexpected provisioning errors
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
elif actor.actor_type == "agent" and actor.agent and actor.agent.id != agent.id:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
@@ -395,27 +394,28 @@ async def heartbeat_or_create_agent(
|
||||
agent.agent_token_hash = hash_agent_token(raw_token)
|
||||
if agent.heartbeat_config is None:
|
||||
agent.heartbeat_config = DEFAULT_HEARTBEAT_CONFIG.copy()
|
||||
provision_token = generate_agent_token()
|
||||
agent.provision_confirm_token_hash = hash_agent_token(provision_token)
|
||||
agent.provision_requested_at = datetime.utcnow()
|
||||
agent.provision_action = "provision"
|
||||
session.add(agent)
|
||||
session.commit()
|
||||
session.refresh(agent)
|
||||
try:
|
||||
board = _require_board(session, str(agent.board_id) if agent.board_id else None)
|
||||
config = _require_gateway_config(board)
|
||||
await send_provisioning_message(agent, board, raw_token)
|
||||
await _send_wakeup_message(agent, config)
|
||||
await send_provisioning_message(agent, board, raw_token, provision_token)
|
||||
record_activity(
|
||||
session,
|
||||
event_type="agent.wakeup.sent",
|
||||
message=f"Wakeup message sent to {agent.name}.",
|
||||
event_type="agent.provision.requested",
|
||||
message=f"Provisioning requested for {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
except OpenClawGatewayError as exc:
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
except Exception as exc: # pragma: no cover - unexpected provisioning errors
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
elif not agent.openclaw_session_id:
|
||||
board = _require_board(session, str(agent.board_id) if agent.board_id else None)
|
||||
@@ -480,6 +480,8 @@ def delete_agent(
|
||||
|
||||
async def _gateway_cleanup_request() -> None:
|
||||
main_session = board.gateway_main_session_key
|
||||
if not main_session:
|
||||
raise OpenClawGatewayError("Board gateway_main_session_key is required")
|
||||
workspace_path = _workspace_path(agent.name, board.gateway_workspace_root)
|
||||
base_url = settings.base_url or "REPLACE_WITH_BASE_URL"
|
||||
cleanup_message = (
|
||||
@@ -519,6 +521,63 @@ def delete_agent(
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@router.post("/{agent_id}/provision/confirm")
|
||||
def confirm_provision_agent(
|
||||
agent_id: str,
|
||||
payload: AgentProvisionConfirm,
|
||||
session: Session = Depends(get_session),
|
||||
) -> dict[str, bool]:
|
||||
agent = session.get(Agent, agent_id)
|
||||
if agent is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
if not agent.provision_confirm_token_hash:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Provisioning confirmation not requested.",
|
||||
)
|
||||
if not verify_agent_token(payload.token, agent.provision_confirm_token_hash):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid token.")
|
||||
if agent.board_id is None:
|
||||
raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)
|
||||
board = _require_board(session, str(agent.board_id))
|
||||
config = _require_gateway_config(board)
|
||||
|
||||
action = payload.action or agent.provision_action or "provision"
|
||||
verb = "updated" if action == "update" else "provisioned"
|
||||
|
||||
try:
|
||||
import asyncio
|
||||
|
||||
asyncio.run(_send_wakeup_message(agent, config, verb=verb))
|
||||
except OpenClawGatewayError as exc:
|
||||
_record_wakeup_failure(session, agent, str(exc))
|
||||
session.commit()
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_502_BAD_GATEWAY,
|
||||
detail=f"Wakeup message failed: {exc}",
|
||||
) from exc
|
||||
|
||||
agent.provision_confirm_token_hash = None
|
||||
agent.provision_requested_at = None
|
||||
agent.provision_action = None
|
||||
agent.updated_at = datetime.utcnow()
|
||||
session.add(agent)
|
||||
record_activity(
|
||||
session,
|
||||
event_type=f"agent.{action}.confirmed",
|
||||
message=f"{action.capitalize()} confirmed for {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
record_activity(
|
||||
session,
|
||||
event_type="agent.wakeup.sent",
|
||||
message=f"Wakeup message sent to {agent.name}.",
|
||||
agent_id=agent.id,
|
||||
)
|
||||
session.commit()
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
@router.post("/{agent_id}/delete/confirm")
|
||||
def confirm_delete_agent(
|
||||
agent_id: str,
|
||||
|
||||
@@ -20,6 +20,9 @@ class Agent(SQLModel, table=True):
|
||||
heartbeat_config: dict[str, Any] | None = Field(
|
||||
default=None, sa_column=Column(JSON)
|
||||
)
|
||||
provision_requested_at: datetime | None = Field(default=None)
|
||||
provision_confirm_token_hash: str | None = Field(default=None, index=True)
|
||||
provision_action: str | None = Field(default=None, index=True)
|
||||
delete_requested_at: datetime | None = Field(default=None)
|
||||
delete_confirm_token_hash: str | None = Field(default=None, index=True)
|
||||
last_seen_at: datetime | None = Field(default=None)
|
||||
|
||||
@@ -44,3 +44,8 @@ class AgentHeartbeatCreate(AgentHeartbeat):
|
||||
|
||||
class AgentDeleteConfirm(SQLModel):
|
||||
token: str
|
||||
|
||||
|
||||
class AgentProvisionConfirm(SQLModel):
|
||||
token: str
|
||||
action: str | None = None
|
||||
|
||||
@@ -125,7 +125,9 @@ def _build_file_blocks(context: dict[str, str]) -> str:
|
||||
)
|
||||
|
||||
|
||||
def build_provisioning_message(agent: Agent, board: Board, auth_token: str) -> str:
|
||||
def build_provisioning_message(
|
||||
agent: Agent, board: Board, auth_token: str, confirm_token: str
|
||||
) -> str:
|
||||
context = _build_context(agent, board, auth_token)
|
||||
file_blocks = _build_file_blocks(context)
|
||||
heartbeat_snippet = json.dumps(
|
||||
@@ -162,12 +164,17 @@ def build_provisioning_message(agent: Agent, board: Board, auth_token: str) -> s
|
||||
f"{heartbeat_snippet}\n"
|
||||
"```\n"
|
||||
"Note: if any agents.list entry defines heartbeat, only those agents "
|
||||
"run heartbeats.\n\n"
|
||||
"run heartbeats.\n"
|
||||
"7) After provisioning completes, confirm by calling:\n"
|
||||
f" POST {context['base_url']}/api/v1/agents/{context['agent_id']}/provision/confirm\n"
|
||||
f" Body: {{\"token\": \"{confirm_token}\", \"action\": \"provision\"}}\n\n"
|
||||
"Files:" + file_blocks
|
||||
)
|
||||
|
||||
|
||||
def build_update_message(agent: Agent, board: Board, auth_token: str) -> str:
|
||||
def build_update_message(
|
||||
agent: Agent, board: Board, auth_token: str, confirm_token: str
|
||||
) -> str:
|
||||
context = _build_context(agent, board, auth_token)
|
||||
file_blocks = _build_file_blocks(context)
|
||||
heartbeat_snippet = json.dumps(
|
||||
@@ -202,7 +209,10 @@ def build_update_message(agent: Agent, board: Board, auth_token: str) -> str:
|
||||
f"{heartbeat_snippet}\n"
|
||||
"```\n"
|
||||
"Note: if any agents.list entry defines heartbeat, only those agents "
|
||||
"run heartbeats.\n\n"
|
||||
"run heartbeats.\n"
|
||||
"7) After the update completes, confirm by calling:\n"
|
||||
f" POST {context['base_url']}/api/v1/agents/{context['agent_id']}/provision/confirm\n"
|
||||
f" Body: {{\"token\": \"{confirm_token}\", \"action\": \"update\"}}\n\n"
|
||||
"Files:" + file_blocks
|
||||
)
|
||||
|
||||
@@ -211,6 +221,7 @@ async def send_provisioning_message(
|
||||
agent: Agent,
|
||||
board: Board,
|
||||
auth_token: str,
|
||||
confirm_token: str,
|
||||
) -> None:
|
||||
if not board.gateway_url:
|
||||
return
|
||||
@@ -219,7 +230,7 @@ async def send_provisioning_message(
|
||||
main_session = board.gateway_main_session_key
|
||||
config = GatewayConfig(url=board.gateway_url, token=board.gateway_token)
|
||||
await ensure_session(main_session, config=config, label="Main Agent")
|
||||
message = build_provisioning_message(agent, board, auth_token)
|
||||
message = build_provisioning_message(agent, board, auth_token, confirm_token)
|
||||
await send_message(message, session_key=main_session, config=config, deliver=False)
|
||||
|
||||
|
||||
@@ -227,6 +238,7 @@ async def send_update_message(
|
||||
agent: Agent,
|
||||
board: Board,
|
||||
auth_token: str,
|
||||
confirm_token: str,
|
||||
) -> None:
|
||||
if not board.gateway_url:
|
||||
return
|
||||
@@ -235,5 +247,5 @@ async def send_update_message(
|
||||
main_session = board.gateway_main_session_key
|
||||
config = GatewayConfig(url=board.gateway_url, token=board.gateway_token)
|
||||
await ensure_session(main_session, config=config, label="Main Agent")
|
||||
message = build_update_message(agent, board, auth_token)
|
||||
message = build_update_message(agent, board, auth_token, confirm_token)
|
||||
await send_message(message, session_key=main_session, config=config, deliver=False)
|
||||
|
||||
Reference in New Issue
Block a user