refactor: update gateway agent session key handling and improve related logic
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
34
backend/app/services/gateway_agents.py
Normal file
34
backend/app/services/gateway_agents.py
Normal 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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user