From fd3b3637ac5fa57d13805a8c0f0a56661176d21d Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 6 May 2026 15:06:37 -0700 Subject: [PATCH] feat(cloud-agents): add /v2/agent-workers register/list/revoke endpoints --- .../worker_endpoints.py | 309 ++++++++++++++++++ 1 file changed, 309 insertions(+) create mode 100644 litellm/proxy/agent_settings_endpoints/worker_endpoints.py diff --git a/litellm/proxy/agent_settings_endpoints/worker_endpoints.py b/litellm/proxy/agent_settings_endpoints/worker_endpoints.py new file mode 100644 index 00000000000..6a9ff9df442 --- /dev/null +++ b/litellm/proxy/agent_settings_endpoints/worker_endpoints.py @@ -0,0 +1,309 @@ +""" +`/v2/agent-workers` endpoints (LIT-2891 / Screen 3). + +Self-hosted worker registration flow: + +1. Operator clicks "Add Machine" → UI calls `POST /v2/agent-workers/pair-token`. + We mint a 32-byte urlsafe token, persist only its sha256, return the raw + token + an install one-liner. TTL = 15 min, single use. +2. Operator runs the install command on the worker box. The worker calls + `POST /v2/agent-workers/register` with the raw token + its hostname. We + re-hash, atomically mark the pairing-token row as consumed, mint a long- + lived worker JWT, persist its sha256, and return the raw JWT to the worker. + Both raw values are returned exactly once and never persisted. +3. The dashboard lists workers via `GET /v2/agent-workers` and revokes them + via `DELETE /v2/agent-workers/{id}`. + +The actual long-poll / hydrate transport is owned by Epic B2 — this module +only owns the auth handshake and CRUD surface. +""" + +import secrets as _stdlib_secrets +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from fastapi import APIRouter, Depends, HTTPException, Request + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.agent_settings_endpoints.pair_tokens import ( + build_install_command, + hash_pair_token, + hash_worker_jwt, + is_expired, + issue_pair_token, +) +from litellm.proxy.agent_settings_endpoints.types import ( + AgentWorkerListResponse, + AgentWorkerRegisterRequest, + AgentWorkerRegisterResponse, + AgentWorkerResponse, + PairTokenResponse, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +def _resolve_team_id(user_api_key_dict: UserAPIKeyAuth) -> str: + """Pick the team to scope this request to. Raise 400 if missing.""" + team_id = user_api_key_dict.team_id or (user_api_key_dict.metadata or {}).get( + "team_id" + ) + if not team_id: + raise HTTPException( + status_code=400, + detail=( + "Cloud Agent workers are scoped to a team. Pick a team from the " + "header switcher and try again." + ), + ) + return team_id + + +def _row_to_worker_response(row: Dict[str, Any]) -> AgentWorkerResponse: + """Map a Prisma worker row to the public response shape.""" + last_seen = row.get("last_seen_at") + return AgentWorkerResponse( + id=row["id"], + hostname=row.get("hostname") or "", + status=row.get("status") or "offline", + last_seen_at=str(last_seen) if last_seen is not None else None, + cpu_pct=row.get("cpu_pct"), + mem_gb=row.get("mem_gb"), + active_sessions=int(row.get("active_sessions") or 0), + ) + + +def _resolve_proxy_url(request: Request) -> str: + """Best-effort guess at the public proxy URL for the install one-liner. + + The UI shows this URL inside the `--proxy` flag of the install command, + so it has to point at the host the worker box can reach. We start from + the request's `Host` header (works for the common case of the dashboard + being served from the same proxy) and fall back to whatever forwarded + headers the operator's reverse proxy set. If the UI overrides this in + the future, just stop calling this helper and pass the URL in. + """ + forwarded_proto = request.headers.get("x-forwarded-proto") + forwarded_host = request.headers.get("x-forwarded-host") + host = forwarded_host or request.headers.get("host") or "localhost:4000" + scheme = forwarded_proto or request.url.scheme or "https" + return f"{scheme}://{host}" + + +@router.get( + "/v2/agent-workers", + dependencies=[Depends(user_api_key_auth)], + response_model=AgentWorkerListResponse, + tags=["cloud agents"], +) +async def list_agent_workers( + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> AgentWorkerListResponse: + """List the team's self-hosted workers.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + team_id = _resolve_team_id(user_api_key_dict) + rows = await prisma_client.db.litellm_agentworker.find_many( + where={"team_id": team_id}, + order={"created_at": "desc"}, + ) + workers: List[AgentWorkerResponse] = [ + _row_to_worker_response(dict(r) if not isinstance(r, dict) else r) for r in rows + ] + return AgentWorkerListResponse(workers=workers) + + +@router.post( + "/v2/agent-workers/pair-token", + dependencies=[Depends(user_api_key_auth)], + response_model=PairTokenResponse, + tags=["cloud agents"], +) +async def create_pair_token( + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> PairTokenResponse: + """Mint a single-use 15-minute pairing token. Raw token returned ONCE.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + team_id = _resolve_team_id(user_api_key_dict) + issued = issue_pair_token() + + try: + await prisma_client.db.litellm_agentworkerpairingtoken.create( + data={ + "token_hash": issued.token_hash, + "team_id": team_id, + "created_by": user_api_key_dict.user_id or "unknown", + "expires_at": issued.expires_at, + } + ) + except Exception as exc: + verbose_proxy_logger.exception( + "Failed to persist pair token for team=%s: %s", team_id, exc + ) + raise HTTPException(status_code=500, detail="Failed to mint pair token.") + + install_command = build_install_command( + proxy_url=_resolve_proxy_url(request), + raw_token=issued.raw_token, + ) + return PairTokenResponse( + token=issued.raw_token, + expires_at=issued.expires_at.isoformat(), + install_command=install_command, + ) + + +@router.post( + "/v2/agent-workers/register", + response_model=AgentWorkerRegisterResponse, + tags=["cloud agents"], +) +async def register_agent_worker( + request: Request, + body: AgentWorkerRegisterRequest, +) -> AgentWorkerRegisterResponse: + """Worker exchanges its pair token for a long-lived JWT. + + NOTE: this endpoint is NOT behind `user_api_key_auth` — the worker has + no API key yet, the pair token *is* its proof of authorization. We + enforce single-use atomically by checking `used_at IS NULL` on update. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + token_hash = hash_pair_token(body.pair_token) + pair_row = await prisma_client.db.litellm_agentworkerpairingtoken.find_unique( + where={"token_hash": token_hash} + ) + if pair_row is None: + raise HTTPException(status_code=401, detail="Invalid pairing token.") + + pair = dict(pair_row) if not isinstance(pair_row, dict) else pair_row + if pair.get("used_at") is not None: + raise HTTPException( + status_code=401, detail="Pairing token has already been used." + ) + if is_expired(pair.get("expires_at")): + raise HTTPException(status_code=401, detail="Pairing token has expired.") + + team_id = pair["team_id"] + + # Atomically consume the pair token. `update_many` with the `used_at: None` + # filter is the closest Prisma gives us to a CAS — if a second register + # raced us, the second one matches 0 rows and we 401. + consumed = await prisma_client.db.litellm_agentworkerpairingtoken.update_many( + where={"token_hash": token_hash, "used_at": None}, + data={"used_at": datetime.now(timezone.utc)}, + ) + if not consumed: + raise HTTPException( + status_code=401, detail="Pairing token has already been used." + ) + + # Mint the worker JWT. This is a short opaque urlsafe string; the daemon + # presents it on every long-poll. We persist only its sha256. + raw_jwt = _stdlib_secrets.token_urlsafe(48) + jwt_hash = hash_worker_jwt(raw_jwt) + + try: + worker = await prisma_client.db.litellm_agentworker.create( + data={ + "team_id": team_id, + "hostname": body.hostname, + "status": "online", + "last_seen_at": datetime.now(timezone.utc), + "active_sessions": 0, + "worker_jwt_hash": jwt_hash, + } + ) + except Exception as exc: + verbose_proxy_logger.exception( + "Failed to create agent worker hostname=%s team=%s: %s", + body.hostname, + team_id, + exc, + ) + raise HTTPException(status_code=500, detail="Failed to register worker.") + + worker_dict = dict(worker) if not isinstance(worker, dict) else worker + return AgentWorkerRegisterResponse( + worker_id=worker_dict["id"], + worker_jwt=raw_jwt, + ) + + +@router.delete( + "/v2/agent-workers/{worker_id}", + dependencies=[Depends(user_api_key_auth)], + tags=["cloud agents"], +) +async def delete_agent_worker( + request: Request, + worker_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> Dict[str, Any]: + """Revoke a worker. Idempotent — 404 if not found, no-op if already gone.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + team_id = _resolve_team_id(user_api_key_dict) + existing = await prisma_client.db.litellm_agentworker.find_unique( + where={"id": worker_id} + ) + if existing is None: + raise HTTPException(status_code=404, detail="Worker not found.") + + existing_dict = dict(existing) if not isinstance(existing, dict) else existing + if existing_dict.get("team_id") != team_id: + # Don't leak existence cross-team; same status code as not-found. + raise HTTPException(status_code=404, detail="Worker not found.") + + await prisma_client.db.litellm_agentworker.delete(where={"id": worker_id}) + return {"deleted": True, "id": worker_id} + + +# Re-exported for tests + B2 hydrate path: given a raw worker JWT, look up the +# worker row. Centralized here so the auth scheme lives in exactly one place. +async def find_worker_by_jwt( + prisma_client: Any, raw_jwt: str +) -> Optional[Dict[str, Any]]: + """Look up a worker row by its (raw) JWT. Returns None if not found. + + We hash and `find_first` instead of indexing on `worker_jwt_hash` directly + because the column is currently unindexed — fine for v1 since the team + pool is tiny. Add an index when we have >100 workers per team. + """ + jwt_hash = hash_worker_jwt(raw_jwt) + worker = await prisma_client.db.litellm_agentworker.find_first( + where={"worker_jwt_hash": jwt_hash} + ) + if worker is None: + return None + return dict(worker) if not isinstance(worker, dict) else worker