unchained/mission-control
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: remove main session key references and enhance agent-gateway associations

Browse Source
This commit is contained in:
Abhimanyu Saharan
2026-02-10 02:58:58 +05:30
parent 50f71960de
commit 6f070df74b
27 changed files with 552 additions and 672 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
backend/app/api/agent.py
View File

@@ -57,7 +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.gateway_agents import gateway_agent_session_key
from app.services.task_dependencies import (
blocked_by_dependency_ids,
dependency_status_by_id,
@@ -172,25 +172,19 @@ async def _require_gateway_main(
session: AsyncSession,
agent: Agent,
) -> tuple[Gateway, GatewayClientConfig]:
session_key = (agent.openclaw_session_id or "").strip()
if not session_key:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Agent missing session key",
)
gateway_id = parse_gateway_agent_session_key(session_key)
if gateway_id is None:
if agent.board_id is not None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Only the dedicated gateway agent may call this endpoint.",
)
gateway_id = agent.gateway_id
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 dedicated gateway agent may call this endpoint.",
)
if gateway_agent_session_key(gateway) != session_key:
if agent.openclaw_session_id != gateway_agent_session_key(gateway):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Only the dedicated gateway agent may call this endpoint.",
@@ -257,7 +251,6 @@ async def list_agents(
statement = statement.where(Agent.board_id == agent_ctx.agent.board_id)
elif board_id:
statement = statement.where(Agent.board_id == board_id)
main_session_keys = await agents_api.get_gateway_main_session_keys(session)
statement = statement.order_by(col(Agent.created_at).desc())
def _transform(items: Sequence[Any]) -> Sequence[Any]:
@@ -265,7 +258,6 @@ async def list_agents(
return [
agents_api.to_agent_read(
agents_api.with_computed_status(agent),
main_session_keys,
)
for agent in agents
]
@@ -758,11 +750,6 @@ async def ask_user_via_gateway_main(
detail="Gateway is not configured for this board",
)
main_session_key = gateway_agent_session_key(gateway)
if not main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway agent session key is required",
)
config = GatewayClientConfig(url=gateway.url, token=gateway.token)
correlation = payload.correlation_id.strip() if payload.correlation_id else ""
@@ -818,7 +805,8 @@ async def ask_user_via_gateway_main(
)
main_agent = await Agent.objects.filter_by(
openclaw_session_id=main_session_key,
gateway_id=gateway.id,
board_id=None,
).first(session)
await session.commit()

119
backend/app/api/agents.py
View File

@@ -201,32 +201,20 @@ def _gateway_client_config(gateway: Gateway) -> GatewayClientConfig:
return GatewayClientConfig(url=gateway.url, token=gateway.token)
async def _get_gateway_main_session_keys(session: AsyncSession) -> set[str]:
gateways = await Gateway.objects.all().all(session)
return {gateway_agent_session_key(gateway) for gateway in gateways}
def _is_gateway_main(agent: Agent) -> bool:
return agent.board_id is None
def _is_gateway_main(agent: Agent, main_session_keys: set[str]) -> bool:
return bool(
agent.openclaw_session_id and agent.openclaw_session_id in main_session_keys,
)
def _to_agent_read(agent: Agent, main_session_keys: set[str]) -> AgentRead:
def _to_agent_read(agent: Agent) -> AgentRead:
model = AgentRead.model_validate(agent, from_attributes=True)
return model.model_copy(
update={"is_gateway_main": _is_gateway_main(agent, main_session_keys)},
update={"is_gateway_main": _is_gateway_main(agent)},
)
async def get_gateway_main_session_keys(session: AsyncSession) -> set[str]:
"""Return gateway main-session keys used to compute `is_gateway_main`."""
return await _get_gateway_main_session_keys(session)
def to_agent_read(agent: Agent, main_session_keys: set[str]) -> AgentRead:
def to_agent_read(agent: Agent) -> AgentRead:
"""Convert an `Agent` model into its API read representation."""
return _to_agent_read(agent, main_session_keys)
return _to_agent_read(agent)
def _coerce_agent_items(items: Sequence[Any]) -> list[Agent]:
@@ -239,17 +227,10 @@ def _coerce_agent_items(items: Sequence[Any]) -> list[Agent]:
return agents
async def _find_gateway_for_main_session(
session: AsyncSession,
session_key: str | None,
) -> Gateway | None:
if not session_key:
async def _main_agent_gateway(session: AsyncSession, agent: Agent) -> Gateway | None:
if agent.board_id is not None:
return None
gateways = await Gateway.objects.all().all(session)
for gateway in gateways:
if gateway_agent_session_key(gateway) == session_key:
return gateway
return None
return await Gateway.objects.by_id(agent.gateway_id).first(session)
async def _ensure_gateway_session(
@@ -281,8 +262,8 @@ def with_computed_status(agent: Agent) -> Agent:
return _with_computed_status(agent)
def _serialize_agent(agent: Agent, main_session_keys: set[str]) -> dict[str, object]:
return _to_agent_read(_with_computed_status(agent), main_session_keys).model_dump(
def _serialize_agent(agent: Agent) -> dict[str, object]:
return _to_agent_read(_with_computed_status(agent)).model_dump(
mode="json",
)
@@ -331,10 +312,7 @@ async def _require_agent_access(
if agent.board_id is None:
if not is_org_admin(ctx.member):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
gateway = await _find_gateway_for_main_session(
session,
agent.openclaw_session_id,
)
gateway = await _main_agent_gateway(session, agent)
if gateway is None or gateway.organization_id != ctx.organization.id:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
return
@@ -593,10 +571,7 @@ async def _apply_agent_update_mutations(
updates: dict[str, Any],
make_main: bool | None,
) -> tuple[Gateway | None, Gateway | None]:
main_gateway = await _find_gateway_for_main_session(
session,
agent.openclaw_session_id,
)
main_gateway = await _main_agent_gateway(session, agent)
gateway_for_main: Gateway | None = None
if make_main:
@@ -604,20 +579,48 @@ async def _apply_agent_update_mutations(
board_for_main = await _require_board(session, board_source)
gateway_for_main, _ = await _require_gateway(session, board_for_main)
updates["board_id"] = None
updates["gateway_id"] = gateway_for_main.id
agent.is_board_lead = False
agent.openclaw_session_id = gateway_agent_session_key(gateway_for_main)
main_gateway = gateway_for_main
elif make_main is not None:
if "board_id" not in updates or updates["board_id"] is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="board_id is required when converting a gateway-main agent to board scope",
)
board = await _require_board(session, updates["board_id"])
if board.gateway_id is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Board gateway_id is required",
)
updates["gateway_id"] = board.gateway_id
agent.openclaw_session_id = None
if not make_main and "board_id" in updates:
await _require_board(session, updates["board_id"])
if make_main is None and "board_id" in updates:
board = await _require_board(session, updates["board_id"])
if board.gateway_id is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Board gateway_id is required",
)
updates["gateway_id"] = board.gateway_id
for key, value in updates.items():
setattr(agent, key, value)
if make_main is None and main_gateway is not None:
agent.board_id = None
agent.gateway_id = main_gateway.id
agent.is_board_lead = False
if make_main is False and agent.board_id is not None:
board = await _require_board(session, agent.board_id)
if board.gateway_id is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Board gateway_id is required",
)
agent.gateway_id = board.gateway_id
agent.updated_at = utcnow()
if agent.heartbeat_config is None:
agent.heartbeat_config = DEFAULT_HEARTBEAT_CONFIG.copy()
@@ -812,6 +815,7 @@ async def _create_agent_from_heartbeat(
data: dict[str, Any] = {
"name": payload.name,
"board_id": board.id,
"gateway_id": gateway.id,
"heartbeat_config": DEFAULT_HEARTBEAT_CONFIG.copy(),
}
agent, raw_token, session_error = await _persist_new_agent(
@@ -925,8 +929,7 @@ async def _commit_heartbeat(
session.add(agent)
await session.commit()
await session.refresh(agent)
main_session_keys = await _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
return _to_agent_read(_with_computed_status(agent))
async def _send_wakeup_message(
@@ -952,7 +955,6 @@ async def list_agents(
ctx: OrganizationContext = ORG_ADMIN_DEP,
) -> LimitOffsetPage[AgentRead]:
"""List agents visible to the active organization admin."""
main_session_keys = await _get_gateway_main_session_keys(session)
board_ids = await list_accessible_board_ids(session, member=ctx.member, write=False)
if board_id is not None and board_id not in set(board_ids):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
@@ -963,9 +965,11 @@ async def list_agents(
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_filters.append(col(Agent.openclaw_session_id).in_(gateway_keys))
gateway_ids = [gateway.id for gateway in gateways]
if gateway_ids:
base_filters.append(
(col(Agent.gateway_id).in_(gateway_ids)) & (col(Agent.board_id).is_(None)),
)
if base_filters:
if len(base_filters) == 1:
statement = select(Agent).where(base_filters[0])
@@ -979,19 +983,18 @@ async def list_agents(
gateway = await Gateway.objects.by_id(gateway_id).first(session)
if gateway is None or gateway.organization_id != ctx.organization.id:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
gateway_main_key = gateway_agent_session_key(gateway)
gateway_board_ids = select(Board.id).where(col(Board.gateway_id) == gateway_id)
statement = statement.where(
or_(
col(Agent.board_id).in_(gateway_board_ids),
col(Agent.openclaw_session_id) == gateway_main_key,
(col(Agent.gateway_id) == gateway_id) & (col(Agent.board_id).is_(None)),
),
)
statement = statement.order_by(col(Agent.created_at).desc())
def _transform(items: Sequence[Any]) -> Sequence[Any]:
agents = _coerce_agent_items(items)
return [_to_agent_read(_with_computed_status(agent), main_session_keys) for agent in agents]
return [_to_agent_read(_with_computed_status(agent)) for agent in agents]
return await paginate(session, statement, transformer=_transform)
@@ -1029,13 +1032,10 @@ async def stream_agents(
agents = [agent for agent in agents if agent.board_id in allowed_ids]
else:
agents = []
main_session_keys = (
await _get_gateway_main_session_keys(stream_session) if agents else set()
)
for agent in agents:
updated_at = agent.updated_at or agent.last_seen_at or utcnow()
last_seen = max(updated_at, last_seen)
payload = {"agent": _serialize_agent(agent, main_session_keys)}
payload = {"agent": _serialize_agent(agent)}
yield {"event": "agent", "data": json.dumps(payload)}
await asyncio.sleep(2)
@@ -1059,6 +1059,7 @@ async def create_agent(
)
gateway, client_config = await _require_gateway(session, board)
data = payload.model_dump()
data["gateway_id"] = gateway.id
requested_name = (data.get("name") or "").strip()
await _ensure_unique_agent_name(
session,
@@ -1089,8 +1090,7 @@ async def create_agent(
request=provision_request,
client_config=client_config,
)
main_session_keys = await _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
return _to_agent_read(_with_computed_status(agent))
@router.get("/{agent_id}", response_model=AgentRead)
@@ -1104,8 +1104,7 @@ async def get_agent(
if agent is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
await _require_agent_access(session, agent=agent, ctx=ctx, write=False)
main_session_keys = await _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
return _to_agent_read(_with_computed_status(agent))
@router.patch("/{agent_id}", response_model=AgentRead)
@@ -1129,8 +1128,7 @@ async def update_agent(
make_main=make_main,
)
if not updates and not params.force and make_main is None:
main_session_keys = await _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
return _to_agent_read(_with_computed_status(agent))
main_gateway, gateway_for_main = await _apply_agent_update_mutations(
session,
agent=agent,
@@ -1164,8 +1162,7 @@ async def update_agent(
agent=agent,
request=provision_request,
)
main_session_keys = await _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
return _to_agent_read(_with_computed_status(agent))
@router.post("/{agent_id}/heartbeat", response_model=AgentRead)

3
backend/app/api/board_onboarding.py
View File

@@ -339,8 +339,7 @@ async def agent_onboarding_update(
gateway = await Gateway.objects.by_id(board.gateway_id).first(session)
if (
gateway
and agent.openclaw_session_id
and agent.openclaw_session_id != gateway_agent_session_key(gateway)
and (agent.gateway_id != gateway.id or agent.board_id is not None)
):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)

13
backend/app/api/gateway.py
View File

@@ -6,6 +6,7 @@ from collections.abc import Iterable
from typing import TYPE_CHECKING
from fastapi import APIRouter, Depends, HTTPException, Query, status
from sqlmodel import col
from app.api.deps import require_org_admin
from app.core.auth import AuthContext, get_auth_context
@@ -23,6 +24,7 @@ from app.integrations.openclaw_gateway_protocol import (
GATEWAY_METHODS,
PROTOCOL_VERSION,
)
from app.models.agents import Agent
from app.models.boards import Board
from app.models.gateways import Gateway
from app.schemas.common import OkResponse
@@ -35,7 +37,6 @@ 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:
@@ -120,10 +121,16 @@ async def _resolve_gateway(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway url is required",
)
main_agent = (
await Agent.objects.filter_by(gateway_id=gateway.id)
.filter(col(Agent.board_id).is_(None))
.first(session)
)
main_session = main_agent.openclaw_session_id if main_agent else None
return (
board,
GatewayClientConfig(url=gateway.url, token=gateway.token),
gateway_agent_session_key(gateway),
main_session,
)
@@ -186,7 +193,6 @@ async def gateways_status(
gateway_url=config.url,
sessions_count=len(sessions_list),
sessions=sessions_list,
main_session_key=main_session,
main_session=main_session_entry,
main_session_error=main_session_error,
)
@@ -241,7 +247,6 @@ async def list_gateway_sessions(
return GatewaySessionsResponse(
sessions=sessions_list,
main_session_key=main_session,
main_session=main_session_entry,
)

96
backend/app/api/gateways.py
View File

@@ -45,7 +45,6 @@ from app.services.agent_provisioning import (
)
from app.services.gateway_agents import (
gateway_agent_session_key,
gateway_agent_session_key_for_id,
gateway_openclaw_agent_id,
)
from app.services.template_sync import GatewayTemplateSyncOptions
@@ -134,68 +133,27 @@ async def _require_gateway(
async def _find_main_agent(
session: AsyncSession,
gateway: Gateway,
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,
).first(
session,
)
if agent:
return agent
if previous_session_key:
agent = await Agent.objects.filter_by(
openclaw_session_id=previous_session_key,
).first(
session,
)
if agent:
return agent
names = {_main_agent_name(gateway)}
if previous_name:
names.add(f"{previous_name} Main")
for name in names:
agent = await Agent.objects.filter_by(name=name).first(session)
if agent:
return agent
return None
return (
await Agent.objects.filter_by(gateway_id=gateway.id)
.filter(col(Agent.board_id).is_(None))
.first(session)
)
async def _upsert_main_agent_record(
session: AsyncSession,
gateway: Gateway,
*,
previous: tuple[str | None, str | None] | None = None,
) -> tuple[Agent, bool]:
changed = False
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)
changed = True
agent = await _find_main_agent(
session,
gateway,
previous_name=previous[0] if previous else None,
previous_session_key=previous[1] if previous else None,
)
agent = await _find_main_agent(session, gateway)
if agent is None:
agent = Agent(
name=_main_agent_name(gateway),
status="provisioning",
board_id=None,
gateway_id=gateway.id,
is_board_lead=False,
openclaw_session_id=session_key,
heartbeat_config=DEFAULT_HEARTBEAT_CONFIG.copy(),
@@ -206,6 +164,9 @@ async def _upsert_main_agent_record(
if agent.board_id is not None:
agent.board_id = None
changed = True
if agent.gateway_id != gateway.id:
agent.gateway_id = gateway.id
changed = True
if agent.is_board_lead:
agent.is_board_lead = False
changed = True
@@ -262,21 +223,15 @@ def _extract_agent_id_from_entry(item: object) -> str | None:
return None
def _extract_config_agents_list(payload: object) -> list[object]:
def _extract_agents_list(payload: object) -> list[object]:
if isinstance(payload, list):
return [item for item in payload]
if not isinstance(payload, dict):
return []
data = payload.get("config") or payload.get("parsed") or {}
if not isinstance(data, dict):
agents = payload.get("agents") or []
if not isinstance(agents, list):
return []
agents = data.get("agents") or {}
if isinstance(agents, list):
return [item for item in agents]
if not isinstance(agents, dict):
return []
agents_list = agents.get("list") or []
if not isinstance(agents_list, list):
return []
return [item for item in agents_list]
return [item for item in agents]
async def _gateway_has_main_agent_entry(gateway: Gateway) -> bool:
@@ -285,11 +240,11 @@ async def _gateway_has_main_agent_entry(gateway: Gateway) -> bool:
config = GatewayClientConfig(url=gateway.url, token=gateway.token)
target_id = gateway_openclaw_agent_id(gateway)
try:
payload = await openclaw_call("config.get", config=config)
payload = await openclaw_call("agents.list", config=config)
except OpenClawGatewayError:
# Avoid treating transient gateway connectivity issues as a missing agent entry.
return True
for item in _extract_config_agents_list(payload):
for item in _extract_agents_list(payload):
if _extract_agent_id_from_entry(item) == target_id:
return True
return False
@@ -376,14 +331,9 @@ async def _ensure_main_agent(
gateway: Gateway,
auth: AuthContext,
*,
previous: tuple[str | None, str | None] | None = None,
action: str = "provision",
) -> Agent:
agent, _ = await _upsert_main_agent_record(
session,
gateway,
previous=previous,
)
agent, _ = await _upsert_main_agent_record(session, gateway)
return await _provision_main_agent_record(
session,
gateway,
@@ -464,7 +414,6 @@ async def create_gateway(
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
@@ -500,15 +449,12 @@ async def update_gateway(
gateway_id=gateway_id,
organization_id=ctx.organization.id,
)
previous_name = gateway.name
previous_session_key = gateway.main_session_key
updates = payload.model_dump(exclude_unset=True)
await crud.patch(session, gateway, updates)
await _ensure_main_agent(
session,
gateway,
auth,
previous=(previous_name, previous_session_key),
action="update",
)
return gateway
@@ -555,14 +501,14 @@ async def delete_gateway(
gateway_id=gateway_id,
organization_id=ctx.organization.id,
)
gateway_session_key = gateway_agent_session_key(gateway)
main_agent = await _find_main_agent(session, gateway)
if main_agent is not None:
await _clear_agent_foreign_keys(session, agent_id=main_agent.id)
await session.delete(main_agent)
duplicate_main_agents = await Agent.objects.filter_by(
openclaw_session_id=gateway_session_key,
gateway_id=gateway.id,
board_id=None,
).all(session)
for agent in duplicate_main_agents:
if main_agent is not None and agent.id == main_agent.id: