refactor: update gateway agent session key handling and improve related logic

This commit is contained in:
Abhimanyu Saharan
2026-02-10 00:45:15 +05:30
parent 79f7ad8ba3
commit ba73ce8bfd
27 changed files with 233 additions and 208 deletions

View File

@@ -16,6 +16,7 @@ from jinja2 import Environment, FileSystemLoader, StrictUndefined, select_autoes
from app.core.config import settings
from app.integrations.openclaw_gateway import GatewayConfig as GatewayClientConfig
from app.integrations.openclaw_gateway import OpenClawGatewayError, ensure_session, openclaw_call
from app.services.gateway_agents import gateway_agent_session_key
if TYPE_CHECKING:
from app.models.agents import Agent
@@ -124,6 +125,7 @@ class MainAgentProvisionRequest:
gateway: Gateway
auth_token: str
user: User | None
session_key: str | None = None
options: ProvisionOptions = field(default_factory=ProvisionOptions)
@@ -307,15 +309,12 @@ def _build_context(
if not gateway.workspace_root:
msg = "gateway_workspace_root is required"
raise ValueError(msg)
if not gateway.main_session_key:
msg = "gateway_main_session_key is required"
raise ValueError(msg)
agent_id = str(agent.id)
workspace_root = gateway.workspace_root
workspace_path = _workspace_path(agent, workspace_root)
session_key = agent.openclaw_session_id or ""
base_url = settings.base_url or "REPLACE_WITH_BASE_URL"
main_session_key = gateway.main_session_key
main_session_key = gateway_agent_session_key(gateway)
identity_profile: dict[str, Any] = {}
if isinstance(agent.identity_profile, dict):
identity_profile = agent.identity_profile
@@ -411,7 +410,7 @@ def _build_main_context(
"session_key": agent.openclaw_session_id or "",
"base_url": base_url,
"auth_token": auth_token,
"main_session_key": gateway.main_session_key or "",
"main_session_key": gateway_agent_session_key(gateway),
"workspace_root": gateway.workspace_root or "",
"user_name": (user.name or "") if user else "",
"user_preferred_name": preferred_name,
@@ -870,20 +869,30 @@ async def provision_main_agent(
gateway = request.gateway
if not gateway.url:
return
if not gateway.main_session_key:
msg = "gateway main_session_key is required"
session_key = (request.session_key or gateway.main_session_key or "").strip()
if not session_key:
msg = "gateway main agent session_key is required"
raise ValueError(msg)
client_config = GatewayClientConfig(url=gateway.url, token=gateway.token)
await ensure_session(
gateway.main_session_key,
session_key,
config=client_config,
label="Main Agent",
label=agent.name or "Gateway Agent",
)
agent_id = await _gateway_default_agent_id(
client_config,
fallback_session_key=gateway.main_session_key,
)
agent_id = _agent_id_from_session_key(session_key)
if agent_id:
if not gateway.workspace_root:
msg = "gateway_workspace_root is required"
raise ValueError(msg)
workspace_path = _workspace_path(agent, gateway.workspace_root)
heartbeat = _heartbeat_config(agent)
await _patch_gateway_agent_list(agent_id, workspace_path, heartbeat, client_config)
else:
agent_id = await _gateway_default_agent_id(
client_config,
fallback_session_key=session_key,
)
if not agent_id:
msg = "Unable to resolve gateway main agent id"
raise OpenClawGatewayError(msg)
@@ -912,7 +921,7 @@ async def provision_main_agent(
client_config=client_config,
)
if request.options.reset_session:
await _reset_session(gateway.main_session_key, client_config)
await _reset_session(session_key, client_config)
async def cleanup_agent(

View File

@@ -19,6 +19,7 @@ from app.schemas.approvals import ApprovalRead
from app.schemas.board_memory import BoardMemoryRead
from app.schemas.boards import BoardRead
from app.schemas.view_models import BoardSnapshot, TaskCardRead
from app.services.gateway_agents import gateway_agent_session_key
from app.services.task_dependencies import (
blocked_by_dependency_ids,
dependency_ids_by_task_id,
@@ -47,8 +48,8 @@ def _computed_agent_status(agent: Agent) -> str:
async def _gateway_main_session_keys(session: AsyncSession) -> set[str]:
keys = (await session.exec(select(Gateway.main_session_key))).all()
return {key for key in keys if key}
gateways = await Gateway.objects.all().all(session)
return {gateway_agent_session_key(gateway) for gateway in gateways}
def _agent_to_read(agent: Agent, main_session_keys: set[str]) -> AgentRead:

View File

@@ -0,0 +1,34 @@
"""Helpers for dedicated gateway-scoped agent identity/session keys."""
from __future__ import annotations
from uuid import UUID
from app.models.gateways import Gateway
_GATEWAY_AGENT_PREFIX = "agent:gateway-"
_GATEWAY_AGENT_SUFFIX = ":main"
def gateway_agent_session_key_for_id(gateway_id: UUID) -> str:
"""Return the dedicated Mission Control gateway-agent session key for an id."""
return f"{_GATEWAY_AGENT_PREFIX}{gateway_id}{_GATEWAY_AGENT_SUFFIX}"
def gateway_agent_session_key(gateway: Gateway) -> str:
"""Return the dedicated Mission Control gateway-agent session key."""
return gateway_agent_session_key_for_id(gateway.id)
def parse_gateway_agent_session_key(session_key: str | None) -> UUID | None:
"""Parse a gateway id from a dedicated gateway-agent session key."""
value = (session_key or "").strip()
if not (value.startswith(_GATEWAY_AGENT_PREFIX) and value.endswith(_GATEWAY_AGENT_SUFFIX)):
return None
gateway_id = value[len(_GATEWAY_AGENT_PREFIX) : -len(_GATEWAY_AGENT_SUFFIX)]
if not gateway_id:
return None
try:
return UUID(gateway_id)
except ValueError:
return None

View File

@@ -31,6 +31,7 @@ from app.services.agent_provisioning import (
provision_agent,
provision_main_agent,
)
from app.services.gateway_agents import gateway_agent_session_key
_TOOLS_KV_RE = re.compile(r"^(?P<key>[A-Z0-9_]+)=(?P<value>.*)$")
SESSION_KEY_PARTS_MIN = 2
@@ -520,21 +521,22 @@ async def _sync_main_agent(
ctx: _SyncContext,
result: GatewayTemplatesSyncResult,
) -> bool:
main_session_key = gateway_agent_session_key(ctx.gateway)
main_agent = (
await Agent.objects.all()
.filter(col(Agent.openclaw_session_id) == ctx.gateway.main_session_key)
.filter(col(Agent.openclaw_session_id) == 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."),
message=("Gateway agent record not found; " "skipping gateway 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,
fallback_session_key=main_session_key,
backoff=ctx.backoff,
)
except TimeoutError as exc:
@@ -544,7 +546,7 @@ async def _sync_main_agent(
_append_sync_error(
result,
agent=main_agent,
message="Unable to resolve gateway default agent id for main agent.",
message="Unable to resolve gateway agent id.",
)
return True
@@ -561,7 +563,7 @@ async def _sync_main_agent(
_append_sync_error(
result,
agent=main_agent,
message="Skipping main agent: unable to read AUTH_TOKEN from TOOLS.md.",
message="Skipping gateway agent: unable to read AUTH_TOKEN from TOOLS.md.",
)
return True
stop_sync = False
@@ -574,6 +576,7 @@ async def _sync_main_agent(
gateway=ctx.gateway,
auth_token=token,
user=ctx.options.user,
session_key=main_session_key,
options=ProvisionOptions(
action="update",
force_bootstrap=ctx.options.force_bootstrap,
@@ -590,7 +593,7 @@ async def _sync_main_agent(
_append_sync_error(
result,
agent=main_agent,
message=f"Failed to sync main agent templates: {exc}",
message=f"Failed to sync gateway agent templates: {exc}",
)
else:
result.main_updated = True