feat: implement gateway main agent functionality and update agent handling

This commit is contained in:
Abhimanyu Saharan
2026-02-05 19:29:17 +05:30
parent 4cc6c42440
commit 28b4a5bba8
4 changed files with 218 additions and 47 deletions

View File

@@ -19,18 +19,13 @@ from app.models.agents import Agent
from app.models.boards import Board
from app.models.gateways import Gateway
from app.models.tasks import Task
from app.schemas.agents import (
AgentCreate,
AgentHeartbeat,
AgentHeartbeatCreate,
AgentRead,
AgentUpdate,
)
from app.schemas.agents import AgentCreate, AgentHeartbeat, AgentHeartbeatCreate, AgentRead, AgentUpdate
from app.services.activity_log import record_activity
from app.services.agent_provisioning import (
DEFAULT_HEARTBEAT_CONFIG,
cleanup_agent,
provision_agent,
provision_main_agent,
)
router = APIRouter(prefix="/agents", tags=["agents"])
@@ -121,6 +116,39 @@ def _require_gateway(session: Session, board: Board) -> tuple[Gateway, GatewayCl
return gateway, GatewayClientConfig(url=gateway.url, token=gateway.token)
def _gateway_client_config(gateway: Gateway) -> GatewayClientConfig:
if not gateway.url:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway url is required",
)
return GatewayClientConfig(url=gateway.url, token=gateway.token)
def _get_gateway_main_session_keys(session: Session) -> set[str]:
keys = session.exec(select(Gateway.main_session_key)).all()
return {key for key in keys if key}
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:
model = AgentRead.model_validate(agent, from_attributes=True)
return model.model_copy(update={"is_gateway_main": _is_gateway_main(agent, main_session_keys)})
def _find_gateway_for_main_session(
session: Session, session_key: str | None
) -> Gateway | None:
if not session_key:
return None
return session.exec(
select(Gateway).where(Gateway.main_session_key == session_key)
).first()
async def _ensure_gateway_session(
agent_name: str,
config: GatewayClientConfig,
@@ -182,7 +210,11 @@ def list_agents(
auth: AuthContext = Depends(require_admin_auth),
) -> list[Agent]:
agents = list(session.exec(select(Agent)))
return [_with_computed_status(agent) for agent in agents]
main_session_keys = _get_gateway_main_session_keys(session)
return [
_to_agent_read(_with_computed_status(agent), main_session_keys)
for agent in agents
]
@router.post("", response_model=AgentRead)
@@ -294,7 +326,8 @@ def get_agent(
agent = session.get(Agent, agent_id)
if agent is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
return _with_computed_status(agent)
main_session_keys = _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
@router.patch("/{agent_id}", response_model=AgentRead)
@@ -309,6 +342,7 @@ async def update_agent(
if agent is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
updates = payload.model_dump(exclude_unset=True)
make_main = updates.pop("is_gateway_main", None)
if "status" in updates:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
@@ -322,22 +356,71 @@ async def update_agent(
updates["identity_profile"] = _normalize_identity_profile(
updates.get("identity_profile")
)
if not updates and not force:
return _with_computed_status(agent)
if "board_id" in updates:
if not updates and not force and make_main is None:
main_session_keys = _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
main_gateway = _find_gateway_for_main_session(session, agent.openclaw_session_id)
gateway_for_main: Gateway | None = None
if make_main is True:
board_source = updates.get("board_id") or agent.board_id
board_for_main = _require_board(session, board_source)
gateway_for_main, _ = _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
main_gateway = gateway_for_main
elif make_main is False:
agent.openclaw_session_id = None
if make_main is not True and "board_id" in updates:
_require_board(session, updates["board_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.is_board_lead = False
agent.updated_at = datetime.utcnow()
if agent.heartbeat_config is None:
agent.heartbeat_config = DEFAULT_HEARTBEAT_CONFIG.copy()
session.add(agent)
session.commit()
session.refresh(agent)
is_main_agent = False
board: Board | None = None
gateway: Gateway | None = None
client_config: GatewayClientConfig | None = None
if make_main is True:
is_main_agent = True
gateway = gateway_for_main
elif make_main is None and agent.board_id is None and main_gateway is not None:
is_main_agent = True
gateway = main_gateway
if is_main_agent:
if gateway is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Main agent requires a gateway main_session_key",
)
if not gateway.main_session_key:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway main_session_key is required",
)
client_config = _gateway_client_config(gateway)
else:
if agent.board_id is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="board_id is required for non-main agents",
)
board = _require_board(session, agent.board_id)
gateway, client_config = _require_gateway(session, board)
session_key = agent.openclaw_session_id or _build_session_key(agent.name)
try:
if client_config is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway configuration is required",
)
await ensure_session(session_key, config=client_config, label=agent.name)
if not agent.openclaw_session_id:
agent.openclaw_session_id = session_key
@@ -356,6 +439,14 @@ async def update_agent(
session.commit()
session.refresh(agent)
try:
if gateway is None:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail="Gateway configuration is required",
)
if is_main_agent:
await provision_main_agent(agent, gateway, raw_token, auth.user, action="update")
else:
await provision_agent(agent, board, gateway, raw_token, auth.user, action="update")
await _send_wakeup_message(agent, client_config, verb="updated")
agent.provision_confirm_token_hash = None
@@ -392,7 +483,8 @@ async def update_agent(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Unexpected error updating agent provisioning.",
) from exc
return _with_computed_status(agent)
main_session_keys = _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
@router.post("/{agent_id}/heartbeat", response_model=AgentRead)
@@ -417,7 +509,8 @@ def heartbeat_agent(
session.add(agent)
session.commit()
session.refresh(agent)
return _with_computed_status(agent)
main_session_keys = _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
@router.post("/heartbeat", response_model=AgentRead)
@@ -562,7 +655,8 @@ async def heartbeat_or_create_agent(
session.add(agent)
session.commit()
session.refresh(agent)
return _with_computed_status(agent)
main_session_keys = _get_gateway_main_session_keys(session)
return _to_agent_read(_with_computed_status(agent), main_session_keys)
@router.delete("/{agent_id}")

View File

@@ -23,6 +23,7 @@ class AgentCreate(AgentBase):
class AgentUpdate(SQLModel):
board_id: UUID | None = None
is_gateway_main: bool | None = None
name: str | None = None
status: str | None = None
heartbeat_config: dict[str, Any] | None = None
@@ -34,6 +35,7 @@ class AgentUpdate(SQLModel):
class AgentRead(AgentBase):
id: UUID
is_board_lead: bool = False
is_gateway_main: bool = False
openclaw_session_id: str | None = None
last_seen_at: datetime | None
created_at: datetime

View File

@@ -32,6 +32,7 @@ type Agent = {
id: string;
name: string;
board_id?: string | null;
is_gateway_main?: boolean;
heartbeat_config?: {
every?: string;
target?: string;
@@ -109,6 +110,8 @@ export default function EditAgentPage() {
const [name, setName] = useState("");
const [boards, setBoards] = useState<Board[]>([]);
const [boardId, setBoardId] = useState("");
const [boardTouched, setBoardTouched] = useState(false);
const [isGatewayMain, setIsGatewayMain] = useState(false);
const [heartbeatEvery, setHeartbeatEvery] = useState("10m");
const [heartbeatTarget, setHeartbeatTarget] = useState("none");
const [identityProfile, setIdentityProfile] = useState<IdentityProfile>({
@@ -150,9 +153,13 @@ export default function EditAgentPage() {
const data = (await response.json()) as Agent;
setAgent(data);
setName(data.name);
if (data.board_id) {
setIsGatewayMain(Boolean(data.is_gateway_main));
if (!data.is_gateway_main && data.board_id) {
setBoardId(data.board_id);
} else {
setBoardId("");
}
setBoardTouched(false);
if (data.heartbeat_config?.every) {
setHeartbeatEvery(data.heartbeat_config.every);
}
@@ -175,7 +182,7 @@ export default function EditAgentPage() {
}, [isSignedIn, agentId]);
useEffect(() => {
if (boardId) return;
if (boardTouched || boardId || isGatewayMain) return;
if (agent?.board_id) {
setBoardId(agent.board_id);
return;
@@ -183,7 +190,7 @@ export default function EditAgentPage() {
if (boards.length > 0) {
setBoardId(boards[0].id);
}
}, [agent, boards, boardId]);
}, [agent, boards, boardId, isGatewayMain, boardTouched]);
const handleSubmit = async (event: React.FormEvent<HTMLFormElement>) => {
event.preventDefault();
@@ -193,14 +200,37 @@ export default function EditAgentPage() {
setError("Agent name is required.");
return;
}
if (!boardId) {
setError("Select a board before saving.");
if (!isGatewayMain && !boardId) {
setError("Select a board or mark this agent as the gateway main.");
return;
}
if (isGatewayMain && !boardId && !agent?.is_gateway_main && !agent?.board_id) {
setError(
"Select a board once so we can resolve the gateway main session key."
);
return;
}
setIsLoading(true);
setError(null);
try {
const token = await getToken();
const payload: Record<string, unknown> = {
name: trimmed,
heartbeat_config: {
every: heartbeatEvery.trim() || "10m",
target: heartbeatTarget,
},
identity_profile: normalizeIdentityProfile(identityProfile),
soul_template: soulTemplate.trim() || null,
};
if (!isGatewayMain) {
payload.board_id = boardId || null;
} else if (boardId) {
payload.board_id = boardId;
}
if (agent?.is_gateway_main !== isGatewayMain) {
payload.is_gateway_main = isGatewayMain;
}
const response = await fetch(
`${apiBase}/api/v1/agents/${agentId}?force=true`,
{
@@ -209,16 +239,7 @@ export default function EditAgentPage() {
"Content-Type": "application/json",
Authorization: token ? `Bearer ${token}` : "",
},
body: JSON.stringify({
name: trimmed,
board_id: boardId,
heartbeat_config: {
every: heartbeatEvery.trim() || "10m",
target: heartbeatTarget,
},
identity_profile: normalizeIdentityProfile(identityProfile),
soul_template: soulTemplate.trim() || null,
}),
body: JSON.stringify(payload),
}
);
if (!response.ok) {
@@ -303,15 +324,40 @@ export default function EditAgentPage() {
</div>
<div className="grid gap-6 md:grid-cols-2">
<div className="space-y-2">
<div className="flex items-center justify-between">
<label className="text-sm font-medium text-slate-900">
Board <span className="text-red-500">*</span>
Board
{isGatewayMain ? (
<span className="ml-2 text-xs font-normal text-slate-500">
optional
</span>
) : (
<span className="text-red-500"> *</span>
)}
</label>
{boardId ? (
<button
type="button"
className="text-xs font-medium text-slate-600 hover:text-slate-900"
onClick={() => {
setBoardTouched(true);
setBoardId("");
}}
disabled={isLoading}
>
Clear board
</button>
) : null}
</div>
<SearchableSelect
ariaLabel="Select board"
value={boardId}
onValueChange={setBoardId}
onValueChange={(value) => {
setBoardTouched(true);
setBoardId(value);
}}
options={getBoardOptions(boards)}
placeholder="Select board"
placeholder={isGatewayMain ? "No board (main agent)" : "Select board"}
searchPlaceholder="Search boards..."
emptyMessage="No matching boards."
triggerClassName="w-full h-11 rounded-xl border border-slate-300 bg-white px-3 py-2 text-sm font-medium text-slate-900 shadow-sm focus:border-blue-500 focus:ring-2 focus:ring-blue-200"
@@ -319,7 +365,13 @@ export default function EditAgentPage() {
itemClassName="px-4 py-3 text-sm text-slate-700 data-[selected=true]:bg-slate-50 data-[selected=true]:text-slate-900"
disabled={boards.length === 0}
/>
{boards.length === 0 ? (
{isGatewayMain ? (
<p className="text-xs text-slate-500">
Main agents are not attached to a board. If a board is
selected, it is only used to resolve the gateway main
session key and will be cleared on save.
</p>
) : boards.length === 0 ? (
<p className="text-xs text-slate-500">
Create a board before assigning agents.
</p>
@@ -353,6 +405,26 @@ export default function EditAgentPage() {
</div>
</div>
</div>
<div className="mt-6 rounded-xl border border-slate-200 bg-slate-50 p-4">
<label className="flex items-start gap-3 text-sm text-slate-700">
<input
type="checkbox"
className="mt-1 h-4 w-4 rounded border-slate-300 text-blue-600 focus:ring-blue-200"
checked={isGatewayMain}
onChange={(event) => setIsGatewayMain(event.target.checked)}
disabled={isLoading}
/>
<span>
<span className="block font-medium text-slate-900">
Gateway main agent
</span>
<span className="block text-xs text-slate-500">
Uses the gateway main session key and is not tied to a
single board.
</span>
</span>
</label>
</div>
</div>
<div>

View File

@@ -32,6 +32,7 @@ type Agent = {
updated_at: string;
board_id?: string | null;
is_board_lead?: boolean;
is_gateway_main?: boolean;
};
type Board = {
@@ -103,9 +104,9 @@ export default function AgentDetailPage() {
return events.filter((event) => event.agent_id === agent.id);
}, [events, agent]);
const linkedBoard = useMemo(() => {
if (!agent?.board_id) return null;
if (!agent?.board_id || agent?.is_gateway_main) return null;
return boards.find((board) => board.id === agent.board_id) ?? null;
}, [boards, agent?.board_id]);
}, [boards, agent?.board_id, agent?.is_gateway_main]);
const loadAgent = async () => {
@@ -267,7 +268,9 @@ export default function AgentDetailPage() {
<p className="text-xs font-semibold uppercase tracking-[0.2em] text-quiet">
Board
</p>
{linkedBoard ? (
{agent.is_gateway_main ? (
<p className="mt-1 text-sm text-strong">Gateway main (no board)</p>
) : linkedBoard ? (
<Link
href={`/boards/${linkedBoard.id}`}
className="mt-1 inline-flex text-sm font-medium text-[color:var(--accent)] transition hover:underline"