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

@@ -57,6 +57,7 @@ from app.schemas.pagination import DefaultLimitOffsetPage
from app.schemas.tasks import TaskCommentCreate, TaskCommentRead, TaskCreate, TaskRead, TaskUpdate
from app.services.activity_log import record_activity
from app.services.board_leads import LeadAgentOptions, LeadAgentRequest, ensure_board_lead_agent
from app.services.gateway_agents import gateway_agent_session_key, parse_gateway_agent_session_key
from app.services.task_dependencies import (
blocked_by_dependency_ids,
dependency_status_by_id,
@@ -177,13 +178,22 @@ async def _require_gateway_main(
status_code=status.HTTP_403_FORBIDDEN,
detail="Agent missing session key",
)
gateway = await Gateway.objects.filter_by(main_session_key=session_key).first(
session,
)
gateway_id = parse_gateway_agent_session_key(session_key)
if gateway_id is None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Only the dedicated gateway agent may call this endpoint.",
)
gateway = await Gateway.objects.by_id(gateway_id).first(session)
if gateway is None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Only the gateway main agent may call this endpoint.",
detail="Only the dedicated gateway agent may call this endpoint.",
)
if gateway_agent_session_key(gateway) != session_key:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Only the dedicated gateway agent may call this endpoint.",
)
if not gateway.url:
raise HTTPException(
@@ -729,7 +739,7 @@ async def ask_user_via_gateway_main(
session: AsyncSession = SESSION_DEP,
agent_ctx: AgentAuthContext = AGENT_CTX_DEP,
) -> GatewayMainAskUserResponse:
"""Route a lead's ask-user request through the gateway main agent."""
"""Route a lead's ask-user request through the dedicated gateway agent."""
import json
_guard_board_access(agent_ctx, board)
@@ -747,11 +757,11 @@ async def ask_user_via_gateway_main(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway is not configured for this board",
)
main_session_key = (gateway.main_session_key or "").strip()
main_session_key = gateway_agent_session_key(gateway)
if not main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway main session key is required",
detail="Gateway agent session key is required",
)
config = GatewayClientConfig(url=gateway.url, token=gateway.token)
@@ -785,13 +795,8 @@ async def ask_user_via_gateway_main(
)
try:
await ensure_session(main_session_key, config=config, label="Main Agent")
await send_message(
message,
session_key=main_session_key,
config=config,
deliver=True,
)
await ensure_session(main_session_key, config=config, label="Gateway Agent")
await send_message(message, session_key=main_session_key, config=config, deliver=True)
except OpenClawGatewayError as exc:
record_activity(
session,
@@ -808,7 +813,7 @@ async def ask_user_via_gateway_main(
record_activity(
session,
event_type="gateway.lead.ask_user.sent",
message=f"Lead requested user info via gateway main for board: {board.name}.",
message=f"Lead requested user info via gateway agent for board: {board.name}.",
agent_id=agent_ctx.agent.id,
)
@@ -871,7 +876,7 @@ async def message_gateway_board_lead(
f"From agent: {agent_ctx.agent.name}\n"
f"{correlation_line}\n"
f"{payload.content.strip()}\n\n"
"Reply to the gateway main by writing a NON-chat memory item on this board:\n"
"Reply to the gateway agent by writing a NON-chat memory item on this board:\n"
f"POST {base_url}/api/v1/agent/boards/{board.id}/memory\n"
f'Body: {{"content":"...","tags":{tags_json},"source":"{reply_source}"}}\n'
"Do NOT reply in OpenClaw chat."
@@ -964,7 +969,7 @@ async def broadcast_gateway_lead_message(
f"From agent: {agent_ctx.agent.name}\n"
f"{correlation_line}\n"
f"{payload.content.strip()}\n\n"
"Reply to the gateway main by writing a NON-chat memory item "
"Reply to the gateway agent by writing a NON-chat memory item "
"on this board:\n"
f"POST {base_url}/api/v1/agent/boards/{target_board.id}/memory\n"
f'Body: {{"content":"...","tags":{tags_json},'

View File

@@ -49,6 +49,7 @@ from app.services.agent_provisioning import (
provision_agent,
provision_main_agent,
)
from app.services.gateway_agents import gateway_agent_session_key
from app.services.organizations import (
OrganizationContext,
get_active_membership,
@@ -178,11 +179,6 @@ async def _require_gateway(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Board gateway_id is invalid",
)
if not gateway.main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway main_session_key is required",
)
if not gateway.url:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
@@ -206,8 +202,8 @@ def _gateway_client_config(gateway: Gateway) -> GatewayClientConfig:
async def _get_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 _is_gateway_main(agent: Agent, main_session_keys: set[str]) -> bool:
@@ -249,7 +245,11 @@ async def _find_gateway_for_main_session(
) -> Gateway | None:
if not session_key:
return None
return await Gateway.objects.filter_by(main_session_key=session_key).first(session)
gateways = await Gateway.objects.all().all(session)
for gateway in gateways:
if gateway_agent_session_key(gateway) == session_key:
return gateway
return None
async def _ensure_gateway_session(
@@ -605,7 +605,7 @@ async def _apply_agent_update_mutations(
gateway_for_main, _ = await _require_gateway(session, board_for_main)
updates["board_id"] = None
agent.is_board_lead = False
agent.openclaw_session_id = gateway_for_main.main_session_key
agent.openclaw_session_id = gateway_agent_session_key(gateway_for_main)
main_gateway = gateway_for_main
elif make_main is not None:
agent.openclaw_session_id = None
@@ -639,12 +639,7 @@ async def _resolve_agent_update_target(
if gateway_for_main is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Main agent requires a gateway main_session_key",
)
if not gateway_for_main.main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway main_session_key is required",
detail="Gateway agent requires a gateway configuration",
)
return _AgentUpdateProvisionTarget(
is_main_agent=True,
@@ -654,11 +649,6 @@ async def _resolve_agent_update_target(
)
if make_main is None and agent.board_id is None and main_gateway is not None:
if not main_gateway.main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway main_session_key is required",
)
return _AgentUpdateProvisionTarget(
is_main_agent=True,
board=None,
@@ -723,6 +713,7 @@ async def _provision_updated_agent(
gateway=request.target.gateway,
auth_token=request.raw_token,
user=request.user,
session_key=agent.openclaw_session_id,
options=ProvisionOptions(
action="update",
force_bootstrap=request.force_bootstrap,
@@ -970,13 +961,15 @@ async def list_agents(
else:
base_filter: ColumnElement[bool] = col(Agent.board_id).in_(board_ids)
if is_org_admin(ctx.member):
gateway_keys = select(Gateway.main_session_key).where(
col(Gateway.organization_id) == ctx.organization.id,
)
base_filter = or_(
base_filter,
col(Agent.openclaw_session_id).in_(gateway_keys),
)
gateways = await Gateway.objects.filter_by(
organization_id=ctx.organization.id,
).all(session)
gateway_keys = [gateway_agent_session_key(gateway) for gateway in gateways]
if gateway_keys:
base_filter = or_(
base_filter,
col(Agent.openclaw_session_id).in_(gateway_keys),
)
statement = select(Agent).where(base_filter)
if board_id is not None:
statement = statement.where(col(Agent.board_id) == board_id)
@@ -1309,9 +1302,9 @@ async def delete_agent(
await session.delete(agent)
await session.commit()
# Always ask the main agent to confirm workspace cleanup.
# Always ask the gateway agent to confirm workspace cleanup.
try:
main_session = gateway.main_session_key
main_session = gateway_agent_session_key(gateway)
if main_session and workspace_path:
cleanup_message = (
"Cleanup request for deleted agent.\n\n"
@@ -1322,7 +1315,7 @@ async def delete_agent(
"1) Remove the workspace directory.\n"
"2) Reply NO_REPLY.\n"
)
await ensure_session(main_session, config=client_config, label="Main Agent")
await ensure_session(main_session, config=client_config, label="Gateway Agent")
await send_message(
cleanup_message,
session_key=main_session,

View File

@@ -36,6 +36,7 @@ from app.schemas.board_onboarding import (
)
from app.schemas.boards import BoardRead
from app.services.board_leads import LeadAgentOptions, LeadAgentRequest, ensure_board_lead_agent
from app.services.gateway_agents import gateway_agent_session_key
if TYPE_CHECKING:
from sqlmodel.ext.asyncio.session import AsyncSession
@@ -60,7 +61,7 @@ async def _gateway_config(
if not board.gateway_id:
raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)
gateway = await Gateway.objects.by_id(board.gateway_id).first(session)
if gateway is None or not gateway.url or not gateway.main_session_key:
if gateway is None or not gateway.url:
raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)
return gateway, GatewayClientConfig(url=gateway.url, token=gateway.token)
@@ -168,7 +169,7 @@ async def start_onboarding(
board: Board = BOARD_USER_WRITE_DEP,
session: AsyncSession = SESSION_DEP,
) -> BoardOnboardingSession:
"""Start onboarding and send instructions to the gateway main agent."""
"""Start onboarding and send instructions to the gateway agent."""
onboarding = (
await BoardOnboardingSession.objects.filter_by(board_id=board.id)
.filter(col(BoardOnboardingSession.status) == "active")
@@ -178,12 +179,12 @@ async def start_onboarding(
return onboarding
gateway, config = await _gateway_config(session, board)
session_key = gateway.main_session_key
session_key = gateway_agent_session_key(gateway)
base_url = settings.base_url or "http://localhost:8000"
prompt = (
"BOARD ONBOARDING REQUEST\n\n"
f"Board Name: {board.name}\n"
"You are the main agent. Ask the user 6-10 focused questions total:\n"
"You are the gateway agent. Ask the user 6-10 focused questions total:\n"
"- 3-6 questions to clarify the board goal.\n"
"- 1 question to choose a unique name for the board lead agent "
"(first-name style).\n"
@@ -246,7 +247,7 @@ async def start_onboarding(
)
try:
await ensure_session(session_key, config=config, label="Main Agent")
await ensure_session(session_key, config=config, label="Gateway Agent")
await send_message(
prompt,
session_key=session_key,
@@ -279,7 +280,7 @@ async def answer_onboarding(
board: Board = BOARD_USER_WRITE_DEP,
session: AsyncSession = SESSION_DEP,
) -> BoardOnboardingSession:
"""Send a user onboarding answer to the gateway main agent."""
"""Send a user onboarding answer to the gateway agent."""
onboarding = (
await BoardOnboardingSession.objects.filter_by(board_id=board.id)
.order_by(col(BoardOnboardingSession.updated_at).desc())
@@ -299,7 +300,7 @@ async def answer_onboarding(
)
try:
await ensure_session(onboarding.session_key, config=config, label="Main Agent")
await ensure_session(onboarding.session_key, config=config, label="Gateway Agent")
await send_message(
answer_text,
session_key=onboarding.session_key,
@@ -327,7 +328,7 @@ async def agent_onboarding_update(
session: AsyncSession = SESSION_DEP,
actor: ActorContext = ACTOR_DEP,
) -> BoardOnboardingSession:
"""Store onboarding updates submitted by the gateway main agent."""
"""Store onboarding updates submitted by the gateway agent."""
if actor.actor_type != "agent" or actor.agent is None:
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
agent = actor.agent
@@ -338,9 +339,8 @@ async def agent_onboarding_update(
gateway = await Gateway.objects.by_id(board.gateway_id).first(session)
if (
gateway
and gateway.main_session_key
and agent.openclaw_session_id
and agent.openclaw_session_id != gateway.main_session_key
and agent.openclaw_session_id != gateway_agent_session_key(gateway)
):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)

View File

@@ -47,6 +47,7 @@ from app.schemas.pagination import DefaultLimitOffsetPage
from app.schemas.view_models import BoardGroupSnapshot, BoardSnapshot
from app.services.board_group_snapshot import build_board_group_snapshot
from app.services.board_snapshot import build_board_snapshot
from app.services.gateway_agents import gateway_agent_session_key
from app.services.organizations import OrganizationContext, board_access_filter
if TYPE_CHECKING:
@@ -195,11 +196,6 @@ async def _board_gateway(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Board gateway_id is invalid",
)
if not config.main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway main_session_key is required",
)
if not config.url:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
@@ -220,7 +216,7 @@ async def _cleanup_agent_on_gateway(
) -> None:
if agent.openclaw_session_id:
await delete_session(agent.openclaw_session_id, config=client_config)
main_session = config.main_session_key
main_session = gateway_agent_session_key(config)
workspace_root = config.workspace_root
workspace_path = f"{workspace_root.rstrip('/')}/workspace-{_slugify(agent.name)}"
cleanup_message = (
@@ -234,7 +230,7 @@ async def _cleanup_agent_on_gateway(
"2) Delete any lingering session artifacts.\n"
"Reply NO_REPLY."
)
await ensure_session(main_session, config=client_config, label="Main Agent")
await ensure_session(main_session, config=client_config, label="Gateway Agent")
await send_message(
cleanup_message,
session_key=main_session,

View File

@@ -35,6 +35,7 @@ from app.schemas.gateway_api import (
GatewaySessionsResponse,
GatewaysStatusResponse,
)
from app.services.gateway_agents import gateway_agent_session_key
from app.services.organizations import OrganizationContext, require_board_access
if TYPE_CHECKING:
@@ -47,11 +48,18 @@ SESSION_DEP = Depends(get_session)
AUTH_DEP = Depends(get_auth_context)
ORG_ADMIN_DEP = Depends(require_org_admin)
BOARD_ID_QUERY = Query(default=None)
RESOLVE_QUERY_DEP = Depends()
def _query_to_resolve_input(params: GatewayResolveQuery) -> GatewayResolveQuery:
return params
def _query_to_resolve_input(
board_id: str | None = Query(default=None),
gateway_url: str | None = Query(default=None),
gateway_token: str | None = Query(default=None),
) -> GatewayResolveQuery:
return GatewayResolveQuery(
board_id=board_id,
gateway_url=gateway_url,
gateway_token=gateway_token,
)
RESOLVE_INPUT_DEP = Depends(_query_to_resolve_input)
@@ -81,7 +89,7 @@ async def _resolve_gateway(
return (
None,
GatewayClientConfig(url=params.gateway_url, token=params.gateway_token),
params.gateway_main_session_key,
None,
)
if not params.board_id:
raise HTTPException(
@@ -115,7 +123,7 @@ async def _resolve_gateway(
return (
board,
GatewayClientConfig(url=gateway.url, token=gateway.token),
gateway.main_session_key,
gateway_agent_session_key(gateway),
)
@@ -167,7 +175,7 @@ async def gateways_status(
ensured = await ensure_session(
main_session,
config=config,
label="Main Agent",
label="Gateway Agent",
)
if isinstance(ensured, dict):
main_session_entry = ensured.get("entry") or ensured
@@ -224,7 +232,7 @@ async def list_gateway_sessions(
ensured = await ensure_session(
main_session,
config=config,
label="Main Agent",
label="Gateway Agent",
)
if isinstance(ensured, dict):
main_session_entry = ensured.get("entry") or ensured
@@ -256,7 +264,7 @@ async def _with_main_session(
if not main_session or any(item.get("key") == main_session for item in sessions_list):
return sessions_list
try:
await ensure_session(main_session, config=config, label="Main Agent")
await ensure_session(main_session, config=config, label="Gateway Agent")
return await _list_sessions(config)
except OpenClawGatewayError:
return sessions_list
@@ -300,7 +308,7 @@ async def get_gateway_session(
ensured = await ensure_session(
main_session,
config=config,
label="Main Agent",
label="Gateway Agent",
)
if isinstance(ensured, dict):
session_entry = ensured.get("entry") or ensured
@@ -360,7 +368,7 @@ async def send_gateway_session_message(
await require_board_access(session, user=auth.user, board=board, write=True)
try:
if main_session and session_id == main_session:
await ensure_session(main_session, config=config, label="Main Agent")
await ensure_session(main_session, config=config, label="Gateway Agent")
await send_message(payload.content, session_key=session_id, config=config)
except OpenClawGatewayError as exc:
raise HTTPException(

View File

@@ -4,7 +4,7 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING
from uuid import UUID
from uuid import UUID, uuid4
from fastapi import APIRouter, Depends, HTTPException, Query, status
from sqlmodel import col
@@ -34,6 +34,7 @@ from app.services.agent_provisioning import (
ProvisionOptions,
provision_main_agent,
)
from app.services.gateway_agents import gateway_agent_session_key, gateway_agent_session_key_for_id
from app.services.template_sync import GatewayTemplateSyncOptions
from app.services.template_sync import sync_gateway_templates as sync_gateway_templates_service
@@ -85,7 +86,7 @@ SYNC_QUERY_DEP = Depends(_template_sync_query)
def _main_agent_name(gateway: Gateway) -> str:
return f"{gateway.name} Main"
return f"{gateway.name} Gateway Agent"
async def _require_gateway(
@@ -113,6 +114,15 @@ async def _find_main_agent(
previous_name: str | None = None,
previous_session_key: str | None = None,
) -> Agent | None:
preferred_session_key = gateway_agent_session_key(gateway)
if preferred_session_key:
agent = await Agent.objects.filter_by(
openclaw_session_id=preferred_session_key,
).first(
session,
)
if agent:
return agent
if gateway.main_session_key:
agent = await Agent.objects.filter_by(
openclaw_session_id=gateway.main_session_key,
@@ -147,8 +157,13 @@ async def _ensure_main_agent(
previous: tuple[str | None, str | None] | None = None,
action: str = "provision",
) -> Agent | None:
if not gateway.url or not gateway.main_session_key:
if not gateway.url:
return None
session_key = gateway_agent_session_key(gateway)
if gateway.main_session_key != session_key:
gateway.main_session_key = session_key
gateway.updated_at = utcnow()
session.add(gateway)
agent = await _find_main_agent(
session,
gateway,
@@ -161,17 +176,17 @@ async def _ensure_main_agent(
status="provisioning",
board_id=None,
is_board_lead=False,
openclaw_session_id=gateway.main_session_key,
openclaw_session_id=session_key,
heartbeat_config=DEFAULT_HEARTBEAT_CONFIG.copy(),
identity_profile={
"role": "Main Agent",
"role": "Gateway Agent",
"communication_style": "direct, concise, practical",
"emoji": ":compass:",
},
)
session.add(agent)
agent.name = _main_agent_name(gateway)
agent.openclaw_session_id = gateway.main_session_key
agent.openclaw_session_id = session_key
raw_token = generate_agent_token()
agent.agent_token_hash = hash_agent_token(raw_token)
agent.provision_requested_at = utcnow()
@@ -189,11 +204,12 @@ async def _ensure_main_agent(
gateway=gateway,
auth_token=raw_token,
user=auth.user,
session_key=session_key,
options=ProvisionOptions(action=action),
),
)
await ensure_session(
gateway.main_session_key,
session_key,
config=GatewayClientConfig(url=gateway.url, token=gateway.token),
label=agent.name,
)
@@ -204,7 +220,7 @@ async def _ensure_main_agent(
"If BOOTSTRAP.md exists, run it once then delete it. "
"Begin heartbeats after startup."
),
session_key=gateway.main_session_key,
session_key=session_key,
config=GatewayClientConfig(url=gateway.url, token=gateway.token),
deliver=True,
)
@@ -237,7 +253,10 @@ async def create_gateway(
) -> Gateway:
"""Create a gateway and provision or refresh its main agent."""
data = payload.model_dump()
gateway_id = uuid4()
data["id"] = gateway_id
data["organization_id"] = ctx.organization.id
data["main_session_key"] = gateway_agent_session_key_for_id(gateway_id)
gateway = await crud.create(session, Gateway, **data)
await _ensure_main_agent(session, gateway, auth, action="provision")
return gateway