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 os

12import signal

13import time

14import uuid

15from collections import defaultdict

16from contextlib import asynccontextmanager

17from dataclasses import dataclass, field

18from pathlib import Path

19from typing import Any, AsyncIterator, Dict, List, Optional, Set

21logger = logging.getLogger(__name__)

23# ── Graceful shutdown state ──────────────────────────────────────────────

25_shutting_down: bool = False

26_active_connections: int = 0

27_shutdown_event = asyncio.Event()

29try:

30 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Query

31 from fastapi.middleware.cors import CORSMiddleware

32 from fastapi.responses import StreamingResponse, JSONResponse

33 from pydantic import BaseModel, Field

34 import uvicorn

35 HAS_API_DEPS = True

36except ImportError:

37 HAS_API_DEPS = False

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

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

42# Pydantic models

43# ---------------------------------------------------------------------------

45if HAS_API_DEPS:

47 class AgentConfigRequest(BaseModel):

48 name: str = "default"

49 model: str = "gpt-4o"

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

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

52 memory: bool = False

53 max_tokens: int = 4096

54 temperature: float = 0.7

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

57 class RunRequest(BaseModel):

58 agent_id: str

59 prompt: str

60 stream: bool = False

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

63 class RunResponse(BaseModel):

64 task_id: str

65 agent_id: str

66 result: str

67 elapsed: float

68 tokens_used: int = 0

70 class AgentInfo(BaseModel):

71 id: str

72 name: str

73 model: str

74 status: str

75 tasks_completed: int = 0

76 uptime: float = 0.0

78 class WorkflowRunRequest(BaseModel):

79 workflow_yaml: str

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

82 class HealthResponse(BaseModel):

83 status: str

84 version: str

85 uptime: float

86 agents_count: int

87 active_websockets: int

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

91# Agent Manager

92# ---------------------------------------------------------------------------

94@dataclass

95class ManagedAgent:

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

97 id: str

98 name: str

99 model: str

100 config: Dict[str, Any]

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

102 tasks_completed: int = 0

103

104

105class AgentManager:

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

107

108 def __init__(self):

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

110 self._start_time = time.time()

111

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

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

114 agent = ManagedAgent(

115 id=agent_id,

116 name=config.name,

117 model=config.model,

118 config=config.model_dump(),

119 )

120 self._agents[agent_id] = agent

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

122 return agent

123

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

125 return self._agents.get(agent_id)

126

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

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

129

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

131 if agent_id in self._agents:

132 del self._agents[agent_id]

133 return True

134 return False

135

136 @property

137 def count(self) -> int:

138 return len(self._agents)

139

140 @property

141 def uptime(self) -> float:

142 return time.time() - self._start_time

143

144

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

146# FastAPI Application

147# ---------------------------------------------------------------------------

148

149if HAS_API_DEPS:

150

151 agent_manager = AgentManager()

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

153

154 # ── Prometheus metrics ───────────────────────────────────────────────

155

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

157 _metrics["agentos_uptime_seconds"] = 0.0

158 _metrics["agentos_requests_total"] = 0

159 _metrics["agentos_errors_total"] = 0

160 _metrics["agentos_active_websockets"] = 0

161 _metrics["agentos_agents_created_total"] = 0

162 _metrics_start_time = time.time()

163

164 @asynccontextmanager

165 async def lifespan(app: FastAPI):

166 global _shutting_down

167

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

169

170 # Register signal handlers for graceful shutdown

171 loop = asyncio.get_running_loop()

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

173 try:

174 loop.add_signal_handler(

175 sig,

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

177 )

178 except NotImplementedError:

179 pass # Windows doesn't support add_signal_handler

180

181 yield

182

183 # Shutdown sequence

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

185 _shutting_down = True

186 _shutdown_event.set()

187

188 # Close all active WebSocket connections

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

190 try:

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

192 except Exception:

193 pass

194 active_ws.clear()

195

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

197 grace_start = time.time()

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

199 await asyncio.sleep(0.5)

200

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

202

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

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

205 _shutting_down = True

206 _shutdown_event.set()

207

208 app = FastAPI(

209 title="AgentOS API",

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

211 version="1.18.0",

212 lifespan=lifespan,

213 )

214

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

216

217 @app.middleware("http")

218 async def request_middleware(request, call_next):

219 global _active_connections

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

221 from fastapi.responses import JSONResponse

222 return JSONResponse(

223 status_code=503,

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

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

226 )

227 _active_connections += 1

228 _metrics["agentos_requests_total"] += 1

229 try:

230 response = await call_next(request)

231 if response.status_code >= 500:

232 _metrics["agentos_errors_total"] += 1

233 return response

234 except Exception:

235 _metrics["agentos_errors_total"] += 1

236 raise

237 finally:

238 _active_connections -= 1

239

240 app.add_middleware(

241 CORSMiddleware,

242 allow_origins=["*"],

243 allow_credentials=True,

244 allow_methods=["*"],

245 allow_headers=["*"],

246 )

247

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

249 # REST Endpoints

250 # -----------------------------------------------------------------------

251

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

253 async def health():

254 from agentos import __version__

255 return HealthResponse(

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

257 version=__version__,

258 uptime=agent_manager.uptime,

259 agents_count=agent_manager.count,

260 active_websockets=len(active_ws),

261 )

262

263 @app.get("/metrics")

264 async def metrics():

265 """Prometheus-compatible metrics endpoint."""

266 from fastapi.responses import PlainTextResponse

267

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

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

270 _metrics["agentos_agents_created_total"] = agent_manager.count

271

272 lines = [

273 "# HELP agentos_uptime_seconds Server uptime in seconds",

274 "# TYPE agentos_uptime_seconds gauge",

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

276 "# HELP agentos_requests_total Total HTTP requests",

277 "# TYPE agentos_requests_total counter",

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

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

280 "# TYPE agentos_errors_total counter",

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

282 "# HELP agentos_active_websockets Active WebSocket connections",

283 "# TYPE agentos_active_websockets gauge",

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

285 "# HELP agentos_agents_created_total Total agents created",

286 "# TYPE agentos_agents_created_total counter",

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

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

289 "# TYPE agentos_active_requests gauge",

290 f"agentos_active_requests {_active_connections}",

291 f"",

292 ]

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

294

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

296 async def create_agent(config: AgentConfigRequest):

297 agent = agent_manager.create(config)

298 return AgentInfo(

299 id=agent.id,

300 name=agent.name,

301 model=agent.model,

302 status="ready",

303 )

304

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

306 async def list_agents():

307 return [

308 AgentInfo(

309 id=a.id, name=a.name, model=a.model,

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

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

312 )

313 for a in agent_manager.list_all()

314 ]

315

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

317 async def get_agent(agent_id: str):

318 agent = agent_manager.get(agent_id)

319 if not agent:

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

321 return AgentInfo(

322 id=agent.id, name=agent.name, model=agent.model,

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

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

325 )

326

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

328 async def delete_agent(agent_id: str):

329 if not agent_manager.delete(agent_id):

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

331 return {"deleted": agent_id}

332

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

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

335 agent = agent_manager.get(agent_id)

336 if not agent:

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

338

339 t0 = time.time()

340 try:

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

342 await asyncio.sleep(0.1)

343 agent.tasks_completed += 1

344 elapsed = time.time() - t0

345

346 return RunResponse(

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

348 agent_id=agent_id,

349 result=result,

350 elapsed=elapsed,

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

352 )

353 except Exception as e:

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

355

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

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

358 agent = agent_manager.get(agent_id)

359 if not agent:

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

361

362 async def event_generator():

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

364 for i, word in enumerate(words):

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

366 await asyncio.sleep(0.05)

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

368

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

370

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

372 async def run_workflow(request: WorkflowRunRequest):

373 try:

374 import yaml

375 from agentos.workflow import WorkflowParser, WorkflowEngine

376

377 wf_data = yaml.safe_load(request.workflow_yaml)

378 wf = WorkflowParser.parse_dict(wf_data)

379 wf.variables.update(request.variables)

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

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

382 except Exception as e:

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

384

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

386 async def validate_workflow(request: WorkflowRunRequest):

387 try:

388 import yaml

389 from agentos.workflow import WorkflowParser, WorkflowEngine

390 wf_data = yaml.safe_load(request.workflow_yaml)

391 wf = WorkflowParser.parse_dict(wf_data)

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

393 return result

394 except Exception as e:

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

396

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

398 # WebSocket endpoint

399 # -----------------------------------------------------------------------

400

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

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

403 agent = agent_manager.get(agent_id)

404 if not agent:

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

406 return

407

408 await websocket.accept()

409 active_ws[agent_id] = websocket

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

411

412 try:

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

414

415 while True:

416 data = await websocket.receive_text()

417 msg = json.loads(data)

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

419

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

421 for i, word in enumerate(words):

422 await websocket.send_json({

423 "type": "token",

424 "data": word,

425 "seq": i,

426 })

427 await asyncio.sleep(0.03)

428

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

430 agent.tasks_completed += 1

431

432 except WebSocketDisconnect:

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

434 except Exception as e:

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

436 finally:

437 active_ws.pop(agent_id, None)

438

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

440 # Marketplace endpoints

441 # -----------------------------------------------------------------------

442

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

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

445 try:

446 from agentos.marketplace import MarketplaceManager, MarketSearchQuery, TemplateCategory

447 manager = MarketplaceManager()

448 cat = TemplateCategory(category) if category else None

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

450 return {

451 "results": [

452 {

453 "id": r.template.id,

454 "name": r.template.name,

455 "description": r.template.description,

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

457 "rating": r.template.rating,

458 "stars": r.template.stars,

459 "downloads": r.template.downloads,

460 "tags": r.template.tags,

461 }

462 for r in results

463 ]

464 }

465 except Exception as e:

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

467

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

469 async def marketplace_stats():

470 from agentos.marketplace import MarketplaceManager, seed_default_templates

471 manager = MarketplaceManager()

472 seed_default_templates(manager)

473 return await manager.get_stats()

474

475else:

476 app = None

477

478

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

480 """Start the API server."""

481 if not HAS_API_DEPS:

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

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

484 return

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

486

487

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

489

490

491# ── Auto-generated compat stubs ──

492

493# Auto-generated compat stubs

494class AgentAPI: pass

495class RunResponse: pass