feat: enhance agent heartbeat handling and validate onboarding goal fields

This commit is contained in:
Abhimanyu Saharan
2026-02-06 11:50:14 +05:30
parent 814c041785
commit 7c569c150d
7 changed files with 78 additions and 13 deletions

View File

@@ -20,7 +20,7 @@ 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, AgentHeartbeatCreate, AgentNudge, AgentRead
from app.schemas.agents import AgentCreate, AgentHeartbeat, AgentHeartbeatCreate, AgentNudge, AgentRead
from app.schemas.approvals import ApprovalCreate, ApprovalRead
from app.schemas.board_memory import BoardMemoryCreate, BoardMemoryRead
from app.schemas.board_onboarding import BoardOnboardingRead
@@ -376,14 +376,10 @@ async def agent_heartbeat(
session: Session = Depends(get_session),
agent_ctx: AgentAuthContext = Depends(get_agent_auth_context),
) -> AgentRead:
if agent_ctx.agent.name != payload.name:
payload = AgentHeartbeatCreate(
name=agent_ctx.agent.name,
status=payload.status,
board_id=payload.board_id,
)
return await agents_api.heartbeat_or_create_agent( # type: ignore[attr-defined]
payload=payload,
# Heartbeats must apply to the authenticated agent; agent names are not unique.
return agents_api.heartbeat_agent( # type: ignore[attr-defined]
agent_id=str(agent_ctx.agent.id),
payload=AgentHeartbeat(status=payload.status),
session=session,
actor=_actor(agent_ctx),
)

View File

@@ -588,7 +588,7 @@ def heartbeat_agent(
payload: AgentHeartbeat,
session: Session = Depends(get_session),
actor: ActorContext = Depends(require_admin_or_agent),
) -> Agent:
) -> AgentRead:
agent = session.get(Agent, agent_id)
if agent is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
@@ -613,8 +613,20 @@ async def heartbeat_or_create_agent(
payload: AgentHeartbeatCreate,
session: Session = Depends(get_session),
actor: ActorContext = Depends(require_admin_or_agent),
) -> Agent:
agent = session.exec(select(Agent).where(Agent.name == payload.name)).first()
) -> AgentRead:
# Agent tokens must heartbeat their authenticated agent record. Names are not unique.
if actor.actor_type == "agent" and actor.agent:
return heartbeat_agent(
agent_id=str(actor.agent.id),
payload=AgentHeartbeat(status=payload.status),
session=session,
actor=actor,
)
statement = select(Agent).where(Agent.name == payload.name)
if payload.board_id is not None:
statement = statement.where(Agent.board_id == payload.board_id)
agent = session.exec(statement).first()
if agent is None:
if actor.actor_type == "agent":
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)

View File

@@ -544,6 +544,20 @@ def update_task(
)
session.commit()
session.refresh(task)
if task.assigned_agent_id and task.assigned_agent_id != previous_assigned:
if actor.actor_type == "agent" and actor.agent and task.assigned_agent_id == actor.agent.id:
return task
assigned_agent = session.get(Agent, task.assigned_agent_id)
if assigned_agent:
board = session.get(Board, task.board_id) if task.board_id else None
if board:
_notify_agent_on_task_assign(
session=session,
board=board,
task=task,
agent=assigned_agent,
)
return task
if actor.actor_type == "agent":
if actor.agent and actor.agent.board_id and task.board_id: