From 8a5ce630695e9d912cf7b8de9a393c52e792583b Mon Sep 17 00:00:00 2001 From: AIOS GitOps Bot Date: Tue, 7 Jul 2026 10:12:27 +0800 Subject: [PATCH] feat(agent-os): implement runtime framework --- runtime/agent-os/agent_evaluator.py | 63 +++++++++ runtime/agent-os/agent_registry.py | 114 ++++++++++++++++ runtime/agent-os/agent_runtime_manager.py | 154 +++++++++++++++++++++ runtime/agent-os/api.py | 98 +++++++++++++ runtime/agent-os/consensus_engine.py | 75 ++++++++++ runtime/agent-os/evolution_engine.py | 91 +++++++++++++ runtime/agent-os/memory_service.py | 107 +++++++++++++++ runtime/agent-os/message_bus.py | 93 +++++++++++++ runtime/agent-os/skill_loader.py | 105 ++++++++++++++ runtime/agent-os/task_scheduler.py | 159 ++++++++++++++++++++++ 10 files changed, 1059 insertions(+) create mode 100644 runtime/agent-os/agent_evaluator.py create mode 100644 runtime/agent-os/agent_registry.py create mode 100644 runtime/agent-os/agent_runtime_manager.py create mode 100644 runtime/agent-os/api.py create mode 100644 runtime/agent-os/consensus_engine.py create mode 100644 runtime/agent-os/evolution_engine.py create mode 100644 runtime/agent-os/memory_service.py create mode 100644 runtime/agent-os/message_bus.py create mode 100644 runtime/agent-os/skill_loader.py create mode 100644 runtime/agent-os/task_scheduler.py diff --git a/runtime/agent-os/agent_evaluator.py b/runtime/agent-os/agent_evaluator.py new file mode 100644 index 0000000..8a1d542 --- /dev/null +++ b/runtime/agent-os/agent_evaluator.py @@ -0,0 +1,63 @@ +"""Agent quality, acceptance, and performance evaluation.""" + +from __future__ import annotations + +from dataclasses import dataclass +from statistics import mean +from typing import Any + + +@dataclass +class EvaluationInput: + correctness: float = 0.0 + completeness: float = 0.0 + safety: float = 0.0 + latency_ms: float = 0.0 + retries: int = 0 + accepted: bool = False + notes: str = "" + + +class AgentEvaluator: + def quality_score(self, item: EvaluationInput) -> float: + return self._bounded(mean([item.correctness, item.completeness, item.safety])) + + def acceptance_score(self, item: EvaluationInput) -> float: + base = 1.0 if item.accepted else 0.0 + return self._bounded((base * 0.5) + (item.correctness * 0.3) + (item.safety * 0.2)) + + def performance_feedback(self, item: EvaluationInput) -> dict[str, Any]: + latency_penalty = min(item.latency_ms / 60000.0, 1.0) + retry_penalty = min(item.retries * 0.1, 0.5) + performance = self._bounded(1.0 - latency_penalty - retry_penalty) + return { + "performance_score": performance, + "latency_ms": item.latency_ms, + "retries": item.retries, + "recommendation": self._recommendation(item, performance), + } + + def evaluate(self, item: EvaluationInput) -> dict[str, Any]: + quality = self.quality_score(item) + acceptance = self.acceptance_score(item) + feedback = self.performance_feedback(item) + overall = self._bounded((quality * 0.45) + (acceptance * 0.35) + (feedback["performance_score"] * 0.20)) + return { + "quality_score": quality, + "acceptance_score": acceptance, + "overall_score": overall, + "feedback": feedback, + "notes": item.notes, + } + + def _recommendation(self, item: EvaluationInput, performance: float) -> str: + if item.safety < 0.8: + return "route_to_review" + if not item.accepted: + return "retry_or_handoff" + if performance < 0.5: + return "optimize_or_parallelize" + return "accepted" + + def _bounded(self, value: float) -> float: + return max(0.0, min(1.0, float(value))) diff --git a/runtime/agent-os/agent_registry.py b/runtime/agent-os/agent_registry.py new file mode 100644 index 0000000..dd05cfc --- /dev/null +++ b/runtime/agent-os/agent_registry.py @@ -0,0 +1,114 @@ +"""Agent registration and discovery for Agent OS.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import time +from typing import Any + + +@dataclass +class AgentRecord: + agent_id: str + name: str + capabilities: set[str] = field(default_factory=set) + skills: set[str] = field(default_factory=set) + status: str = "registered" + metadata: dict[str, Any] = field(default_factory=dict) + registered_at: float = field(default_factory=time.time) + updated_at: float = field(default_factory=time.time) + + def snapshot(self) -> dict[str, Any]: + return { + "agent_id": self.agent_id, + "name": self.name, + "capabilities": sorted(self.capabilities), + "skills": sorted(self.skills), + "status": self.status, + "metadata": dict(self.metadata), + "registered_at": self.registered_at, + "updated_at": self.updated_at, + } + + +class AgentRegistry: + def __init__(self) -> None: + self._agents: dict[str, AgentRecord] = {} + + def register_agent( + self, + agent_id: str, + name: str, + capabilities: list[str] | set[str] | None = None, + skills: list[str] | set[str] | None = None, + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + if agent_id in self._agents: + raise ValueError(f"agent already registered: {agent_id}") + record = AgentRecord( + agent_id=agent_id, + name=name, + capabilities=set(capabilities or []), + skills=set(skills or []), + metadata=dict(metadata or {}), + ) + self._agents[agent_id] = record + return record.snapshot() + + def get_agent(self, agent_id: str) -> dict[str, Any] | None: + record = self._agents.get(agent_id) + return record.snapshot() if record else None + + def list_agents(self, status: str | None = None) -> list[dict[str, Any]]: + records = self._agents.values() + if status is not None: + records = [record for record in records if record.status == status] + return [record.snapshot() for record in records] + + def update_status(self, agent_id: str, status: str) -> dict[str, Any]: + record = self._require(agent_id) + record.status = status + record.updated_at = time.time() + return record.snapshot() + + def add_capability(self, agent_id: str, capability: str) -> dict[str, Any]: + record = self._require(agent_id) + record.capabilities.add(capability) + record.updated_at = time.time() + return record.snapshot() + + def add_skill(self, agent_id: str, skill: str) -> dict[str, Any]: + record = self._require(agent_id) + record.skills.add(skill) + record.updated_at = time.time() + return record.snapshot() + + def find_by_capability(self, capability: str, only_active: bool = True) -> list[dict[str, Any]]: + matches = [] + for record in self._agents.values(): + if capability in record.capabilities and (not only_active or record.status in {"registered", "idle", "running"}): + matches.append(record.snapshot()) + return matches + + def discover_capabilities(self) -> dict[str, list[str]]: + capabilities: dict[str, list[str]] = {} + for record in self._agents.values(): + for capability in record.capabilities: + capabilities.setdefault(capability, []).append(record.agent_id) + return {name: sorted(agent_ids) for name, agent_ids in capabilities.items()} + + def discover_skills(self) -> dict[str, list[str]]: + skills: dict[str, list[str]] = {} + for record in self._agents.values(): + for skill in record.skills: + skills.setdefault(skill, []).append(record.agent_id) + return {name: sorted(agent_ids) for name, agent_ids in skills.items()} + + def unregister_agent(self, agent_id: str) -> bool: + return self._agents.pop(agent_id, None) is not None + + def _require(self, agent_id: str) -> AgentRecord: + try: + return self._agents[agent_id] + except KeyError as exc: + raise KeyError(f"agent not registered: {agent_id}") from exc diff --git a/runtime/agent-os/agent_runtime_manager.py b/runtime/agent-os/agent_runtime_manager.py new file mode 100644 index 0000000..70d7e68 --- /dev/null +++ b/runtime/agent-os/agent_runtime_manager.py @@ -0,0 +1,154 @@ +"""Agent lifecycle management for the Agent OS runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +import threading +import time +from typing import Any, Callable + + +class AgentState(str, Enum): + REGISTERED = "registered" + STARTING = "starting" + RUNNING = "running" + STOPPING = "stopping" + STOPPED = "stopped" + UNHEALTHY = "unhealthy" + FAILED = "failed" + + +LifecycleHook = Callable[["ManagedAgent"], None] +HealthHook = Callable[["ManagedAgent"], bool] + + +@dataclass +class ManagedAgent: + 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, + } + + +class AgentRuntimeManager: + """Track and control agent lifecycle without owning process credentials.""" + + def __init__(self) -> None: + self._agents: dict[str, ManagedAgent] = {} + self._lock = threading.RLock() + + 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 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 list_agents(self) -> list[dict[str, Any]]: + with self._lock: + return [agent.snapshot() for agent in self._agents.values()] + + 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) diff --git a/runtime/agent-os/api.py b/runtime/agent-os/api.py new file mode 100644 index 0000000..784ddfc --- /dev/null +++ b/runtime/agent-os/api.py @@ -0,0 +1,98 @@ +"""Public API facade for Agent OS runtime.""" + +from __future__ import annotations + +from pathlib import Path +import sys +from typing import Any + + +RUNTIME_DIR = Path(__file__).resolve().parent +if str(RUNTIME_DIR) not in sys.path: + sys.path.insert(0, str(RUNTIME_DIR)) + +from agent_registry import AgentRegistry # noqa: E402 +from memory_service import MemoryService # noqa: E402 +from skill_loader import SkillLoader # noqa: E402 +from task_scheduler import TaskScheduler # noqa: E402 + + +class AgentOSApi: + def __init__(self) -> None: + self.registry = AgentRegistry() + self.memory = MemoryService() + self.skills = SkillLoader() + self.scheduler = TaskScheduler(registry=self.registry) + + def register_agent( + self, + agent_id: str, + name: str, + capabilities: list[str] | None = None, + skills: list[str] | None = None, + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return self.registry.register_agent(agent_id, name, capabilities, skills, metadata) + + def list_agents(self, status: str | None = None) -> list[dict[str, Any]]: + return self.registry.list_agents(status=status) + + def submit_task( + self, + capability: str, + payload: dict[str, Any] | None = None, + priority: int = 5, + timeout_seconds: float = 60.0, + max_retries: int = 1, + ) -> dict[str, Any]: + return self.scheduler.submit_task(capability, payload, priority, timeout_seconds, max_retries) + + def query_status(self, task_id: str) -> dict[str, Any]: + return self.scheduler.query_status(task_id) + + def memory_store( + self, + key: str, + value: Any, + kind: str = "working", + scope: str = "default", + tags: list[str] | None = None, + ) -> dict[str, Any]: + return self.memory.store(key=key, value=value, kind=kind, scope=scope, tags=tags) + + def skill_list(self, capability: str | None = None) -> list[dict[str, Any]]: + return self.skills.list_skills(capability=capability) + + +api = AgentOSApi() + + +def register_agent(*args: Any, **kwargs: Any) -> dict[str, Any]: + return api.register_agent(*args, **kwargs) + + +def list_agents(*args: Any, **kwargs: Any) -> list[dict[str, Any]]: + return api.list_agents(*args, **kwargs) + + +def submit_task(*args: Any, **kwargs: Any) -> dict[str, Any]: + return api.submit_task(*args, **kwargs) + + +def query_status(*args: Any, **kwargs: Any) -> dict[str, Any]: + return api.query_status(*args, **kwargs) + + +def memory_store(*args: Any, **kwargs: Any) -> dict[str, Any]: + return api.memory_store(*args, **kwargs) + + +def skill_list(*args: Any, **kwargs: Any) -> list[dict[str, Any]]: + return api.skill_list(*args, **kwargs) + + +if __name__ == "__main__": + register_agent("agent.local", "Local Agent", capabilities=["self-test"], skills=["core"]) + task = submit_task("self-test", {"message": "hello"}) + memory_store("last_task", task["task_id"], kind="working") + print({"agents": list_agents(), "task": query_status(task["task_id"]), "skills": skill_list()}) diff --git a/runtime/agent-os/consensus_engine.py b/runtime/agent-os/consensus_engine.py new file mode 100644 index 0000000..150f638 --- /dev/null +++ b/runtime/agent-os/consensus_engine.py @@ -0,0 +1,75 @@ +"""Weighted consensus engine for multi-agent decisions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +DEFAULT_WEIGHTS = { + "custom_gpt": 2, + "windows_codex": 2, + "server_agent": 1, +} + + +@dataclass +class Vote: + voter_id: str + voter_type: str + choice: str + confidence: float = 1.0 + rationale: str = "" + + def weight(self, weights: dict[str, int]) -> float: + return float(weights.get(self.voter_type, weights.get("server_agent", 1))) * max(0.0, min(1.0, self.confidence)) + + +class ConsensusEngine: + def __init__(self, weights: dict[str, int] | None = None) -> None: + self.weights = dict(DEFAULT_WEIGHTS) + if weights: + self.weights.update(weights) + + def weighted_vote(self, votes: list[Vote]) -> dict[str, Any]: + totals: dict[str, float] = {} + for vote in votes: + totals[vote.choice] = totals.get(vote.choice, 0.0) + vote.weight(self.weights) + if not totals: + return {"decision": "owner_escalation_required", "totals": {}, "tie": True} + ordered = sorted(totals.items(), key=lambda item: (-item[1], item[0])) + top_score = ordered[0][1] + winners = [choice for choice, score in ordered if score == top_score] + if len(winners) > 1: + return self.owner_escalation(votes, totals, reason="weighted_tie") + return {"decision": winners[0], "totals": totals, "tie": False} + + def conflict_detection(self, votes: list[Vote]) -> dict[str, Any]: + choices = sorted({vote.choice for vote in votes}) + return {"conflict": len(choices) > 1, "choices": choices, "vote_count": len(votes)} + + def owner_escalation(self, votes: list[Vote], totals: dict[str, float] | None = None, reason: str = "conflict") -> dict[str, Any]: + return { + "decision": "owner_escalation_required", + "reason": reason, + "tie": True, + "totals": totals or {}, + "review_required": "owner/admin", + "votes": [ + { + "voter_id": vote.voter_id, + "voter_type": vote.voter_type, + "choice": vote.choice, + "confidence": vote.confidence, + "rationale": vote.rationale, + } + for vote in votes + ], + } + + def decide(self, votes: list[Vote]) -> dict[str, Any]: + conflict = self.conflict_detection(votes) + result = self.weighted_vote(votes) + result["conflict"] = conflict["conflict"] + result["choices"] = conflict["choices"] + return result diff --git a/runtime/agent-os/evolution_engine.py b/runtime/agent-os/evolution_engine.py new file mode 100644 index 0000000..0812cd1 --- /dev/null +++ b/runtime/agent-os/evolution_engine.py @@ -0,0 +1,91 @@ +"""Closed-loop self-evolution engine.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +import time +import uuid +from typing import Any, Callable + + +class EvolutionStage(str, Enum): + DETECT = "detect" + PLAN = "plan" + SCHEDULE = "schedule" + EXECUTE = "execute" + VALIDATE = "validate" + VOTE = "vote" + PATCH = "patch" + LEARN = "learn" + + +@dataclass +class EvolutionRun: + run_id: str = field(default_factory=lambda: str(uuid.uuid4())) + stage: EvolutionStage = EvolutionStage.DETECT + context: dict[str, Any] = field(default_factory=dict) + timeline: list[dict[str, Any]] = field(default_factory=list) + status: str = "created" + + def record(self, stage: EvolutionStage, result: Any) -> None: + self.stage = stage + self.timeline.append({"stage": stage.value, "result": result, "ts": time.time()}) + + +StageHandler = Callable[[EvolutionRun], Any] + + +class EvolutionEngine: + ORDER = [ + EvolutionStage.DETECT, + EvolutionStage.PLAN, + EvolutionStage.SCHEDULE, + EvolutionStage.EXECUTE, + EvolutionStage.VALIDATE, + EvolutionStage.VOTE, + EvolutionStage.PATCH, + EvolutionStage.LEARN, + ] + + def __init__(self) -> None: + self._handlers: dict[EvolutionStage, StageHandler] = {} + + def register_handler(self, stage: EvolutionStage | str, handler: StageHandler) -> None: + self._handlers[EvolutionStage(stage)] = handler + + def detect(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.DETECT, run, {"detected": True, "signals": run.context.get("signals", [])}) + + def plan(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.PLAN, run, {"plan": run.context.get("plan", [])}) + + def schedule(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.SCHEDULE, run, {"scheduled": True}) + + def execute(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.EXECUTE, run, {"executed": True}) + + def validate(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.VALIDATE, run, {"valid": True}) + + def vote(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.VOTE, run, {"decision": "accepted"}) + + def patch(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.PATCH, run, {"patch_ready": True}) + + def learn(self, run: EvolutionRun) -> Any: + return self._handle(EvolutionStage.LEARN, run, {"learned": True}) + + def run_cycle(self, context: dict[str, Any] | None = None) -> dict[str, Any]: + run = EvolutionRun(context=dict(context or {}), status="running") + for stage in self.ORDER: + result = getattr(self, stage.value)(run) + run.record(stage, result) + run.status = "completed" + return {"run_id": run.run_id, "status": run.status, "timeline": run.timeline} + + def _handle(self, stage: EvolutionStage, run: EvolutionRun, default: Any) -> Any: + handler = self._handlers.get(stage) + return handler(run) if handler else default diff --git a/runtime/agent-os/memory_service.py b/runtime/agent-os/memory_service.py new file mode 100644 index 0000000..e5f0c7f --- /dev/null +++ b/runtime/agent-os/memory_service.py @@ -0,0 +1,107 @@ +"""Memory service with secret-safe storage checks.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +import re +import time +from typing import Any + + +SECRET_PATTERNS = [ + re.compile(pattern, re.IGNORECASE) + for pattern in ( + r"api[_-]?key", + r"password", + r"secret", + r"token", + r"bearer\s+[a-z0-9._-]+", + r"-----BEGIN [A-Z ]*PRIVATE KEY-----", + ) +] + + +@dataclass +class MemoryRecord: + key: str + value: Any + kind: str + scope: str + tags: set[str] = field(default_factory=set) + created_at: float = field(default_factory=time.time) + updated_at: float = field(default_factory=time.time) + + def snapshot(self) -> dict[str, Any]: + return { + "key": self.key, + "value": self.value, + "kind": self.kind, + "scope": self.scope, + "tags": sorted(self.tags), + "created_at": self.created_at, + "updated_at": self.updated_at, + } + + +class MemoryService: + VALID_KINDS = {"working", "long_term", "shared_knowledge", "reflection"} + + def __init__(self) -> None: + self._records: dict[str, MemoryRecord] = {} + + def store( + self, + key: str, + value: Any, + kind: str = "working", + scope: str = "default", + tags: list[str] | set[str] | None = None, + ) -> dict[str, Any]: + if kind not in self.VALID_KINDS: + raise ValueError(f"unsupported memory kind: {kind}") + self._reject_secret(key, value) + record_key = self._record_key(scope, key) + now = time.time() + existing = self._records.get(record_key) + if existing: + existing.value = value + existing.kind = kind + existing.tags = set(tags or existing.tags) + existing.updated_at = now + return existing.snapshot() + record = MemoryRecord(key=key, value=value, kind=kind, scope=scope, tags=set(tags or [])) + self._records[record_key] = record + return record.snapshot() + + def get(self, key: str, scope: str = "default") -> dict[str, Any] | None: + record = self._records.get(self._record_key(scope, key)) + return record.snapshot() if record else None + + def search(self, query: str = "", kind: str | None = None, scope: str | None = None) -> list[dict[str, Any]]: + query_lower = query.lower() + results = [] + for record in self._records.values(): + if kind and record.kind != kind: + continue + if scope and record.scope != scope: + continue + haystack = f"{record.key} {json.dumps(record.value, ensure_ascii=False, default=str)}".lower() + if query_lower in haystack: + results.append(record.snapshot()) + return results + + def delete(self, key: str, scope: str = "default") -> bool: + return self._records.pop(self._record_key(scope, key), None) is not None + + def reflect(self, key: str, value: Any, scope: str = "default") -> dict[str, Any]: + return self.store(key=key, value=value, kind="reflection", scope=scope, tags={"reflection"}) + + def _record_key(self, scope: str, key: str) -> str: + return f"{scope}:{key}" + + def _reject_secret(self, key: str, value: Any) -> None: + serialized = f"{key} {json.dumps(value, ensure_ascii=False, default=str)}" + for pattern in SECRET_PATTERNS: + if pattern.search(serialized): + raise ValueError("memory refuses to store secret-like material") diff --git a/runtime/agent-os/message_bus.py b/runtime/agent-os/message_bus.py new file mode 100644 index 0000000..4067705 --- /dev/null +++ b/runtime/agent-os/message_bus.py @@ -0,0 +1,93 @@ +"""In-process message bus for Agent OS agents and task events.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import time +import uuid +from typing import Any, Callable + + +MessageHandler = Callable[["AgentMessage"], None] + + +@dataclass +class AgentMessage: + sender: str + recipient: str + message_type: str + payload: dict[str, Any] = field(default_factory=dict) + message_id: str = field(default_factory=lambda: str(uuid.uuid4())) + created_at: float = field(default_factory=time.time) + + def snapshot(self) -> dict[str, Any]: + return { + "message_id": self.message_id, + "sender": self.sender, + "recipient": self.recipient, + "message_type": self.message_type, + "payload": dict(self.payload), + "created_at": self.created_at, + } + + +@dataclass +class TaskEvent: + task_id: str + event_type: str + source: str + detail: dict[str, Any] = field(default_factory=dict) + created_at: float = field(default_factory=time.time) + + def to_message(self) -> AgentMessage: + return AgentMessage( + sender=self.source, + recipient="task-events", + message_type=self.event_type, + payload={"task_id": self.task_id, "detail": self.detail, "created_at": self.created_at}, + ) + + +class MessageBus: + def __init__(self) -> None: + self._subscribers: dict[str, list[MessageHandler]] = {} + self._routes: dict[str, str] = {} + self._history: list[dict[str, Any]] = [] + + def subscribe(self, topic: str, handler: MessageHandler) -> None: + self._subscribers.setdefault(topic, []).append(handler) + + def unsubscribe(self, topic: str, handler: MessageHandler) -> bool: + handlers = self._subscribers.get(topic, []) + if handler not in handlers: + return False + handlers.remove(handler) + return True + + def route(self, message_type: str, topic: str) -> None: + self._routes[message_type] = topic + + def publish(self, topic: str, message: AgentMessage) -> dict[str, Any]: + self._history.append({"topic": topic, **message.snapshot()}) + delivered = 0 + for handler in list(self._subscribers.get(topic, [])): + handler(message) + delivered += 1 + return {"topic": topic, "message_id": message.message_id, "delivered": delivered} + + def send(self, message: AgentMessage) -> dict[str, Any]: + topic = self._routes.get(message.message_type, message.recipient) + return self.publish(topic, message) + + def broadcast(self, sender: str, message_type: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: + message = AgentMessage(sender=sender, recipient="broadcast", message_type=message_type, payload=dict(payload or {})) + total = 0 + for topic in list(self._subscribers): + total += self.publish(topic, message)["delivered"] + return {"message_id": message.message_id, "delivered": total} + + def emit_task_event(self, event: TaskEvent) -> dict[str, Any]: + return self.send(event.to_message()) + + def history(self, limit: int = 100) -> list[dict[str, Any]]: + return self._history[-limit:] diff --git a/runtime/agent-os/skill_loader.py b/runtime/agent-os/skill_loader.py new file mode 100644 index 0000000..8344d35 --- /dev/null +++ b/runtime/agent-os/skill_loader.py @@ -0,0 +1,105 @@ +"""Skill installation and validation registry for Agent OS.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import re +from typing import Any + + +NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") +VERSION_RE = re.compile(r"^[0-9]+(\.[0-9]+){0,2}([+-][A-Za-z0-9._-]+)?$") + + +@dataclass +class SkillManifest: + name: str + version: str + capabilities: set[str] = field(default_factory=set) + dependencies: set[str] = field(default_factory=set) + metadata: dict[str, Any] = field(default_factory=dict) + enabled: bool = True + + def snapshot(self) -> dict[str, Any]: + return { + "name": self.name, + "version": self.version, + "capabilities": sorted(self.capabilities), + "dependencies": sorted(self.dependencies), + "metadata": dict(self.metadata), + "enabled": self.enabled, + } + + +class SkillLoader: + def __init__(self) -> None: + self._skills: dict[str, SkillManifest] = {} + + def install_skill( + self, + name: str, + version: str, + capabilities: list[str] | set[str] | None = None, + dependencies: list[str] | set[str] | None = None, + metadata: dict[str, Any] | None = None, + ) -> dict[str, Any]: + manifest = SkillManifest( + name=name, + version=version, + capabilities=set(capabilities or []), + dependencies=set(dependencies or []), + metadata=dict(metadata or {}), + ) + self.validate_skill(manifest) + missing = self.dependency_check(manifest) + if missing: + raise ValueError(f"missing skill dependencies: {', '.join(missing)}") + self._skills[name] = manifest + return manifest.snapshot() + + def validate_skill(self, manifest: SkillManifest) -> bool: + if not NAME_RE.match(manifest.name): + raise ValueError("invalid skill name") + if not VERSION_RE.match(manifest.version): + raise ValueError("invalid skill version") + return True + + def dependency_check(self, manifest: SkillManifest) -> list[str]: + return sorted(dep for dep in manifest.dependencies if dep not in self._skills) + + def set_version(self, name: str, version: str) -> dict[str, Any]: + manifest = self._require(name) + candidate = SkillManifest( + name=manifest.name, + version=version, + capabilities=set(manifest.capabilities), + dependencies=set(manifest.dependencies), + metadata=dict(manifest.metadata), + enabled=manifest.enabled, + ) + self.validate_skill(candidate) + manifest.version = version + return manifest.snapshot() + + def list_skills(self, capability: str | None = None, enabled_only: bool = True) -> list[dict[str, Any]]: + skills = self._skills.values() + if capability: + skills = [skill for skill in skills if capability in skill.capabilities] + if enabled_only: + skills = [skill for skill in skills if skill.enabled] + return [skill.snapshot() for skill in skills] + + def get_skill(self, name: str) -> dict[str, Any] | None: + manifest = self._skills.get(name) + return manifest.snapshot() if manifest else None + + def enable_skill(self, name: str, enabled: bool = True) -> dict[str, Any]: + manifest = self._require(name) + manifest.enabled = enabled + return manifest.snapshot() + + def _require(self, name: str) -> SkillManifest: + try: + return self._skills[name] + except KeyError as exc: + raise KeyError(f"skill not installed: {name}") from exc diff --git a/runtime/agent-os/task_scheduler.py b/runtime/agent-os/task_scheduler.py new file mode 100644 index 0000000..f2b7eec --- /dev/null +++ b/runtime/agent-os/task_scheduler.py @@ -0,0 +1,159 @@ +"""Priority task scheduling and handoff for Agent OS.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor, TimeoutError +from dataclasses import dataclass, field +import heapq +import time +import uuid +from typing import Any, Callable + + +TaskHandler = Callable[[dict[str, Any], dict[str, Any]], Any] + + +@dataclass(order=True) +class ScheduledTask: + sort_key: tuple[int, float] = field(init=False, repr=False) + priority: int + created_at: float + capability: str = field(compare=False) + payload: dict[str, Any] = field(default_factory=dict, compare=False) + task_id: str = field(default_factory=lambda: str(uuid.uuid4()), compare=False) + timeout_seconds: float = field(default=60.0, compare=False) + max_retries: int = field(default=1, compare=False) + attempts: int = field(default=0, compare=False) + status: str = field(default="queued", compare=False) + assigned_agent_id: str = field(default="", compare=False) + result: Any = field(default=None, compare=False) + error: str = field(default="", compare=False) + + def __post_init__(self) -> None: + self.sort_key = (-self.priority, self.created_at) + + def snapshot(self) -> dict[str, Any]: + return { + "task_id": self.task_id, + "capability": self.capability, + "priority": self.priority, + "created_at": self.created_at, + "timeout_seconds": self.timeout_seconds, + "max_retries": self.max_retries, + "attempts": self.attempts, + "status": self.status, + "assigned_agent_id": self.assigned_agent_id, + "result": self.result, + "error": self.error, + } + + +class TaskScheduler: + def __init__(self, registry: Any | None = None, max_workers: int = 4) -> None: + self.registry = registry + self.executor = ThreadPoolExecutor(max_workers=max_workers) + self._queue: list[ScheduledTask] = [] + self._tasks: dict[str, ScheduledTask] = {} + self._handlers: dict[str, TaskHandler] = {} + + def register_handler(self, capability: str, handler: TaskHandler) -> None: + self._handlers[capability] = handler + + def submit_task( + self, + capability: str, + payload: dict[str, Any] | None = None, + priority: int = 5, + timeout_seconds: float = 60.0, + max_retries: int = 1, + ) -> dict[str, Any]: + task = ScheduledTask( + priority=priority, + created_at=time.time(), + capability=capability, + payload=dict(payload or {}), + timeout_seconds=timeout_seconds, + max_retries=max_retries, + ) + self._tasks[task.task_id] = task + heapq.heappush(self._queue, task) + return task.snapshot() + + def match_agent(self, capability: str) -> dict[str, Any] | None: + if self.registry is None: + return {"agent_id": "local-handler", "capabilities": [capability], "status": "running"} + matches = self.registry.find_by_capability(capability) + return matches[0] if matches else None + + def run_next(self) -> dict[str, Any] | None: + if not self._queue: + return None + task = heapq.heappop(self._queue) + return self._execute(task).snapshot() + + def run_until_idle(self, parallel: bool = True) -> list[dict[str, Any]]: + if not parallel: + results = [] + while self._queue: + item = self.run_next() + if item: + results.append(item) + return results + + tasks = [heapq.heappop(self._queue) for _ in range(len(self._queue))] + futures = [self.executor.submit(self._execute, task) for task in tasks] + return [future.result().snapshot() for future in futures] + + def handoff(self, task_id: str, target_agent_id: str) -> dict[str, Any]: + task = self._require(task_id) + task.assigned_agent_id = target_agent_id + task.status = "handoff" + return task.snapshot() + + def query_status(self, task_id: str) -> dict[str, Any]: + return self._require(task_id).snapshot() + + def _execute(self, task: ScheduledTask) -> ScheduledTask: + agent = self.match_agent(task.capability) + if not agent: + task.status = "waiting_for_handoff" + task.error = f"no agent for capability: {task.capability}" + return task + + task.assigned_agent_id = str(agent.get("agent_id") or "") + handler = self._handlers.get(task.capability) + if handler is None: + task.status = "waiting_for_handoff" + task.error = f"no handler for capability: {task.capability}" + return task + + while task.attempts <= task.max_retries: + task.attempts += 1 + task.status = "running" + attempt_executor = ThreadPoolExecutor(max_workers=1) + future = attempt_executor.submit(handler, task.payload, agent) + try: + task.result = future.result(timeout=task.timeout_seconds) + task.status = "completed" + task.error = "" + return task + except TimeoutError: + future.cancel() + task.error = "timeout" + except Exception as exc: + task.error = str(exc)[:500] + finally: + attempt_executor.shutdown(wait=False, cancel_futures=True) + + if task.attempts <= task.max_retries: + task.status = "retrying" + time.sleep(0.05) + + task.status = "failed" + return task + + def _require(self, task_id: str) -> ScheduledTask: + try: + return self._tasks[task_id] + except KeyError as exc: + raise KeyError(f"task not found: {task_id}") from exc