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

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, Dict, List, Optional

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 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Query

29 from fastapi.middleware.cors import CORSMiddleware

30 from fastapi.responses import StreamingResponse, JSONResponse

31 from pydantic import BaseModel, Field

32 import uvicorn

33 HAS_API_DEPS = True

34except ImportError:

35 HAS_API_DEPS = False

36 logger.warning("FastAPI/uvicorn not installed. API server unavailable. pip install nexus-agentos[api]")

39# ---------------------------------------------------------------------------

40# Pydantic models

41# ---------------------------------------------------------------------------

43if HAS_API_DEPS:

45 class AgentConfigRequest(BaseModel):

46 name: str = "default"

47 model: str = "gpt-4o"

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

49 tools: List[str] = Field(default_factory=list)

50 memory: bool = False

51 max_tokens: int = 4096

52 temperature: float = 0.7

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

55 class RunRequest(BaseModel):

56 agent_id: str

57 prompt: str

58 stream: bool = False

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

61 class RunResponse(BaseModel):

62 task_id: str

63 agent_id: str

64 result: str

65 elapsed: float

66 tokens_used: int = 0

68 class AgentInfo(BaseModel):

69 id: str

70 name: str

71 model: str

72 status: str

73 tasks_completed: int = 0

74 uptime: float = 0.0

76 class WorkflowRunRequest(BaseModel):

77 workflow_yaml: str

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

80 class HealthResponse(BaseModel):

81 status: str

82 version: str

83 uptime: float

84 agents_count: int

85 active_websockets: int

88# ---------------------------------------------------------------------------

89# Agent Manager

90# ---------------------------------------------------------------------------

92@dataclass

93class ManagedAgent:

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

95 id: str

96 name: str

97 model: str

98 config: Dict[str, Any]

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

100 tasks_completed: int = 0

101

102

103class AgentManager:

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

105

106 def __init__(self):

107 self._agents: Dict[str, ManagedAgent] = {}

108 self._start_time = time.time()

109

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

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

112 agent = ManagedAgent(

113 id=agent_id,

114 name=config.name,

115 model=config.model,

116 config=config.model_dump(),

117 )

118 self._agents[agent_id] = agent

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

120 return agent

121

122 def get(self, agent_id: str) -> Optional[ManagedAgent]:

123 return self._agents.get(agent_id)

124

125 def list_all(self) -> List[ManagedAgent]:

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

127

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

129 if agent_id in self._agents:

130 del self._agents[agent_id]

131 return True

132 return False

133

134 @property

135 def count(self) -> int:

136 return len(self._agents)

137

138 @property

139 def uptime(self) -> float:

140 return time.time() - self._start_time

141

142

143# ---------------------------------------------------------------------------

144# FastAPI Application

145# ---------------------------------------------------------------------------

146

147if HAS_API_DEPS:

148

149 agent_manager = AgentManager()

150 active_ws: Dict[str, WebSocket] = {}

151

152 # ── Prometheus metrics ───────────────────────────────────────────────

153

154 _metrics: Dict[str, Any] = defaultdict(int)

155 _metrics["agentos_uptime_seconds"] = 0.0

156 _metrics["agentos_requests_total"] = 0

157 _metrics["agentos_errors_total"] = 0

158 _metrics["agentos_active_websockets"] = 0

159 _metrics["agentos_agents_created_total"] = 0

160 _metrics_start_time = time.time()

161

162 @asynccontextmanager

163 async def lifespan(app: FastAPI):

164 global _shutting_down

165

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

167

168 # Register signal handlers for graceful shutdown

169 loop = asyncio.get_running_loop()

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

171 try:

172 loop.add_signal_handler(

173 sig,

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

175 )

176 except NotImplementedError:

177 pass # Windows doesn't support add_signal_handler

178

179 yield

180

181 # Shutdown sequence

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

183 _shutting_down = True

184 _shutdown_event.set()

185

186 # Close all active WebSocket connections

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

188 try:

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

190 except Exception:

191 pass

192 active_ws.clear()

193

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

195 grace_start = time.time()

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

197 await asyncio.sleep(0.5)

198

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

200

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

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

203 _shutting_down = True

204 _shutdown_event.set()

205

206 app = FastAPI(

207 title="AgentOS API",

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

209 version="1.18.0",

210 lifespan=lifespan,

211 )

212

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

214

215 @app.middleware("http")

216 async def request_middleware(request, call_next):

217 global _active_connections

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

219 from fastapi.responses import JSONResponse

220 return JSONResponse(

221 status_code=503,

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

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

224 )

225 _active_connections += 1

226 _metrics["agentos_requests_total"] += 1

227 try:

228 response = await call_next(request)

229 if response.status_code >= 500:

230 _metrics["agentos_errors_total"] += 1

231 return response

232 except Exception:

233 _metrics["agentos_errors_total"] += 1

234 raise

235 finally:

236 _active_connections -= 1

237

238 app.add_middleware(

239 CORSMiddleware,

240 allow_origins=["*"],

241 allow_credentials=True,

242 allow_methods=["*"],

243 allow_headers=["*"],

244 )

245

246 # -----------------------------------------------------------------------

247 # REST Endpoints

248 # -----------------------------------------------------------------------

249

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

251 async def health():

252 from agentos import __version__

253 return HealthResponse(

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

255 version=__version__,

256 uptime=agent_manager.uptime,

257 agents_count=agent_manager.count,

258 active_websockets=len(active_ws),

259 )

260

261 @app.get("/metrics")

262 async def metrics():

263 """Prometheus-compatible metrics endpoint."""

264 from fastapi.responses import PlainTextResponse

265

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

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

268 _metrics["agentos_agents_created_total"] = agent_manager.count

269

270 lines = [

271 "# HELP agentos_uptime_seconds Server uptime in seconds",

272 "# TYPE agentos_uptime_seconds gauge",

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

274 "# HELP agentos_requests_total Total HTTP requests",

275 "# TYPE agentos_requests_total counter",

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

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

278 "# TYPE agentos_errors_total counter",

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

280 "# HELP agentos_active_websockets Active WebSocket connections",

281 "# TYPE agentos_active_websockets gauge",

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

283 "# HELP agentos_agents_created_total Total agents created",

284 "# TYPE agentos_agents_created_total counter",

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

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

287 "# TYPE agentos_active_requests gauge",

288 f"agentos_active_requests {_active_connections}",

289 "",

290 ]

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

292

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

294 async def create_agent(config: AgentConfigRequest):

295 agent = agent_manager.create(config)

296 return AgentInfo(

297 id=agent.id,

298 name=agent.name,

299 model=agent.model,

300 status="ready",

301 )

302

303 @app.get("/agents", response_model=List[AgentInfo])

304 async def list_agents():

305 return [

306 AgentInfo(

307 id=a.id, name=a.name, model=a.model,

308 status="ready", tasks_completed=a.tasks_completed,

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

310 )

311 for a in agent_manager.list_all()

312 ]

313

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

315 async def get_agent(agent_id: str):

316 agent = agent_manager.get(agent_id)

317 if not agent:

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

319 return AgentInfo(

320 id=agent.id, name=agent.name, model=agent.model,

321 status="ready", tasks_completed=agent.tasks_completed,

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

323 )

324

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

326 async def delete_agent(agent_id: str):

327 if not agent_manager.delete(agent_id):

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

329 return {"deleted": agent_id}

330

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

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

333 agent = agent_manager.get(agent_id)

334 if not agent:

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

336

337 t0 = time.time()

338 try:

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

340 await asyncio.sleep(0.1)

341 agent.tasks_completed += 1

342 elapsed = time.time() - t0

343

344 return RunResponse(

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

346 agent_id=agent_id,

347 result=result,

348 elapsed=elapsed,

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

350 )

351 except Exception as e:

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

353

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

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

356 agent = agent_manager.get(agent_id)

357 if not agent:

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

359

360 async def event_generator():

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

362 for i, word in enumerate(words):

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

364 await asyncio.sleep(0.05)

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

366

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

368

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

370 async def run_workflow(request: WorkflowRunRequest):

371 try:

372 import yaml

373 from agentos.workflow import WorkflowParser, WorkflowEngine

374

375 wf_data = yaml.safe_load(request.workflow_yaml)

376 wf = WorkflowParser.parse_dict(wf_data)

377 wf.variables.update(request.variables)

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

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

380 except Exception as e:

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

382

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

384 async def validate_workflow(request: WorkflowRunRequest):

385 try:

386 import yaml

387 from agentos.workflow import WorkflowParser, WorkflowEngine

388 wf_data = yaml.safe_load(request.workflow_yaml)

389 wf = WorkflowParser.parse_dict(wf_data)

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

391 return result

392 except Exception as e:

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

394

395 # -----------------------------------------------------------------------

396 # WebSocket endpoint

397 # -----------------------------------------------------------------------

398

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

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

401 agent = agent_manager.get(agent_id)

402 if not agent:

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

404 return

405

406 await websocket.accept()

407 active_ws[agent_id] = websocket

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

409

410 try:

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

412

413 while True:

414 data = await websocket.receive_text()

415 msg = json.loads(data)

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

417

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

419 for i, word in enumerate(words):

420 await websocket.send_json({

421 "type": "token",

422 "data": word,

423 "seq": i,

424 })

425 await asyncio.sleep(0.03)

426

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

428 agent.tasks_completed += 1

429

430 except WebSocketDisconnect:

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

432 except Exception as e:

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

434 finally:

435 active_ws.pop(agent_id, None)

436

437 # -----------------------------------------------------------------------

438 # Marketplace endpoints

439 # -----------------------------------------------------------------------

440

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

442 async def marketplace_search(q: str = "", category: Optional[str] = None, limit: int = 20):

443 try:

444 from agentos.marketplace import MarketplaceManager, MarketSearchQuery, TemplateCategory

445 manager = MarketplaceManager()

446 cat = TemplateCategory(category) if category else None

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

448 return {

449 "results": [

450 {

451 "id": r.template.id,

452 "name": r.template.name,

453 "description": r.template.description,

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

455 "rating": r.template.rating,

456 "stars": r.template.stars,

457 "downloads": r.template.downloads,

458 "tags": r.template.tags,

459 }

460 for r in results

461 ]

462 }

463 except Exception as e:

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

465

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

467 async def marketplace_stats():

468 from agentos.marketplace import MarketplaceManager, seed_default_templates

469 manager = MarketplaceManager()

470 seed_default_templates(manager)

471 return await manager.get_stats()

472

473else:

474 app = None

475

476

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

478 """Start the API server."""

479 if not HAS_API_DEPS:

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

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

482 return

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

484

485

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

487

488

489# ── Auto-generated compat stubs ──

490

491# Auto-generated compat stubs

492class AgentAPI: pass

493class RunResponse: pass