refactor: enhance docstrings for clarity and consistency across multiple files

backend/app/services/template_sync.py

@@ -1,9 +1,12 @@
+"""Gateway template synchronization orchestration."""
+
 from __future__ import annotations
 
 import asyncio
 import random
 import re
 from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 from typing import TypeVar
 from uuid import UUID, uuid4
 
@@ -11,7 +14,11 @@ from sqlalchemy import func
 from sqlmodel import col, select
 from sqlmodel.ext.asyncio.session import AsyncSession
 
-from app.core.agent_tokens import generate_agent_token, hash_agent_token, verify_agent_token
+from app.core.agent_tokens import (
+    generate_agent_token,
+    hash_agent_token,
+    verify_agent_token,
+)
 from app.core.time import utcnow
 from app.integrations.openclaw_gateway import GatewayConfig as GatewayClientConfig
 from app.integrations.openclaw_gateway import OpenClawGatewayError, openclaw_call
@@ -49,6 +56,31 @@ _TRANSIENT_GATEWAY_ERROR_MARKERS = (
 )
 
 T = TypeVar("T")
+_SECURE_RANDOM = random.SystemRandom()
+_RUNTIME_TYPE_REFERENCES = (Awaitable, Callable, AsyncSession, Gateway, User, UUID)
+
+
+@dataclass(frozen=True)
+class GatewayTemplateSyncOptions:
+    """Runtime options controlling gateway template synchronization."""
+
+    user: User | None
+    include_main: bool = True
+    reset_sessions: bool = False
+    rotate_tokens: bool = False
+    force_bootstrap: bool = False
+    board_id: UUID | None = None
+
+
+@dataclass(frozen=True)
+class _SyncContext:
+    """Shared state passed to sync helper functions."""
+
+    session: AsyncSession
+    gateway: Gateway
+    config: GatewayClientConfig
+    backoff: _GatewayBackoff
+    options: GatewayTemplateSyncOptions
 
 
 def _slugify(value: str) -> str:
@@ -70,7 +102,10 @@ def _is_transient_gateway_error(exc: Exception) -> bool:
 
 
 def _gateway_timeout_message(exc: OpenClawGatewayError) -> str:
-    return f"Gateway unreachable after 10 minutes (template sync timeout). Last error: {exc}"
+    return (
+        "Gateway unreachable after 10 minutes (template sync timeout). "
+        f"Last error: {exc}"
+    )
 
 
 class _GatewayBackoff:
@@ -91,16 +126,25 @@ class _GatewayBackoff:
     def reset(self) -> None:
         self._delay_s = self._base_delay_s
 
+    async def _attempt(
+        self,
+        fn: Callable[[], Awaitable[T]],
+    ) -> tuple[T | None, OpenClawGatewayError | None]:
+        try:
+            return await fn(), None
+        except OpenClawGatewayError as exc:
+            return None, exc
+
     async def run(self, fn: Callable[[], Awaitable[T]]) -> T:
         # Use per-call deadlines so long-running syncs can still tolerate a later
         # gateway restart without having an already-expired retry window.
         deadline_s = asyncio.get_running_loop().time() + self._timeout_s
         while True:
-            try:
-                value = await fn()
-            except OpenClawGatewayError as exc:
+            value, error = await self._attempt(fn)
+            if error is not None:
+                exc = error
                 if not _is_transient_gateway_error(exc):
-                    raise
+                    raise exc
                 now = asyncio.get_running_loop().time()
                 remaining = deadline_s - now
                 if remaining <= 0:
@@ -108,13 +152,16 @@ class _GatewayBackoff:
 
                 sleep_s = min(self._delay_s, remaining)
                 if self._jitter:
-                    sleep_s *= 1.0 + random.uniform(-self._jitter, self._jitter)
+                    sleep_s *= 1.0 + _SECURE_RANDOM.uniform(
+                        -self._jitter,
+                        self._jitter,
+                    )
                 sleep_s = max(0.0, min(sleep_s, remaining))
                 await asyncio.sleep(sleep_s)
                 self._delay_s = min(self._delay_s * 2.0, self._max_delay_s)
-            else:
-                self.reset()
-                return value
+                continue
+            self.reset()
+            return value
 
 
 async def _with_gateway_retry(
@@ -138,23 +185,25 @@ 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:
-    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 _extract_agent_id_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()
+    return None
 
+
+def _extract_agent_id(payload: object) -> str | None:
+    """Extract a default gateway agent id from common list payload shapes."""
     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"):
@@ -162,7 +211,7 @@ def _extract_agent_id(payload: object) -> str | None:
         if isinstance(raw, str) and raw.strip():
             return raw.strip()
     for key in ("agents", "items", "list", "data"):
-        agent_id = _from_list(payload.get(key))
+        agent_id = _extract_agent_id_from_list(payload.get(key))
         if agent_id:
             return agent_id
     return None
@@ -212,9 +261,6 @@ async def _get_agent_file(
     if isinstance(payload, str):
         return payload
     if isinstance(payload, dict):
-        # Common shapes:
-        # - {"name": "...", "content": "..."}
-        # - {"file": {"name": "...", "content": "..." }}
         content = payload.get("content")
         if isinstance(content, str):
             return content
@@ -291,18 +337,53 @@ async def _paused_board_ids(session: AsyncSession, board_ids: list[UUID]) -> set
     return paused
 
 
-async def sync_gateway_templates(
-    session: AsyncSession,
+def _append_sync_error(
+    result: GatewayTemplatesSyncResult,
+    *,
+    message: str,
+    agent: Agent | None = None,
+    board: Board | None = None,
+) -> None:
+    result.errors.append(
+        GatewayTemplatesSyncError(
+            agent_id=agent.id if agent else None,
+            agent_name=agent.name if agent else None,
+            board_id=board.id if board else None,
+            message=message,
+        ),
+    )
+
+
+async def _rotate_agent_token(session: AsyncSession, agent: Agent) -> str:
+    token = generate_agent_token()
+    agent.agent_token_hash = hash_agent_token(token)
+    agent.updated_at = utcnow()
+    session.add(agent)
+    await session.commit()
+    await session.refresh(agent)
+    return token
+
+
+async def _ping_gateway(ctx: _SyncContext, result: GatewayTemplatesSyncResult) -> bool:
+    try:
+        async def _do_ping() -> object:
+            return await openclaw_call("agents.list", config=ctx.config)
+
+        await ctx.backoff.run(_do_ping)
+    except (TimeoutError, OpenClawGatewayError) as exc:
+        _append_sync_error(result, message=str(exc))
+        return False
+    else:
+        return True
+
+
+def _base_result(
     gateway: Gateway,
     *,
-    user: User | None,
-    include_main: bool = True,
-    reset_sessions: bool = False,
-    rotate_tokens: bool = False,
-    force_bootstrap: bool = False,
-    board_id: UUID | None = None,
+    include_main: bool,
+    reset_sessions: bool,
 ) -> GatewayTemplatesSyncResult:
-    result = GatewayTemplatesSyncResult(
+    return GatewayTemplatesSyncResult(
         gateway_id=gateway.id,
         include_main=include_main,
         reset_sessions=reset_sessions,
@@ -310,45 +391,239 @@ async def sync_gateway_templates(
         agents_skipped=0,
         main_updated=False,
     )
+
+
+def _boards_by_id(
+    boards: list[Board],
+    *,
+    board_id: UUID | None,
+) -> dict[UUID, Board] | None:
+    boards_by_id = {board.id: board for board in boards}
+    if board_id is None:
+        return boards_by_id
+    board = boards_by_id.get(board_id)
+    if board is None:
+        return None
+    return {board_id: board}
+
+
+async def _resolve_agent_auth_token(
+    ctx: _SyncContext,
+    result: GatewayTemplatesSyncResult,
+    agent: Agent,
+    board: Board | None,
+    *,
+    agent_gateway_id: str,
+) -> tuple[str | None, bool]:
+    try:
+        auth_token = await _get_existing_auth_token(
+            agent_gateway_id=agent_gateway_id,
+            config=ctx.config,
+            backoff=ctx.backoff,
+        )
+    except TimeoutError as exc:
+        _append_sync_error(result, agent=agent, board=board, message=str(exc))
+        return None, True
+
+    if not auth_token:
+        if not ctx.options.rotate_tokens:
+            result.agents_skipped += 1
+            _append_sync_error(
+                result,
+                agent=agent,
+                board=board,
+                message=(
+                    "Skipping agent: unable to read AUTH_TOKEN from TOOLS.md "
+                    "(run with rotate_tokens=true to re-key)."
+                ),
+            )
+            return None, False
+        auth_token = await _rotate_agent_token(ctx.session, agent)
+
+    if agent.agent_token_hash and not verify_agent_token(
+        auth_token,
+        agent.agent_token_hash,
+    ):
+        if ctx.options.rotate_tokens:
+            auth_token = await _rotate_agent_token(ctx.session, agent)
+        else:
+            _append_sync_error(
+                result,
+                agent=agent,
+                board=board,
+                message=(
+                    "Warning: AUTH_TOKEN in TOOLS.md does not match backend "
+                    "token hash (agent auth may be broken)."
+                ),
+            )
+    return auth_token, False
+
+
+async def _sync_one_agent(
+    ctx: _SyncContext,
+    result: GatewayTemplatesSyncResult,
+    agent: Agent,
+    board: Board,
+) -> bool:
+    auth_token, fatal = await _resolve_agent_auth_token(
+        ctx,
+        result,
+        agent,
+        board,
+        agent_gateway_id=_gateway_agent_id(agent),
+    )
+    if fatal:
+        return True
+    if not auth_token:
+        return False
+    try:
+        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,
+            )
+
+        await _with_gateway_retry(_do_provision, backoff=ctx.backoff)
+        result.agents_updated += 1
+    except TimeoutError as exc:  # pragma: no cover - gateway/network dependent
+        result.agents_skipped += 1
+        _append_sync_error(result, agent=agent, board=board, message=str(exc))
+        return True
+    except (OSError, RuntimeError, ValueError) as exc:  # pragma: no cover
+        result.agents_skipped += 1
+        _append_sync_error(
+            result,
+            agent=agent,
+            board=board,
+            message=f"Failed to sync templates: {exc}",
+        )
+        return False
+    else:
+        return False
+
+
+async def _sync_main_agent(
+    ctx: _SyncContext,
+    result: GatewayTemplatesSyncResult,
+) -> bool:
+    main_agent = (
+        await Agent.objects.all()
+        .filter(col(Agent.openclaw_session_id) == ctx.gateway.main_session_key)
+        .first(ctx.session)
+    )
+    if main_agent is None:
+        _append_sync_error(
+            result,
+            message=(
+                "Gateway main agent record not found; "
+                "skipping main agent template sync."
+            ),
+        )
+        return True
+    try:
+        main_gateway_agent_id = await _gateway_default_agent_id(
+            ctx.config,
+            fallback_session_key=ctx.gateway.main_session_key,
+            backoff=ctx.backoff,
+        )
+    except TimeoutError as exc:
+        _append_sync_error(result, agent=main_agent, message=str(exc))
+        return True
+    if not main_gateway_agent_id:
+        _append_sync_error(
+            result,
+            agent=main_agent,
+            message="Unable to resolve gateway default agent id for main agent.",
+        )
+        return True
+
+    token, fatal = await _resolve_agent_auth_token(
+        ctx,
+        result,
+        main_agent,
+        board=None,
+        agent_gateway_id=main_gateway_agent_id,
+    )
+    if fatal:
+        return True
+    if not token:
+        _append_sync_error(
+            result,
+            agent=main_agent,
+            message="Skipping main agent: unable to read AUTH_TOKEN from TOOLS.md.",
+        )
+        return True
+    stop_sync = False
+    try:
+        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,
+            )
+
+        await _with_gateway_retry(_do_provision_main, backoff=ctx.backoff)
+    except TimeoutError as exc:  # pragma: no cover - gateway/network dependent
+        _append_sync_error(result, agent=main_agent, message=str(exc))
+        stop_sync = True
+    except (OSError, RuntimeError, ValueError) as exc:  # pragma: no cover
+        _append_sync_error(
+            result,
+            agent=main_agent,
+            message=f"Failed to sync main agent templates: {exc}",
+        )
+    else:
+        result.main_updated = True
+    return stop_sync
+
+
+async def sync_gateway_templates(
+    session: AsyncSession,
+    gateway: Gateway,
+    options: GatewayTemplateSyncOptions,
+) -> GatewayTemplatesSyncResult:
+    """Synchronize AGENTS/TOOLS/etc templates to gateway-connected agents."""
+    result = _base_result(
+        gateway,
+        include_main=options.include_main,
+        reset_sessions=options.reset_sessions,
+    )
     if not gateway.url:
-        result.errors.append(
-            GatewayTemplatesSyncError(message="Gateway URL is not configured for this gateway.")
+        _append_sync_error(
+            result,
+            message="Gateway URL is not configured for this gateway.",
         )
         return result
 
-    client_config = GatewayClientConfig(url=gateway.url, token=gateway.token)
-    backoff = _GatewayBackoff(timeout_s=10 * 60)
-
-    # First, wait for the gateway to be reachable (e.g. while it is restarting).
-    try:
-
-        async def _do_ping() -> object:
-            return await openclaw_call("agents.list", config=client_config)
-
-        await backoff.run(_do_ping)
-    except TimeoutError as exc:
-        result.errors.append(GatewayTemplatesSyncError(message=str(exc)))
-        return result
-    except OpenClawGatewayError as exc:
-        result.errors.append(GatewayTemplatesSyncError(message=str(exc)))
+    ctx = _SyncContext(
+        session=session,
+        gateway=gateway,
+        config=GatewayClientConfig(url=gateway.url, token=gateway.token),
+        backoff=_GatewayBackoff(timeout_s=10 * 60),
+        options=options,
+    )
+    if not await _ping_gateway(ctx, result):
         return result
 
     boards = await Board.objects.filter_by(gateway_id=gateway.id).all(session)
-    boards_by_id = {board.id: board for board in boards}
-    if board_id is not None:
-        board = boards_by_id.get(board_id)
-        if board is None:
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    board_id=board_id,
-                    message="Board does not belong to this gateway.",
-                )
-            )
-            return result
-        boards_by_id = {board_id: board}
-
+    boards_by_id = _boards_by_id(boards, board_id=options.board_id)
+    if boards_by_id is None:
+        _append_sync_error(
+            result,
+            message="Board does not belong to this gateway.",
+        )
+        return result
     paused_board_ids = await _paused_board_ids(session, list(boards_by_id.keys()))
-
     if boards_by_id:
         agents = await (
             Agent.objects.by_field_in("board_id", list(boards_by_id.keys()))
@@ -358,251 +633,24 @@ async def sync_gateway_templates(
     else:
         agents = []
 
+    stop_sync = False
     for agent in agents:
         board = boards_by_id.get(agent.board_id) if agent.board_id is not None else None
         if board is None:
             result.agents_skipped += 1
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=agent.id,
-                    agent_name=agent.name,
-                    board_id=agent.board_id,
-                    message="Skipping agent: board not found for agent.",
-                )
+            _append_sync_error(
+                result,
+                agent=agent,
+                message="Skipping agent: board not found for agent.",
             )
             continue
-
         if board.id in paused_board_ids:
             result.agents_skipped += 1
             continue
+        stop_sync = await _sync_one_agent(ctx, result, agent, board)
+        if stop_sync:
+            break
 
-        agent_gateway_id = _gateway_agent_id(agent)
-        try:
-            auth_token = await _get_existing_auth_token(
-                agent_gateway_id=agent_gateway_id,
-                config=client_config,
-                backoff=backoff,
-            )
-        except TimeoutError as exc:
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=agent.id,
-                    agent_name=agent.name,
-                    board_id=board.id,
-                    message=str(exc),
-                )
-            )
-            return result
-
-        if not auth_token:
-            if not rotate_tokens:
-                result.agents_skipped += 1
-                result.errors.append(
-                    GatewayTemplatesSyncError(
-                        agent_id=agent.id,
-                        agent_name=agent.name,
-                        board_id=board.id,
-                        message="Skipping agent: unable to read AUTH_TOKEN from TOOLS.md (run with rotate_tokens=true to re-key).",
-                    )
-                )
-                continue
-            raw_token = generate_agent_token()
-            agent.agent_token_hash = hash_agent_token(raw_token)
-            agent.updated_at = utcnow()
-            session.add(agent)
-            await session.commit()
-            await session.refresh(agent)
-            auth_token = raw_token
-
-        if agent.agent_token_hash and not verify_agent_token(auth_token, agent.agent_token_hash):
-            # Do not block template sync on token drift; optionally re-key.
-            if rotate_tokens:
-                raw_token = generate_agent_token()
-                agent.agent_token_hash = hash_agent_token(raw_token)
-                agent.updated_at = utcnow()
-                session.add(agent)
-                await session.commit()
-                await session.refresh(agent)
-                auth_token = raw_token
-            else:
-                result.errors.append(
-                    GatewayTemplatesSyncError(
-                        agent_id=agent.id,
-                        agent_name=agent.name,
-                        board_id=board.id,
-                        message="Warning: AUTH_TOKEN in TOOLS.md does not match backend token hash (agent auth may be broken).",
-                    )
-                )
-
-        try:
-            agent_item: Agent = agent
-            board_item: Board = board
-            auth_token_value: str = auth_token
-
-            async def _do_provision(
-                agent_item: Agent = agent_item,
-                board_item: Board = board_item,
-                auth_token_value: str = auth_token_value,
-            ) -> None:
-                await provision_agent(
-                    agent_item,
-                    board_item,
-                    gateway,
-                    auth_token_value,
-                    user,
-                    action="update",
-                    force_bootstrap=force_bootstrap,
-                    reset_session=reset_sessions,
-                )
-
-            await _with_gateway_retry(_do_provision, backoff=backoff)
-            result.agents_updated += 1
-        except TimeoutError as exc:  # pragma: no cover - gateway/network dependent
-            result.agents_skipped += 1
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=agent.id,
-                    agent_name=agent.name,
-                    board_id=board.id,
-                    message=str(exc),
-                )
-            )
-            return result
-        except (OSError, RuntimeError, ValueError) as exc:  # pragma: no cover
-            result.agents_skipped += 1
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=agent.id,
-                    agent_name=agent.name,
-                    board_id=board.id,
-                    message=f"Failed to sync templates: {exc}",
-                )
-            )
-
-    if include_main:
-        main_agent = (
-            await Agent.objects.all()
-            .filter(col(Agent.openclaw_session_id) == gateway.main_session_key)
-            .first(session)
-        )
-        if main_agent is None:
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    message="Gateway main agent record not found; skipping main agent template sync.",
-                )
-            )
-            return result
-
-        try:
-            main_gateway_agent_id = await _gateway_default_agent_id(
-                client_config,
-                fallback_session_key=gateway.main_session_key,
-                backoff=backoff,
-            )
-        except TimeoutError as exc:
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=main_agent.id,
-                    agent_name=main_agent.name,
-                    message=str(exc),
-                )
-            )
-            return result
-        if not main_gateway_agent_id:
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=main_agent.id,
-                    agent_name=main_agent.name,
-                    message="Unable to resolve gateway default agent id for main agent.",
-                )
-            )
-            return result
-
-        try:
-            main_token = await _get_existing_auth_token(
-                agent_gateway_id=main_gateway_agent_id,
-                config=client_config,
-                backoff=backoff,
-            )
-        except TimeoutError as exc:
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=main_agent.id,
-                    agent_name=main_agent.name,
-                    message=str(exc),
-                )
-            )
-            return result
-        if not main_token:
-            if rotate_tokens:
-                raw_token = generate_agent_token()
-                main_agent.agent_token_hash = hash_agent_token(raw_token)
-                main_agent.updated_at = utcnow()
-                session.add(main_agent)
-                await session.commit()
-                await session.refresh(main_agent)
-                main_token = raw_token
-            else:
-                result.errors.append(
-                    GatewayTemplatesSyncError(
-                        agent_id=main_agent.id,
-                        agent_name=main_agent.name,
-                        message="Skipping main agent: unable to read AUTH_TOKEN from TOOLS.md.",
-                    )
-                )
-                return result
-
-        if main_agent.agent_token_hash and not verify_agent_token(
-            main_token, main_agent.agent_token_hash
-        ):
-            if rotate_tokens:
-                raw_token = generate_agent_token()
-                main_agent.agent_token_hash = hash_agent_token(raw_token)
-                main_agent.updated_at = utcnow()
-                session.add(main_agent)
-                await session.commit()
-                await session.refresh(main_agent)
-                main_token = raw_token
-            else:
-                result.errors.append(
-                    GatewayTemplatesSyncError(
-                        agent_id=main_agent.id,
-                        agent_name=main_agent.name,
-                        message="Warning: AUTH_TOKEN in TOOLS.md does not match backend token hash (main agent auth may be broken).",
-                    )
-                )
-
-        try:
-
-            async def _do_provision_main() -> None:
-                await provision_main_agent(
-                    main_agent,
-                    gateway,
-                    main_token,
-                    user,
-                    action="update",
-                    force_bootstrap=force_bootstrap,
-                    reset_session=reset_sessions,
-                )
-
-            await _with_gateway_retry(_do_provision_main, backoff=backoff)
-            result.main_updated = True
-        except TimeoutError as exc:  # pragma: no cover - gateway/network dependent
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=main_agent.id,
-                    agent_name=main_agent.name,
-                    message=str(exc),
-                )
-            )
-            return result
-        except (OSError, RuntimeError, ValueError) as exc:  # pragma: no cover
-            result.errors.append(
-                GatewayTemplatesSyncError(
-                    agent_id=main_agent.id,
-                    agent_name=main_agent.name,
-                    message=f"Failed to sync main agent templates: {exc}",
-                )
-            )
-
+    if not stop_sync and options.include_main:
+        await _sync_main_agent(ctx, result)
     return result