refactor: standardize runtime annotation types across multiple files

This commit is contained in:
Abhimanyu Saharan
2026-02-09 17:24:21 +05:30
parent 7706943209
commit f5d592f61a
47 changed files with 2203 additions and 1413 deletions

View File

@@ -6,6 +6,7 @@ import hashlib
import json
import re
from contextlib import suppress
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4
@@ -88,6 +89,36 @@ MAIN_TEMPLATE_MAP = {
}
@dataclass(frozen=True, slots=True)
class ProvisionOptions:
"""Toggles controlling provisioning write/reset behavior."""
action: str = "provision"
force_bootstrap: bool = False
reset_session: bool = False
@dataclass(frozen=True, slots=True)
class AgentProvisionRequest:
"""Inputs required to provision a board-scoped agent."""
board: Board
gateway: Gateway
auth_token: str
user: User | None
options: ProvisionOptions = field(default_factory=ProvisionOptions)
@dataclass(frozen=True, slots=True)
class MainAgentProvisionRequest:
"""Inputs required to provision a gateway main agent."""
gateway: Gateway
auth_token: str
user: User | None
options: ProvisionOptions = field(default_factory=ProvisionOptions)
def _repo_root() -> Path:
return Path(__file__).resolve().parents[3]
@@ -114,31 +145,48 @@ def _agent_id_from_session_key(session_key: str | None) -> str | None:
return agent_id or None
def _extract_agent_id(payload: object) -> str | None: # noqa: C901
def _from_list(items: object) -> str | None:
if not isinstance(items, list):
return None
for item in items:
if isinstance(item, str) and item.strip():
return item.strip()
if not isinstance(item, dict):
continue
for key in ("id", "agentId", "agent_id"):
raw = item.get(key)
if isinstance(raw, str) and raw.strip():
return raw.strip()
def _clean_str(value: object) -> str | None:
if isinstance(value, str) and value.strip():
return value.strip()
return None
def _extract_agent_id_from_item(item: object) -> str | None:
if isinstance(item, str):
return _clean_str(item)
if not isinstance(item, dict):
return None
for key in ("id", "agentId", "agent_id"):
agent_id = _clean_str(item.get(key))
if agent_id:
return agent_id
return None
def _extract_agent_id_from_list(items: object) -> str | None:
if not isinstance(items, list):
return None
for item in items:
agent_id = _extract_agent_id_from_item(item)
if agent_id:
return agent_id
return None
def _extract_agent_id(payload: object) -> str | None:
default_keys = ("defaultId", "default_id", "defaultAgentId", "default_agent_id")
collection_keys = ("agents", "items", "list", "data")
if isinstance(payload, list):
return _from_list(payload)
return _extract_agent_id_from_list(payload)
if not isinstance(payload, dict):
return None
for key in ("defaultId", "default_id", "defaultAgentId", "default_agent_id"):
raw = payload.get(key)
if isinstance(raw, str) and raw.strip():
return raw.strip()
for key in ("agents", "items", "list", "data"):
agent_id = _from_list(payload.get(key))
for key in default_keys:
agent_id = _clean_str(payload.get(key))
if agent_id:
return agent_id
for key in collection_keys:
agent_id = _extract_agent_id_from_list(payload.get(key))
if agent_id:
return agent_id
return None
@@ -523,42 +571,44 @@ async def _patch_gateway_agent_list(
await openclaw_call("config.patch", params, config=config)
async def patch_gateway_agent_heartbeats( # noqa: C901
gateway: Gateway,
*,
entries: list[tuple[str, str, dict[str, Any]]],
) -> None:
"""Patch multiple agent heartbeat configs in a single gateway config.patch call.
Each entry is (agent_id, workspace_path, heartbeat_dict).
"""
if not gateway.url:
msg = "Gateway url is required"
raise OpenClawGatewayError(msg)
config = GatewayClientConfig(url=gateway.url, token=gateway.token)
async def _gateway_config_agent_list(
config: GatewayClientConfig,
) -> tuple[str | None, list[object]]:
cfg = await openclaw_call("config.get", config=config)
if not isinstance(cfg, dict):
msg = "config.get returned invalid payload"
raise OpenClawGatewayError(msg)
base_hash = cfg.get("hash")
data = cfg.get("config") or cfg.get("parsed") or {}
if not isinstance(data, dict):
msg = "config.get returned invalid config"
raise OpenClawGatewayError(msg)
agents_section = data.get("agents") or {}
lst = agents_section.get("list") or []
if not isinstance(lst, list):
agents_list = agents_section.get("list") or []
if not isinstance(agents_list, list):
msg = "config agents.list is not a list"
raise OpenClawGatewayError(msg)
return cfg.get("hash"), agents_list
entry_by_id: dict[str, tuple[str, dict[str, Any]]] = {
def _heartbeat_entry_map(
entries: list[tuple[str, str, dict[str, Any]]],
) -> dict[str, tuple[str, dict[str, Any]]]:
return {
agent_id: (workspace_path, heartbeat)
for agent_id, workspace_path, heartbeat in entries
}
def _updated_agent_list(
raw_list: list[object],
entry_by_id: dict[str, tuple[str, dict[str, Any]]],
) -> list[object]:
updated_ids: set[str] = set()
new_list: list[dict[str, Any]] = []
for raw_entry in lst:
new_list: list[object] = []
for raw_entry in raw_list:
if not isinstance(raw_entry, dict):
new_list.append(raw_entry)
continue
@@ -566,6 +616,7 @@ async def patch_gateway_agent_heartbeats( # noqa: C901
if not isinstance(agent_id, str) or agent_id not in entry_by_id:
new_list.append(raw_entry)
continue
workspace_path, heartbeat = entry_by_id[agent_id]
new_entry = dict(raw_entry)
new_entry["workspace"] = workspace_path
@@ -580,6 +631,26 @@ async def patch_gateway_agent_heartbeats( # noqa: C901
{"id": agent_id, "workspace": workspace_path, "heartbeat": heartbeat},
)
return new_list
async def patch_gateway_agent_heartbeats(
gateway: Gateway,
*,
entries: list[tuple[str, str, dict[str, Any]]],
) -> None:
"""Patch multiple agent heartbeat configs in a single gateway config.patch call.
Each entry is (agent_id, workspace_path, heartbeat_dict).
"""
if not gateway.url:
msg = "Gateway url is required"
raise OpenClawGatewayError(msg)
config = GatewayClientConfig(url=gateway.url, token=gateway.token)
base_hash, raw_list = await _gateway_config_agent_list(config)
entry_by_id = _heartbeat_entry_map(entries)
new_list = _updated_agent_list(raw_list, entry_by_id)
patch = {"agents": {"list": new_list}}
params = {"raw": json.dumps(patch)}
if base_hash:
@@ -656,18 +727,52 @@ async def _get_gateway_agent_entry(
return None
async def provision_agent( # noqa: C901, PLR0912, PLR0913
agent: Agent,
board: Board,
gateway: Gateway,
auth_token: str,
user: User | None,
def _should_include_bootstrap(
*,
action: str = "provision",
force_bootstrap: bool = False,
reset_session: bool = False,
action: str,
force_bootstrap: bool,
existing_files: dict[str, dict[str, Any]],
) -> bool:
if action != "update" or force_bootstrap:
return True
if not existing_files:
return False
entry = existing_files.get("BOOTSTRAP.md")
return not (entry and entry.get("missing") is True)
async def _set_agent_files(
*,
agent_id: str,
rendered: dict[str, str],
existing_files: dict[str, dict[str, Any]],
client_config: GatewayClientConfig,
) -> None:
for name, content in rendered.items():
if content == "":
continue
if name in PRESERVE_AGENT_EDITABLE_FILES:
entry = existing_files.get(name)
if entry and entry.get("missing") is not True:
continue
try:
await openclaw_call(
"agents.files.set",
{"agentId": agent_id, "name": name, "content": content},
config=client_config,
)
except OpenClawGatewayError as exc:
if "unsupported file" in str(exc).lower():
continue
raise
async def provision_agent(
agent: Agent,
request: AgentProvisionRequest,
) -> None:
"""Provision or update a regular board agent workspace."""
gateway = request.gateway
if not gateway.url:
return
if not gateway.workspace_root:
@@ -682,18 +787,21 @@ async def provision_agent( # noqa: C901, PLR0912, PLR0913
heartbeat = _heartbeat_config(agent)
await _patch_gateway_agent_list(agent_id, workspace_path, heartbeat, client_config)
context = _build_context(agent, board, gateway, auth_token, user)
context = _build_context(
agent,
request.board,
gateway,
request.auth_token,
request.user,
)
supported = set(await _supported_gateway_files(client_config))
supported.update({"USER.md", "SELF.md", "AUTONOMY.md"})
existing_files = await _gateway_agent_files_index(agent_id, client_config)
include_bootstrap = True
if action == "update" and not force_bootstrap:
if not existing_files:
include_bootstrap = False
else:
entry = existing_files.get("BOOTSTRAP.md")
if entry and entry.get("missing") is True:
include_bootstrap = False
include_bootstrap = _should_include_bootstrap(
action=request.options.action,
force_bootstrap=request.options.force_bootstrap,
existing_files=existing_files,
)
rendered = _render_agent_files(
context,
@@ -710,41 +818,22 @@ async def provision_agent( # noqa: C901, PLR0912, PLR0913
with suppress(OSError):
# Local workspace may not be writable/available; fall back to gateway API.
_ensure_workspace_file(workspace_path, name, content, overwrite=False)
for name, content in rendered.items():
if content == "":
continue
if name in PRESERVE_AGENT_EDITABLE_FILES:
# Never overwrite; only provision if missing.
entry = existing_files.get(name)
if entry and entry.get("missing") is not True:
continue
try:
await openclaw_call(
"agents.files.set",
{"agentId": agent_id, "name": name, "content": content},
config=client_config,
)
except OpenClawGatewayError as exc:
# Gateways may restrict file names. Skip unsupported files rather than
# failing provisioning for the entire agent.
if "unsupported file" in str(exc).lower():
continue
raise
if reset_session:
await _set_agent_files(
agent_id=agent_id,
rendered=rendered,
existing_files=existing_files,
client_config=client_config,
)
if request.options.reset_session:
await _reset_session(session_key, client_config)
async def provision_main_agent( # noqa: C901, PLR0912, PLR0913
async def provision_main_agent(
agent: Agent,
gateway: Gateway,
auth_token: str,
user: User | None,
*,
action: str = "provision",
force_bootstrap: bool = False,
reset_session: bool = False,
request: MainAgentProvisionRequest,
) -> None:
"""Provision or update the gateway main agent workspace."""
gateway = request.gateway
if not gateway.url:
return
if not gateway.main_session_key:
@@ -763,18 +852,15 @@ async def provision_main_agent( # noqa: C901, PLR0912, PLR0913
msg = "Unable to resolve gateway main agent id"
raise OpenClawGatewayError(msg)
context = _build_main_context(agent, gateway, auth_token, user)
context = _build_main_context(agent, gateway, request.auth_token, request.user)
supported = set(await _supported_gateway_files(client_config))
supported.update({"USER.md", "SELF.md", "AUTONOMY.md"})
existing_files = await _gateway_agent_files_index(agent_id, client_config)
include_bootstrap = action != "update" or force_bootstrap
if action == "update" and not force_bootstrap:
if not existing_files:
include_bootstrap = False
else:
entry = existing_files.get("BOOTSTRAP.md")
if entry and entry.get("missing") is True:
include_bootstrap = False
include_bootstrap = _should_include_bootstrap(
action=request.options.action,
force_bootstrap=request.options.force_bootstrap,
existing_files=existing_files,
)
rendered = _render_agent_files(
context,
@@ -783,24 +869,13 @@ async def provision_main_agent( # noqa: C901, PLR0912, PLR0913
include_bootstrap=include_bootstrap,
template_overrides=MAIN_TEMPLATE_MAP,
)
for name, content in rendered.items():
if content == "":
continue
if name in PRESERVE_AGENT_EDITABLE_FILES:
entry = existing_files.get(name)
if entry and entry.get("missing") is not True:
continue
try:
await openclaw_call(
"agents.files.set",
{"agentId": agent_id, "name": name, "content": content},
config=client_config,
)
except OpenClawGatewayError as exc:
if "unsupported file" in str(exc).lower():
continue
raise
if reset_session:
await _set_agent_files(
agent_id=agent_id,
rendered=rendered,
existing_files=existing_files,
client_config=client_config,
)
if request.options.reset_session:
await _reset_session(gateway.main_session_key, client_config)

View File

@@ -2,6 +2,7 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from sqlmodel import col, select
@@ -15,7 +16,12 @@ from app.integrations.openclaw_gateway import (
send_message,
)
from app.models.agents import Agent
from app.services.agent_provisioning import DEFAULT_HEARTBEAT_CONFIG, provision_agent
from app.services.agent_provisioning import (
DEFAULT_HEARTBEAT_CONFIG,
AgentProvisionRequest,
ProvisionOptions,
provision_agent,
)
if TYPE_CHECKING:
from sqlmodel.ext.asyncio.session import AsyncSession
@@ -35,18 +41,34 @@ def lead_agent_name(_: Board) -> str:
return "Lead Agent"
async def ensure_board_lead_agent( # noqa: PLR0913
@dataclass(frozen=True, slots=True)
class LeadAgentOptions:
"""Optional overrides for board-lead provisioning behavior."""
agent_name: str | None = None
identity_profile: dict[str, str] | None = None
action: str = "provision"
@dataclass(frozen=True, slots=True)
class LeadAgentRequest:
"""Inputs required to ensure or provision a board lead agent."""
board: Board
gateway: Gateway
config: GatewayClientConfig
user: User | None
options: LeadAgentOptions = field(default_factory=LeadAgentOptions)
async def ensure_board_lead_agent(
session: AsyncSession,
*,
board: Board,
gateway: Gateway,
config: GatewayClientConfig,
user: User | None,
agent_name: str | None = None,
identity_profile: dict[str, str] | None = None,
action: str = "provision",
request: LeadAgentRequest,
) -> tuple[Agent, bool]:
"""Ensure a board has a lead agent; return `(agent, created)`."""
board = request.board
config_options = request.options
existing = (
await session.exec(
select(Agent)
@@ -55,7 +77,7 @@ async def ensure_board_lead_agent( # noqa: PLR0913
)
).first()
if existing:
desired_name = agent_name or lead_agent_name(board)
desired_name = config_options.agent_name or lead_agent_name(board)
changed = False
if existing.name != desired_name:
existing.name = desired_name
@@ -76,17 +98,17 @@ async def ensure_board_lead_agent( # noqa: PLR0913
"communication_style": "direct, concise, practical",
"emoji": ":gear:",
}
if identity_profile:
if config_options.identity_profile:
merged_identity_profile.update(
{
key: value.strip()
for key, value in identity_profile.items()
for key, value in config_options.identity_profile.items()
if value.strip()
},
)
agent = Agent(
name=agent_name or lead_agent_name(board),
name=config_options.agent_name or lead_agent_name(board),
status="provisioning",
board_id=board.id,
is_board_lead=True,
@@ -94,7 +116,7 @@ async def ensure_board_lead_agent( # noqa: PLR0913
identity_profile=merged_identity_profile,
openclaw_session_id=lead_session_key(board),
provision_requested_at=utcnow(),
provision_action=action,
provision_action=config_options.action,
)
raw_token = generate_agent_token()
agent.agent_token_hash = hash_agent_token(raw_token)
@@ -103,11 +125,20 @@ async def ensure_board_lead_agent( # noqa: PLR0913
await session.refresh(agent)
try:
await provision_agent(agent, board, gateway, raw_token, user, action=action)
await provision_agent(
agent,
AgentProvisionRequest(
board=board,
gateway=request.gateway,
auth_token=raw_token,
user=request.user,
options=ProvisionOptions(action=config_options.action),
),
)
if agent.openclaw_session_id:
await ensure_session(
agent.openclaw_session_id,
config=config,
config=request.config,
label=agent.name,
)
await send_message(
@@ -118,7 +149,7 @@ async def ensure_board_lead_agent( # noqa: PLR0913
"then delete it. Begin heartbeats after startup."
),
session_key=agent.openclaw_session_id,
config=config,
config=request.config,
deliver=True,
)
except OpenClawGatewayError:

View File

@@ -2,9 +2,10 @@
from __future__ import annotations
import re
import time
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from html import unescape
from typing import Final
import httpx
@@ -14,6 +15,10 @@ SOULS_DIRECTORY_SITEMAP_URL: Final[str] = f"{SOULS_DIRECTORY_BASE_URL}/sitemap.x
_SITEMAP_TTL_SECONDS: Final[int] = 60 * 60
_SOUL_URL_MIN_PARTS: Final[int] = 6
_LOC_PATTERN: Final[re.Pattern[str]] = re.compile(
r"<(?:[A-Za-z0-9_]+:)?loc>(.*?)</(?:[A-Za-z0-9_]+:)?loc>",
flags=re.IGNORECASE | re.DOTALL,
)
@dataclass(frozen=True, slots=True)
@@ -36,17 +41,10 @@ class SoulRef:
def _parse_sitemap_soul_refs(sitemap_xml: str) -> list[SoulRef]:
"""Parse sitemap XML and extract valid souls.directory handle/slug refs."""
try:
# Souls sitemap is fetched from a known trusted host in this service flow.
root = ET.fromstring(sitemap_xml) # noqa: S314
except ET.ParseError:
return []
# Handle both namespaced and non-namespaced sitemap XML.
# Extract <loc> values without XML entity expansion.
urls = [
loc.text.strip()
for loc in root.iter()
if loc.tag.endswith("loc") and loc.text
unescape(match.group(1)).strip()
for match in _LOC_PATTERN.finditer(sitemap_xml)
]
refs: list[SoulRef] = []

View File

@@ -28,7 +28,13 @@ from app.models.boards import Board
from app.models.gateways import Gateway
from app.models.users import User
from app.schemas.gateways import GatewayTemplatesSyncError, GatewayTemplatesSyncResult
from app.services.agent_provisioning import provision_agent, provision_main_agent
from app.services.agent_provisioning import (
AgentProvisionRequest,
MainAgentProvisionRequest,
ProvisionOptions,
provision_agent,
provision_main_agent,
)
_TOOLS_KV_RE = re.compile(r"^(?P<key>[A-Z0-9_]+)=(?P<value>.*)$")
SESSION_KEY_PARTS_MIN = 2
@@ -480,13 +486,17 @@ async def _sync_one_agent(
async def _do_provision() -> None:
await provision_agent(
agent,
board,
ctx.gateway,
auth_token,
ctx.options.user,
action="update",
force_bootstrap=ctx.options.force_bootstrap,
reset_session=ctx.options.reset_sessions,
AgentProvisionRequest(
board=board,
gateway=ctx.gateway,
auth_token=auth_token,
user=ctx.options.user,
options=ProvisionOptions(
action="update",
force_bootstrap=ctx.options.force_bootstrap,
reset_session=ctx.options.reset_sessions,
),
),
)
await _with_gateway_retry(_do_provision, backoff=ctx.backoff)
@@ -564,12 +574,16 @@ async def _sync_main_agent(
async def _do_provision_main() -> None:
await provision_main_agent(
main_agent,
ctx.gateway,
token,
ctx.options.user,
action="update",
force_bootstrap=ctx.options.force_bootstrap,
reset_session=ctx.options.reset_sessions,
MainAgentProvisionRequest(
gateway=ctx.gateway,
auth_token=token,
user=ctx.options.user,
options=ProvisionOptions(
action="update",
force_bootstrap=ctx.options.force_bootstrap,
reset_session=ctx.options.reset_sessions,
),
),
)
await _with_gateway_retry(_do_provision_main, backoff=ctx.backoff)