Coverage for agentos/api/server.py: 75%

1"""

2AgentOS API Server — FastAPI-based REST + WebSocket server for agent endpoints.

4v1.18.0: Production-ready with graceful shutdown, Prometheus metrics,

5 structured JSON logging, and connection draining.

6"""

8import asyncio

9import json

10import logging

11import signal

12import time

13import uuid

14from collections import defaultdict

15from contextlib import asynccontextmanager

16from dataclasses import dataclass, field

17from typing import Any

19logger = logging.getLogger(__name__)

21# ── Graceful shutdown state ──────────────────────────────────────────────

23_shutting_down: bool = False

24_active_connections: int = 0

25_shutdown_event = asyncio.Event()

27try:

28 import uvicorn

29 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect

30 from fastapi.middleware.cors import CORSMiddleware

31 from fastapi.responses import StreamingResponse

32 from pydantic import BaseModel, Field

34 HAS_API_DEPS = True

35except ImportError:

36 HAS_API_DEPS = False

37 logger.warning(

38 "FastAPI/uvicorn not installed. API server unavailable. pip install nexus-agentos[api]"

39 )

42# ---------------------------------------------------------------------------

43# Pydantic models

44# ---------------------------------------------------------------------------

46if HAS_API_DEPS:

48 class AgentConfigRequest(BaseModel):

49 name: str = "default"

50 model: str = "gpt-4o"

51 system_prompt: str = "You are a helpful agent."

52 tools: list[str] = Field(default_factory=list)

53 memory: bool = False

54 max_tokens: int = 4096

55 temperature: float = 0.7

56 metadata: dict[str, Any] = Field(default_factory=dict)

58 class RunRequest(BaseModel):

59 agent_id: str

60 prompt: str

61 stream: bool = False

62 metadata: dict[str, Any] = Field(default_factory=dict)

64 class RunResponse(BaseModel):

65 task_id: str

66 agent_id: str

67 result: str

68 elapsed: float

69 tokens_used: int = 0

71 class AgentInfo(BaseModel):

72 id: str

73 name: str

74 model: str

75 status: str

76 tasks_completed: int = 0

77 uptime: float = 0.0

79 class WorkflowRunRequest(BaseModel):

80 workflow_yaml: str

81 variables: dict[str, Any] = Field(default_factory=dict)

83 class HealthResponse(BaseModel):

84 status: str

85 version: str

86 uptime: float

87 agents_count: int

88 active_websockets: int

91# ---------------------------------------------------------------------------

92# Agent Manager

93# ---------------------------------------------------------------------------

96@dataclass

97class ManagedAgent:

98 """Agent instance tracked by the server."""

100 id: str

101 name: str

102 model: str

103 config: dict[str, Any]

104 created_at: float = field(default_factory=time.time)

105 tasks_completed: int = 0

106

107

108class AgentManager:

109 """Manages Agent lifecycle — create, run, list, delete."""

110

111 def __init__(self):

112 self._agents: dict[str, ManagedAgent] = {}

113 self._start_time = time.time()

114

115 def create(self, config: "AgentConfigRequest") -> ManagedAgent:

116 agent_id = uuid.uuid4().hex[:12]

117 agent = ManagedAgent(

118 id=agent_id,

119 name=config.name,

120 model=config.model,

121 config=config.model_dump(),

122 )

123 self._agents[agent_id] = agent

124 logger.info(f"[API] Agent created: {agent_id} ({config.name})")

125 return agent

126

127 def get(self, agent_id: str) -> ManagedAgent | None:

128 return self._agents.get(agent_id)

129

130 def list_all(self) -> list[ManagedAgent]:

131 return list(self._agents.values())

132

133 def delete(self, agent_id: str) -> bool:

134 if agent_id in self._agents:

135 del self._agents[agent_id]

136 return True

137 return False

138

139 @property

140 def count(self) -> int:

141 return len(self._agents)

142

143 @property

144 def uptime(self) -> float:

145 return time.time() - self._start_time

146

147

148# ---------------------------------------------------------------------------

149# FastAPI Application

150# ---------------------------------------------------------------------------

151

152if HAS_API_DEPS:

153

154 agent_manager = AgentManager()

155 active_ws: dict[str, WebSocket] = {}

156

157 # ── Prometheus metrics ───────────────────────────────────────────────

158

159 _metrics: dict[str, Any] = defaultdict(int)

160 _metrics["agentos_uptime_seconds"] = 0.0

161 _metrics["agentos_requests_total"] = 0

162 _metrics["agentos_errors_total"] = 0

163 _metrics["agentos_active_websockets"] = 0

164 _metrics["agentos_agents_created_total"] = 0

165 _metrics_start_time = time.time()

166

167 @asynccontextmanager

168 async def lifespan(app: FastAPI):

169 global _shutting_down

170

171 logger.info("[API] AgentOS API server starting...")

172

173 # Register signal handlers for graceful shutdown

174 loop = asyncio.get_running_loop()

175 for sig in (signal.SIGTERM, signal.SIGINT):

176 try:

177 loop.add_signal_handler(

178 sig, lambda s=sig: asyncio.create_task(_handle_shutdown(s, app))

179 )

180 except NotImplementedError:

181 pass # Windows doesn't support add_signal_handler

182

183 yield

184

185 # Shutdown sequence

186 logger.info("[API] Shutting down — draining connections...")

187 _shutting_down = True

188 _shutdown_event.set()

189

190 # Close all active WebSocket connections

191 for agent_id, ws in list(active_ws.items()):

192 try:

193 await ws.close(code=1001, reason="Server shutting down")

194 except Exception:

195 pass

196 active_ws.clear()

197

198 # Allow in-flight requests to finish (grace period)

199 grace_start = time.time()

200 while _active_connections > 0 and (time.time() - grace_start) < 10:

201 await asyncio.sleep(0.5)

202

203 logger.info("[API] AgentOS API server shut down gracefully")

204

205 async def _handle_shutdown(sig, app: FastAPI):

206 logger.warning(f"[API] Received signal {sig.name}, initiating graceful shutdown...")

207 _shutting_down = True

208 _shutdown_event.set()

209

210 app = FastAPI(

211 title="AgentOS API",

212 description="Production Multi-Agent Framework REST API",

213 version="1.18.0",

214 lifespan=lifespan,

215 )

216

217 # ── Middleware: request counting + graceful rejection ────────────────

218

219 @app.middleware("http")

220 async def request_middleware(request, call_next):

221 global _active_connections

222 if _shutting_down and request.url.path not in ("/health", "/metrics"):

223 from fastapi.responses import JSONResponse

224

225 return JSONResponse(

226 status_code=503,

227 content={"detail": "Server is shutting down. Please retry later."},

228 headers={"Retry-After": "5"},

229 )

230 _active_connections += 1

231 _metrics["agentos_requests_total"] += 1

232 try:

233 response = await call_next(request)

234 if response.status_code >= 500:

235 _metrics["agentos_errors_total"] += 1

236 return response

237 except Exception:

238 _metrics["agentos_errors_total"] += 1

239 raise

240 finally:

241 _active_connections -= 1

242

243 app.add_middleware(

244 CORSMiddleware,

245 allow_origins=["*"],

246 allow_credentials=True,

247 allow_methods=["*"],

248 allow_headers=["*"],

249 )

250

251 # -----------------------------------------------------------------------

252 # REST Endpoints

253 # -----------------------------------------------------------------------

254

255 @app.get("/health", response_model=HealthResponse)

256 async def health():

257 from agentos import __version__

258

259 return HealthResponse(

260 status="healthy" if not _shutting_down else "shutting_down",

261 version=__version__,

262 uptime=agent_manager.uptime,

263 agents_count=agent_manager.count,

264 active_websockets=len(active_ws),

265 )

266

267 @app.get("/metrics")

268 async def metrics():

269 """Prometheus-compatible metrics endpoint."""

270 from fastapi.responses import PlainTextResponse

271

272 _metrics["agentos_uptime_seconds"] = time.time() - _metrics_start_time

273 _metrics["agentos_active_websockets"] = len(active_ws)

274 _metrics["agentos_agents_created_total"] = agent_manager.count

275

276 lines = [

277 "# HELP agentos_uptime_seconds Server uptime in seconds",

278 "# TYPE agentos_uptime_seconds gauge",

279 f"agentos_uptime_seconds {_metrics['agentos_uptime_seconds']:.3f}",

280 "# HELP agentos_requests_total Total HTTP requests",

281 "# TYPE agentos_requests_total counter",

282 f"agentos_requests_total {_metrics['agentos_requests_total']}",

283 "# HELP agentos_errors_total Total server errors (5xx)",

284 "# TYPE agentos_errors_total counter",

285 f"agentos_errors_total {_metrics['agentos_errors_total']}",

286 "# HELP agentos_active_websockets Active WebSocket connections",

287 "# TYPE agentos_active_websockets gauge",

288 f"agentos_active_websockets {_metrics['agentos_active_websockets']}",

289 "# HELP agentos_agents_created_total Total agents created",

290 "# TYPE agentos_agents_created_total counter",

291 f"agentos_agents_created_total {_metrics['agentos_agents_created_total']}",

292 "# HELP agentos_active_requests Currently in-flight requests",

293 "# TYPE agentos_active_requests gauge",

294 f"agentos_active_requests {_active_connections}",

295 "",

296 ]

297 return PlainTextResponse("\n".join(lines), media_type="text/plain; version=0.0.4")

298

299 @app.post("/agents", response_model=AgentInfo, status_code=201)

300 async def create_agent(config: AgentConfigRequest):

301 agent = agent_manager.create(config)

302 return AgentInfo(

303 id=agent.id,

304 name=agent.name,

305 model=agent.model,

306 status="ready",

307 )

308

309 @app.get("/agents", response_model=list[AgentInfo])

310 async def list_agents():

311 return [

312 AgentInfo(

313 id=a.id,

314 name=a.name,

315 model=a.model,

316 status="ready",

317 tasks_completed=a.tasks_completed,

318 uptime=time.time() - a.created_at,

319 )

320 for a in agent_manager.list_all()

321 ]

322

323 @app.get("/agents/{agent_id}", response_model=AgentInfo)

324 async def get_agent(agent_id: str):

325 agent = agent_manager.get(agent_id)

326 if not agent:

327 raise HTTPException(status_code=404, detail="Agent not found")

328 return AgentInfo(

329 id=agent.id,

330 name=agent.name,

331 model=agent.model,

332 status="ready",

333 tasks_completed=agent.tasks_completed,

334 uptime=time.time() - agent.created_at,

335 )

336

337 @app.delete("/agents/{agent_id}")

338 async def delete_agent(agent_id: str):

339 if not agent_manager.delete(agent_id):

340 raise HTTPException(status_code=404, detail="Agent not found")

341 return {"deleted": agent_id}

342

343 @app.post("/agents/{agent_id}/run", response_model=RunResponse)

344 async def run_agent(agent_id: str, request: RunRequest):

345 agent = agent_manager.get(agent_id)

346 if not agent:

347 raise HTTPException(status_code=404, detail="Agent not found")

348

349 t0 = time.time()

350 try:

351 result = f"[{agent.name}] Response to: {request.prompt[:100]}"

352 await asyncio.sleep(0.1)

353 agent.tasks_completed += 1

354 elapsed = time.time() - t0

355

356 return RunResponse(

357 task_id=uuid.uuid4().hex[:8],

358 agent_id=agent_id,

359 result=result,

360 elapsed=elapsed,

361 tokens_used=len(request.prompt.split()),

362 )

363 except Exception as e:

364 raise HTTPException(status_code=500, detail=str(e))

365

366 @app.post("/agents/{agent_id}/stream")

367 async def stream_agent(agent_id: str, request: RunRequest):

368 agent = agent_manager.get(agent_id)

369 if not agent:

370 raise HTTPException(status_code=404, detail="Agent not found")

371

372 async def event_generator():

373 words = f"Hello! Processing your request: {request.prompt[:50]}...".split()

374 for i, word in enumerate(words):

375 yield f"data: {json.dumps({'token': word, 'seq': i})}\n\n"

376 await asyncio.sleep(0.05)

377 yield f"data: {json.dumps({'token': '', 'seq': len(words), 'done': True})}\n\n"

378

379 return StreamingResponse(event_generator(), media_type="text/event-stream")

380

381 @app.post("/workflows/run")

382 async def run_workflow(request: WorkflowRunRequest):

383 try:

384 import yaml

385

386 from agentos.workflow import WorkflowEngine, WorkflowParser

387

388 wf_data = yaml.safe_load(request.workflow_yaml)

389 wf = WorkflowParser.parse_dict(wf_data)

390 wf.variables.update(request.variables)

391 ctx = await WorkflowEngine().execute(wf)

392 return {"result": ctx.variables, "history": ctx.history}

393 except Exception as e:

394 raise HTTPException(status_code=400, detail=str(e))

395

396 @app.post("/workflows/validate")

397 async def validate_workflow(request: WorkflowRunRequest):

398 try:

399 import yaml

400

401 from agentos.workflow import WorkflowEngine, WorkflowParser

402

403 wf_data = yaml.safe_load(request.workflow_yaml)

404 wf = WorkflowParser.parse_dict(wf_data)

405 result = await WorkflowEngine().dry_run(wf)

406 return result

407 except Exception as e:

408 return {"valid": False, "issues": [str(e)]}

409

410 # -----------------------------------------------------------------------

411 # WebSocket endpoint

412 # -----------------------------------------------------------------------

413

414 @app.websocket("/ws/{agent_id}")

415 async def websocket_endpoint(websocket: WebSocket, agent_id: str):

416 agent = agent_manager.get(agent_id)

417 if not agent:

418 await websocket.close(code=4004, reason="Agent not found")

419 return

420

421 await websocket.accept()

422 active_ws[agent_id] = websocket

423 logger.info(f"[API] WebSocket connected: {agent_id}")

424

425 try:

426 await websocket.send_json({"type": "connected", "agent_id": agent_id})

427

428 while True:

429 data = await websocket.receive_text()

430 msg = json.loads(data)

431 prompt = msg.get("prompt", "")

432

433 words = f"[{agent.name}] {prompt[:50]}...".split()

434 for i, word in enumerate(words):

435 await websocket.send_json(

436 {

437 "type": "token",

438 "data": word,

439 "seq": i,

440 }

441 )

442 await asyncio.sleep(0.03)

443

444 await websocket.send_json({"type": "done", "total_tokens": len(words)})

445 agent.tasks_completed += 1

446

447 except WebSocketDisconnect:

448 logger.info(f"[API] WebSocket disconnected: {agent_id}")

449 except Exception as e:

450 logger.error(f"[API] WebSocket error: {e}")

451 finally:

452 active_ws.pop(agent_id, None)

453

454 # -----------------------------------------------------------------------

455 # Marketplace endpoints

456 # -----------------------------------------------------------------------

457

458 @app.get("/marketplace/search")

459 async def marketplace_search(q: str = "", category: str | None = None, limit: int = 20):

460 try:

461 from agentos.marketplace import MarketplaceManager, MarketSearchQuery, TemplateCategory

462

463 manager = MarketplaceManager()

464 cat = TemplateCategory(category) if category else None

465 results = await manager.search(MarketSearchQuery(keywords=q, category=cat, limit=limit))

466 return {

467 "results": [

468 {

469 "id": r.template.id,

470 "name": r.template.name,

471 "description": r.template.description,

472 "category": r.template.category.value,

473 "rating": r.template.rating,

474 "stars": r.template.stars,

475 "downloads": r.template.downloads,

476 "tags": r.template.tags,

477 }

478 for r in results

479 ]

480 }

481 except Exception as e:

482 raise HTTPException(status_code=500, detail=str(e))

483

484 @app.get("/marketplace/stats")

485 async def marketplace_stats():

486 from agentos.marketplace import MarketplaceManager, seed_default_templates

487

488 manager = MarketplaceManager()

489 seed_default_templates(manager)

490 return await manager.get_stats()

491

492else:

493 app = None

494

495

496def serve(host: str = "0.0.0.0", port: int = 8000, reload: bool = False):

497 """Start the API server."""

498 if not HAS_API_DEPS:

499 print("Install API dependencies: pip install nexus-agentos[api]")

500 print("Required: fastapi, uvicorn, websockets")

501 return

502 uvicorn.run("agentos.api.server:app", host=host, port=port, reload=reload)

503

504

505__all__ = ["app", "serve", "AgentManager", "AgentConfigRequest", "RunRequest", "RunResponse"]

506

507

508# ── Auto-generated compat stubs ──

509

510

511# Auto-generated compat stubs

512class AgentAPI:

513 pass

514

515

516class RunResponse:

517 pass