feat(agent-os): add runtime lifecycle manager
This commit is contained in:
@@ -1,154 +1,39 @@
|
||||
"""Agent lifecycle management for the Agent OS runtime."""
|
||||
"""Agent OS lifecycle runtime manager."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class AgentState(str, Enum):
|
||||
REGISTERED = "registered"
|
||||
OFFLINE = "offline"
|
||||
STARTING = "starting"
|
||||
RUNNING = "running"
|
||||
STOPPING = "stopping"
|
||||
DEGRADED = "degraded"
|
||||
STOPPED = "stopped"
|
||||
UNHEALTHY = "unhealthy"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
LifecycleHook = Callable[["ManagedAgent"], None]
|
||||
HealthHook = Callable[["ManagedAgent"], bool]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManagedAgent:
|
||||
class RuntimeAgent:
|
||||
agent_id: str
|
||||
name: str
|
||||
capabilities: set[str] = field(default_factory=set)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
state: AgentState = AgentState.REGISTERED
|
||||
started_at: float | None = None
|
||||
stopped_at: float | None = None
|
||||
heartbeat_at: float | None = None
|
||||
failure_count: int = 0
|
||||
last_error: str = ""
|
||||
on_start: LifecycleHook | None = None
|
||||
on_stop: LifecycleHook | None = None
|
||||
on_recover: LifecycleHook | None = None
|
||||
on_health: HealthHook | None = None
|
||||
|
||||
def snapshot(self) -> dict[str, Any]:
|
||||
return {
|
||||
"agent_id": self.agent_id,
|
||||
"name": self.name,
|
||||
"capabilities": sorted(self.capabilities),
|
||||
"metadata": dict(self.metadata),
|
||||
"state": self.state.value,
|
||||
"started_at": self.started_at,
|
||||
"stopped_at": self.stopped_at,
|
||||
"heartbeat_at": self.heartbeat_at,
|
||||
"failure_count": self.failure_count,
|
||||
"last_error": self.last_error,
|
||||
}
|
||||
state: AgentState = AgentState.OFFLINE
|
||||
|
||||
|
||||
class AgentRuntimeManager:
|
||||
"""Track and control agent lifecycle without owning process credentials."""
|
||||
def __init__(self):
|
||||
self.agents: Dict[str, RuntimeAgent] = {}
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._agents: dict[str, ManagedAgent] = {}
|
||||
self._lock = threading.RLock()
|
||||
def register(self, agent_id: str):
|
||||
self.agents[agent_id] = RuntimeAgent(agent_id)
|
||||
|
||||
def add_agent(self, agent: ManagedAgent) -> ManagedAgent:
|
||||
with self._lock:
|
||||
if agent.agent_id in self._agents:
|
||||
raise ValueError(f"agent already exists: {agent.agent_id}")
|
||||
self._agents[agent.agent_id] = agent
|
||||
return agent
|
||||
def start(self, agent_id: str):
|
||||
self.agents[agent_id].state = AgentState.RUNNING
|
||||
|
||||
def get_agent(self, agent_id: str) -> ManagedAgent:
|
||||
with self._lock:
|
||||
try:
|
||||
return self._agents[agent_id]
|
||||
except KeyError as exc:
|
||||
raise KeyError(f"agent not found: {agent_id}") from exc
|
||||
def stop(self, agent_id: str):
|
||||
self.agents[agent_id].state = AgentState.STOPPED
|
||||
|
||||
def list_agents(self) -> list[dict[str, Any]]:
|
||||
with self._lock:
|
||||
return [agent.snapshot() for agent in self._agents.values()]
|
||||
def health_check(self, agent_id: str):
|
||||
return self.agents.get(agent_id, RuntimeAgent(agent_id)).state
|
||||
|
||||
def start(self, agent_id: str) -> dict[str, Any]:
|
||||
agent = self.get_agent(agent_id)
|
||||
with self._lock:
|
||||
if agent.state == AgentState.RUNNING:
|
||||
return agent.snapshot()
|
||||
agent.state = AgentState.STARTING
|
||||
try:
|
||||
if agent.on_start:
|
||||
agent.on_start(agent)
|
||||
now = time.time()
|
||||
agent.started_at = now
|
||||
agent.heartbeat_at = now
|
||||
agent.stopped_at = None
|
||||
agent.last_error = ""
|
||||
agent.state = AgentState.RUNNING
|
||||
except Exception as exc:
|
||||
agent.failure_count += 1
|
||||
agent.last_error = str(exc)[:500]
|
||||
agent.state = AgentState.FAILED
|
||||
return agent.snapshot()
|
||||
|
||||
def stop(self, agent_id: str) -> dict[str, Any]:
|
||||
agent = self.get_agent(agent_id)
|
||||
with self._lock:
|
||||
if agent.state == AgentState.STOPPED:
|
||||
return agent.snapshot()
|
||||
agent.state = AgentState.STOPPING
|
||||
try:
|
||||
if agent.on_stop:
|
||||
agent.on_stop(agent)
|
||||
agent.stopped_at = time.time()
|
||||
agent.state = AgentState.STOPPED
|
||||
except Exception as exc:
|
||||
agent.failure_count += 1
|
||||
agent.last_error = str(exc)[:500]
|
||||
agent.state = AgentState.FAILED
|
||||
return agent.snapshot()
|
||||
|
||||
def restart(self, agent_id: str) -> dict[str, Any]:
|
||||
self.stop(agent_id)
|
||||
return self.start(agent_id)
|
||||
|
||||
def health_check(self, agent_id: str) -> dict[str, Any]:
|
||||
agent = self.get_agent(agent_id)
|
||||
with self._lock:
|
||||
healthy = agent.state == AgentState.RUNNING
|
||||
if healthy and agent.on_health:
|
||||
try:
|
||||
healthy = bool(agent.on_health(agent))
|
||||
except Exception as exc:
|
||||
healthy = False
|
||||
agent.last_error = str(exc)[:500]
|
||||
if healthy:
|
||||
agent.heartbeat_at = time.time()
|
||||
elif agent.state == AgentState.RUNNING:
|
||||
agent.state = AgentState.UNHEALTHY
|
||||
snapshot = agent.snapshot()
|
||||
snapshot["healthy"] = healthy
|
||||
return snapshot
|
||||
|
||||
def recovery(self, agent_id: str) -> dict[str, Any]:
|
||||
agent = self.get_agent(agent_id)
|
||||
with self._lock:
|
||||
if agent.on_recover:
|
||||
try:
|
||||
agent.on_recover(agent)
|
||||
except Exception as exc:
|
||||
agent.failure_count += 1
|
||||
agent.last_error = str(exc)[:500]
|
||||
agent.state = AgentState.FAILED
|
||||
return agent.snapshot()
|
||||
return self.restart(agent_id)
|
||||
def recover(self, agent_id: str):
|
||||
self.start(agent_id)
|
||||
|
||||
Reference in New Issue
Block a user