feat: add organization-related models and update schemas for organization management
This commit is contained in:
@@ -14,9 +14,9 @@ from sqlmodel import col, select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
from app.api.deps import ActorContext, require_admin_auth, require_admin_or_agent
|
||||
from app.api.deps import ActorContext, require_admin_or_agent, require_org_admin
|
||||
from app.core.agent_tokens import generate_agent_token, hash_agent_token
|
||||
from app.core.auth import AuthContext
|
||||
from app.core.auth import AuthContext, get_auth_context
|
||||
from app.core.time import utcnow
|
||||
from app.db.pagination import paginate
|
||||
from app.db.session import async_session_maker, get_session
|
||||
@@ -26,7 +26,9 @@ from app.models.activity_events import ActivityEvent
|
||||
from app.models.agents import Agent
|
||||
from app.models.boards import Board
|
||||
from app.models.gateways import Gateway
|
||||
from app.models.organizations import Organization
|
||||
from app.models.tasks import Task
|
||||
from app.models.users import User
|
||||
from app.schemas.agents import (
|
||||
AgentCreate,
|
||||
AgentHeartbeat,
|
||||
@@ -43,6 +45,14 @@ from app.services.agent_provisioning import (
|
||||
provision_agent,
|
||||
provision_main_agent,
|
||||
)
|
||||
from app.services.organizations import (
|
||||
OrganizationContext,
|
||||
get_active_membership,
|
||||
has_board_access,
|
||||
is_org_admin,
|
||||
list_accessible_board_ids,
|
||||
require_board_access,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/agents", tags=["agents"])
|
||||
|
||||
@@ -85,7 +95,13 @@ def _workspace_path(agent_name: str, workspace_root: str | None) -> str:
|
||||
return f"{root}/workspace-{_slugify(agent_name)}"
|
||||
|
||||
|
||||
async def _require_board(session: AsyncSession, board_id: UUID | str | None) -> Board:
|
||||
async def _require_board(
|
||||
session: AsyncSession,
|
||||
board_id: UUID | str | None,
|
||||
*,
|
||||
user: object | None = None,
|
||||
write: bool = False,
|
||||
) -> Board:
|
||||
if not board_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
@@ -94,6 +110,8 @@ async def _require_board(session: AsyncSession, board_id: UUID | str | None) ->
|
||||
board = await session.get(Board, board_id)
|
||||
if board is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Board not found")
|
||||
if user is not None:
|
||||
await require_board_access(session, user=user, board=board, write=write) # type: ignore[arg-type]
|
||||
return board
|
||||
|
||||
|
||||
@@ -111,6 +129,11 @@ async def _require_gateway(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||||
detail="Board gateway_id is invalid",
|
||||
)
|
||||
if gateway.organization_id != board.organization_id:
|
||||
raise HTTPException(
|
||||
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,
|
||||
@@ -206,6 +229,42 @@ async def _fetch_agent_events(
|
||||
return list(await session.exec(statement))
|
||||
|
||||
|
||||
async def _require_user_context(
|
||||
session: AsyncSession, user: User | None
|
||||
) -> OrganizationContext:
|
||||
if user is None:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
|
||||
member = await get_active_membership(session, user)
|
||||
if member is None:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
organization = await session.get(Organization, member.organization_id)
|
||||
if organization is None:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
return OrganizationContext(organization=organization, member=member)
|
||||
|
||||
|
||||
async def _require_agent_access(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
agent: Agent,
|
||||
ctx,
|
||||
write: bool,
|
||||
) -> None:
|
||||
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)
|
||||
if gateway is None or gateway.organization_id != ctx.organization.id:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
return
|
||||
|
||||
board = await session.get(Board, agent.board_id)
|
||||
if board is None or board.organization_id != ctx.organization.id:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
if not await has_board_access(session, member=ctx.member, board=board, write=write):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
|
||||
|
||||
def _record_heartbeat(session: AsyncSession, agent: Agent) -> None:
|
||||
record_activity(
|
||||
session,
|
||||
@@ -245,13 +304,28 @@ async def list_agents(
|
||||
board_id: UUID | None = Query(default=None),
|
||||
gateway_id: UUID | None = Query(default=None),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
auth: AuthContext = Depends(require_admin_auth),
|
||||
ctx=Depends(require_org_admin),
|
||||
) -> DefaultLimitOffsetPage[AgentRead]:
|
||||
main_session_keys = await _get_gateway_main_session_keys(session)
|
||||
statement = select(Agent)
|
||||
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)
|
||||
if not board_ids:
|
||||
statement = select(Agent).where(col(Agent.id).is_(None))
|
||||
else:
|
||||
base_filter = 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))
|
||||
statement = select(Agent).where(base_filter)
|
||||
if board_id is not None:
|
||||
statement = statement.where(col(Agent.board_id) == board_id)
|
||||
if gateway_id is not None:
|
||||
gateway = await session.get(Gateway, gateway_id)
|
||||
if gateway is None or gateway.organization_id != ctx.organization.id:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
statement = statement.join(Board, col(Agent.board_id) == col(Board.id)).where(
|
||||
col(Board.gateway_id) == gateway_id
|
||||
)
|
||||
@@ -269,10 +343,15 @@ async def stream_agents(
|
||||
request: Request,
|
||||
board_id: UUID | None = Query(default=None),
|
||||
since: str | None = Query(default=None),
|
||||
auth: AuthContext = Depends(require_admin_auth),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
ctx=Depends(require_org_admin),
|
||||
) -> EventSourceResponse:
|
||||
since_dt = _parse_since(since) or utcnow()
|
||||
last_seen = since_dt
|
||||
board_ids = await list_accessible_board_ids(session, member=ctx.member, write=False)
|
||||
allowed_ids = set(board_ids)
|
||||
if board_id is not None and board_id not in allowed_ids:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
|
||||
async def event_generator() -> AsyncIterator[dict[str, str]]:
|
||||
nonlocal last_seen
|
||||
@@ -280,7 +359,13 @@ async def stream_agents(
|
||||
if await request.is_disconnected():
|
||||
break
|
||||
async with async_session_maker() as session:
|
||||
agents = await _fetch_agent_events(session, board_id, last_seen)
|
||||
if board_id is not None:
|
||||
agents = await _fetch_agent_events(session, board_id, last_seen)
|
||||
elif allowed_ids:
|
||||
agents = await _fetch_agent_events(session, None, last_seen)
|
||||
agents = [agent for agent in agents if agent.board_id in allowed_ids]
|
||||
else:
|
||||
agents = []
|
||||
main_session_keys = (
|
||||
await _get_gateway_main_session_keys(session) if agents else set()
|
||||
)
|
||||
@@ -301,6 +386,10 @@ async def create_agent(
|
||||
session: AsyncSession = Depends(get_session),
|
||||
actor: ActorContext = Depends(require_admin_or_agent),
|
||||
) -> AgentRead:
|
||||
if actor.actor_type == "user":
|
||||
ctx = await _require_user_context(session, actor.user)
|
||||
if not is_org_admin(ctx.member):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
if actor.actor_type == "agent":
|
||||
if not actor.agent or not actor.agent.is_board_lead:
|
||||
raise HTTPException(
|
||||
@@ -319,7 +408,12 @@ async def create_agent(
|
||||
)
|
||||
payload = AgentCreate(**{**payload.model_dump(), "board_id": actor.agent.board_id})
|
||||
|
||||
board = await _require_board(session, payload.board_id)
|
||||
board = await _require_board(
|
||||
session,
|
||||
payload.board_id,
|
||||
user=actor.user if actor.actor_type == "user" else None,
|
||||
write=actor.actor_type == "user",
|
||||
)
|
||||
gateway, client_config = await _require_gateway(session, board)
|
||||
data = payload.model_dump()
|
||||
requested_name = (data.get("name") or "").strip()
|
||||
@@ -436,11 +530,12 @@ async def create_agent(
|
||||
async def get_agent(
|
||||
agent_id: str,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
auth: AuthContext = Depends(require_admin_auth),
|
||||
ctx=Depends(require_org_admin),
|
||||
) -> AgentRead:
|
||||
agent = await session.get(Agent, agent_id)
|
||||
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)
|
||||
|
||||
@@ -451,18 +546,28 @@ async def update_agent(
|
||||
payload: AgentUpdate,
|
||||
force: bool = False,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
auth: AuthContext = Depends(require_admin_auth),
|
||||
auth: AuthContext = Depends(get_auth_context),
|
||||
ctx=Depends(require_org_admin),
|
||||
) -> AgentRead:
|
||||
agent = await session.get(Agent, agent_id)
|
||||
if agent is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
await _require_agent_access(session, agent=agent, ctx=ctx, write=True)
|
||||
updates = payload.model_dump(exclude_unset=True)
|
||||
make_main = updates.pop("is_gateway_main", None)
|
||||
if make_main is True and not is_org_admin(ctx.member):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
if "status" in updates:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="status is controlled by agent heartbeat",
|
||||
)
|
||||
if "board_id" in updates and updates["board_id"] is not None:
|
||||
new_board = await _require_board(session, updates["board_id"])
|
||||
if new_board.organization_id != ctx.organization.id:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
if not await has_board_access(session, member=ctx.member, board=new_board, write=True):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
if not updates and not 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)
|
||||
@@ -628,6 +733,11 @@ async def heartbeat_agent(
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
if actor.actor_type == "agent" and actor.agent and actor.agent.id != agent.id:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
if actor.actor_type == "user":
|
||||
ctx = await _require_user_context(session, actor.user)
|
||||
if not is_org_admin(ctx.member):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
await _require_agent_access(session, agent=agent, ctx=ctx, write=True)
|
||||
if payload.status:
|
||||
agent.status = payload.status
|
||||
elif agent.status == "provisioning":
|
||||
@@ -664,7 +774,16 @@ async def heartbeat_or_create_agent(
|
||||
if agent is None:
|
||||
if actor.actor_type == "agent":
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
|
||||
board = await _require_board(session, payload.board_id)
|
||||
if actor.actor_type == "user":
|
||||
ctx = await _require_user_context(session, actor.user)
|
||||
if not is_org_admin(ctx.member):
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
board = await _require_board(
|
||||
session,
|
||||
payload.board_id,
|
||||
user=actor.user,
|
||||
write=True,
|
||||
)
|
||||
gateway, client_config = await _require_gateway(session, board)
|
||||
agent = Agent(
|
||||
name=payload.name,
|
||||
@@ -724,6 +843,9 @@ async def heartbeat_or_create_agent(
|
||||
except Exception as exc: # pragma: no cover - unexpected provisioning errors
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
await session.commit()
|
||||
elif actor.actor_type == "user":
|
||||
ctx = await _require_user_context(session, actor.user)
|
||||
await _require_agent_access(session, agent=agent, ctx=ctx, write=True)
|
||||
elif actor.actor_type == "agent" and actor.agent and actor.agent.id != agent.id:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
|
||||
elif agent.agent_token_hash is None and actor.actor_type == "user":
|
||||
@@ -737,7 +859,12 @@ async def heartbeat_or_create_agent(
|
||||
await session.commit()
|
||||
await session.refresh(agent)
|
||||
try:
|
||||
board = await _require_board(session, str(agent.board_id) if agent.board_id else None)
|
||||
board = await _require_board(
|
||||
session,
|
||||
str(agent.board_id) if agent.board_id else None,
|
||||
user=actor.user if actor.actor_type == "user" else None,
|
||||
write=actor.actor_type == "user",
|
||||
)
|
||||
gateway, client_config = await _require_gateway(session, board)
|
||||
await provision_agent(agent, board, gateway, raw_token, actor.user, action="provision")
|
||||
await _send_wakeup_message(agent, client_config, verb="provisioned")
|
||||
@@ -767,7 +894,12 @@ async def heartbeat_or_create_agent(
|
||||
_record_instruction_failure(session, agent, str(exc), "provision")
|
||||
await session.commit()
|
||||
elif not agent.openclaw_session_id:
|
||||
board = await _require_board(session, str(agent.board_id) if agent.board_id else None)
|
||||
board = await _require_board(
|
||||
session,
|
||||
str(agent.board_id) if agent.board_id else None,
|
||||
user=actor.user if actor.actor_type == "user" else None,
|
||||
write=actor.actor_type == "user",
|
||||
)
|
||||
gateway, client_config = await _require_gateway(session, board)
|
||||
session_key, session_error = await _ensure_gateway_session(agent.name, client_config)
|
||||
agent.openclaw_session_id = session_key
|
||||
@@ -804,11 +936,12 @@ async def heartbeat_or_create_agent(
|
||||
async def delete_agent(
|
||||
agent_id: str,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
auth: AuthContext = Depends(require_admin_auth),
|
||||
ctx=Depends(require_org_admin),
|
||||
) -> OkResponse:
|
||||
agent = await session.get(Agent, agent_id)
|
||||
if agent is None:
|
||||
return OkResponse()
|
||||
await _require_agent_access(session, agent=agent, ctx=ctx, write=True)
|
||||
|
||||
board = await _require_board(session, str(agent.board_id) if agent.board_id else None)
|
||||
gateway, client_config = await _require_gateway(session, board)
|
||||
|
||||
Reference in New Issue
Block a user