@agentikos/omega-os 0.1.0 → 0.2.0

package/omega/Agentik_Engine/omega_engine/tools.py

@@ -0,0 +1,272 @@
+"""The tool registry — Agentik_Tools/registry.json + install/remove logic.
+
+Tools are third-party software that lands under `Agentik_Tools/<name>/`. Each
+tool has one home and one invoke command. Source can be the MCP catalog
+(`mcp-catalog.yaml`), a direct npm package spec, or a pre-existing tool the
+installer already shipped (e.g. pdfgen).
+
+The registry is the manifest of what is installed — it is git-tracked, the
+tool payloads are not. On a new machine the bootstrap reads the registry and
+reinstalls exactly the same set.
+"""
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import time
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+# --------------------------------------------------------------------------
+# Tool dataclass + registry
+# --------------------------------------------------------------------------
+
+
+@dataclass
+class Tool:
+    """One row in Agentik_Tools/registry.json."""
+
+    name: str
+    version: str = ""
+    path: str = ""                              # relative to OMEGA_HOME
+    invoke: str = ""                            # relative path to the binary
+    source: str = ""                            # npm package, mcp:<id>, etc.
+    providers: list[str] = field(default_factory=list)  # provider ids that should see it
+    installed_at: str = ""
+    # Optional extras kept for MCP tools (server config + secret refs).
+    secret_refs: list[str] = field(default_factory=list)
+    enabled: bool = True
+
+    def to_dict(self) -> dict[str, Any]:
+        # Match the existing registry.json field naming for back-compat
+        # (camelCase for installedAt — see pdfgen entry shipped in the repo).
+        d = asdict(self)
+        d["installedAt"] = d.pop("installed_at", "")
+        return d
+
+    @staticmethod
+    def from_dict(d: dict[str, Any]) -> "Tool":
+        return Tool(
+            name=str(d.get("name", "")),
+            version=str(d.get("version", "")),
+            path=str(d.get("path", "")),
+            invoke=str(d.get("invoke", "")),
+            source=str(d.get("source", "")),
+            providers=list(d.get("providers", []) or []),
+            installed_at=str(d.get("installedAt", d.get("installed_at", "")) or ""),
+            secret_refs=list(d.get("secret_refs", []) or []),
+            enabled=bool(d.get("enabled", True)),
+        )
+
+
+def _omega_home(explicit: str | Path | None = None) -> Path:
+    return Path(explicit or os.environ.get("OMEGA_HOME", str(Path.home() / "Omega")))
+
+
+class ToolRegistry:
+    """Load, mutate, save Agentik_Tools/registry.json."""
+
+    def __init__(self, omega_home: Path, tools: list[Tool], version: int = 1) -> None:
+        self._home = omega_home
+        self._tools: dict[str, Tool] = {t.name: t for t in tools}
+        self._version = version
+
+    # ----- file io -------------------------------------------------------
+
+    @classmethod
+    def load(cls, omega_home: str | Path | None = None) -> "ToolRegistry":
+        home = _omega_home(omega_home)
+        rfile = home / "Agentik_Tools" / "registry.json"
+        if not rfile.exists():
+            return cls(home, [])
+        data = json.loads(rfile.read_text() or "{}")
+        tools = [Tool.from_dict(t) for t in data.get("tools", [])]
+        return cls(home, tools, version=int(data.get("version", 1)))
+
+    def save(self, omega_home: str | Path | None = None) -> None:
+        home = _omega_home(omega_home) if omega_home is not None else self._home
+        rfile = home / "Agentik_Tools" / "registry.json"
+        rfile.parent.mkdir(parents=True, exist_ok=True)
+        payload = {
+            "version": self._version,
+            "tools": [t.to_dict() for t in self._tools.values()],
+        }
+        rfile.write_text(json.dumps(payload, indent=2) + "\n")
+
+    # ----- queries -------------------------------------------------------
+
+    def list(self) -> list[Tool]:
+        return list(self._tools.values())
+
+    def get(self, name: str) -> Tool | None:
+        return self._tools.get(name)
+
+    def installed_path(self, name: str) -> Path:
+        tool = self.get(name)
+        if tool is None:
+            raise KeyError(f"tool not registered: {name}")
+        return self._home / tool.path
+
+    # ----- mutations -----------------------------------------------------
+
+    def install(self, tool: Tool) -> Tool:
+        """Register a tool (does not run the installer — see install_from_catalog)."""
+        if not tool.installed_at:
+            tool.installed_at = time.strftime("%Y-%m-%dT%H:%M:%S")
+        self._tools[tool.name] = tool
+        return tool
+
+    def remove(self, name: str) -> bool:
+        return self._tools.pop(name, None) is not None
+
+
+# --------------------------------------------------------------------------
+# MCP catalog → tool install
+# --------------------------------------------------------------------------
+
+
+class CatalogError(RuntimeError):
+    """The catalog is missing or the requested id is unknown."""
+
+
+def load_catalog(omega_home: str | Path | None = None) -> dict[str, Any]:
+    """Read Agentik_SSOT/mcp/mcp-catalog.yaml. Raises CatalogError if missing."""
+    import yaml
+    home = _omega_home(omega_home)
+    path = home / "Agentik_SSOT" / "mcp" / "mcp-catalog.yaml"
+    if not path.exists():
+        raise CatalogError(f"MCP catalog missing: {path}")
+    data = yaml.safe_load(path.read_text()) or {}
+    return data
+
+
+def catalog_entry(catalog: dict[str, Any], mcp_id: str) -> dict[str, Any]:
+    for entry in catalog.get("catalog") or []:
+        if str(entry.get("id")) == mcp_id:
+            return entry
+    raise CatalogError(f"unknown MCP id: {mcp_id}")
+
+
+def install_from_catalog(
+    omega_home: str | Path,
+    mcp_id: str,
+    catalog: dict[str, Any] | None = None,
+    registry: ToolRegistry | None = None,
+) -> Tool:
+    """Install one MCP server from the catalog.
+
+    Concrete steps for `install.method == npx`:
+      1. resolve the npm package from `install.package`
+      2. `npm install -g --prefix Agentik_Tools/<id>/ <package>` — binaries land
+         under `Agentik_Tools/<id>/bin/`
+      3. register the tool in registry.json with `source=mcp:<id>`,
+         `invoke=Agentik_Tools/<id>/bin/<binname>` (first binary in bin/), and
+         `secret_refs` copied from the catalog entry.
+
+    Other install methods are recorded but logged as "no installer wired" — they
+    can be added later without changing the call sites.
+    """
+    home = Path(omega_home)
+    if catalog is None:
+        catalog = load_catalog(home)
+    entry = catalog_entry(catalog, mcp_id)
+
+    install = entry.get("install") or {}
+    method = str(install.get("method", "")).lower()
+    package = str(install.get("package", "")).strip()
+    secret_refs = list(entry.get("secrets") or [])
+
+    tool_dir = home / "Agentik_Tools" / mcp_id
+    tool_dir.mkdir(parents=True, exist_ok=True)
+    bin_dir = tool_dir / "bin"
+
+    invoke = ""
+    version = ""
+
+    if method == "npx" and package:
+        npm = shutil.which("npm")
+        if npm is None:
+            raise RuntimeError("npm not found on PATH — install Node.js to use MCP servers")
+        # npm install -g --prefix <dir> <pkg> places binaries in <dir>/bin/.
+        try:
+            subprocess.run(
+                [npm, "install", "-g", "--prefix", str(tool_dir), package],
+                check=True, capture_output=True, text=True, timeout=600,
+            )
+        except subprocess.CalledProcessError as exc:
+            raise RuntimeError(
+                f"npm install failed for {package}: "
+                f"{(exc.stderr or exc.stdout or '')[:500]}"
+            ) from exc
+        if bin_dir.is_dir():
+            bins = sorted(p for p in bin_dir.iterdir() if p.is_file())
+            if bins:
+                invoke = str(bins[0].relative_to(home))
+        # best-effort version detect: read the installed package's package.json
+        try:
+            pkg_dir = tool_dir / "lib" / "node_modules" / package
+            pkg_json = pkg_dir / "package.json"
+            if pkg_json.exists():
+                version = str(json.loads(pkg_json.read_text()).get("version", ""))
+        except (OSError, json.JSONDecodeError):
+            version = ""
+
+    tool = Tool(
+        name=mcp_id,
+        version=version,
+        path=f"Agentik_Tools/{mcp_id}/",
+        invoke=invoke,
+        source=f"mcp:{mcp_id}" if method else f"catalog:{mcp_id}",
+        secret_refs=secret_refs,
+        installed_at=time.strftime("%Y-%m-%dT%H:%M:%S"),
+    )
+
+    reg = registry if registry is not None else ToolRegistry.load(home)
+    reg.install(tool)
+    reg.save(home)
+    return tool
+
+
+# --------------------------------------------------------------------------
+# Canonical MCP config (Agentik_SSOT/mcp/mcp-config.yaml)
+# --------------------------------------------------------------------------
+
+
+def mcp_config_path(omega_home: str | Path | None = None) -> Path:
+    return _omega_home(omega_home) / "Agentik_SSOT" / "mcp" / "mcp-config.yaml"
+
+
+def load_mcp_config(omega_home: str | Path | None = None) -> dict[str, Any]:
+    import yaml
+    path = mcp_config_path(omega_home)
+    if not path.exists():
+        return {"version": 1, "servers": []}
+    data = yaml.safe_load(path.read_text()) or {}
+    data.setdefault("version", 1)
+    data.setdefault("servers", [])
+    return data
+
+
+def merge_mcp_config(
+    omega_home: str | Path | None,
+    entry: dict[str, Any],
+) -> dict[str, Any]:
+    """Append-or-replace one server entry in mcp-config.yaml. Idempotent."""
+    import yaml
+    data = load_mcp_config(omega_home)
+    servers = data.setdefault("servers", [])
+    eid = entry.get("id")
+    for i, existing in enumerate(servers):
+        if existing.get("id") == eid:
+            servers[i] = entry
+            break
+    else:
+        servers.append(entry)
+    path = mcp_config_path(omega_home)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(yaml.safe_dump(data, sort_keys=False))
+    return data

package/omega/Agentik_Engine/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "omega-engine"
-version = "0.1.0"
+version = "0.2.0"
 description = "The Omega OS orchestration engine — event-sourced, verified-completion agent graphs."
 readme = "README.md"
 requires-python = ">=3.11"

package/omega/Agentik_Engine/tests/test_account.py

@@ -0,0 +1,333 @@
+"""The account pool + vault + billing aggregator — the real things.
+
+End-to-end coverage in stdlib:
+  - load a YAML pool from a tempdir-based OMEGA_HOME
+  - round-robin and least-used selection
+  - read_token / write_token round-trip via the vault (mode 600 enforced)
+  - BillingAggregator aggregates per-account totals from seeded events
+  - the new providers (GLM, OpenAI, DeepSeek) import cleanly with no API key
+    and only raise on `.run()` — with a helpful message
+
+Standalone:  python3 tests/test_account.py
+"""
+import os
+import stat
+import sys
+import tempfile
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from omega_engine.account import (             # noqa: E402
+    AccountPool,
+    BillingAggregator,
+    ClaudeAccount,
+    read_token,
+    vault_path,
+    write_token,
+)
+from omega_engine.events import Event, EventType  # noqa: E402
+from omega_engine.provider import (             # noqa: E402
+    AgentRequest,
+    DeepSeekProvider,
+    GLMProvider,
+    OpenAIProvider,
+)
+from omega_engine.store import SQLiteStore     # noqa: E402
+
+
+# ──── helpers ────
+
+
+def _omega_home_with_pool(tmp: Path, selection: str = "least-used") -> Path:
+    """Seed OMEGA_HOME with an accounts.yaml containing two entries."""
+    cfg_dir = tmp / "Agentik_Providers" / "claude"
+    cfg_dir.mkdir(parents=True)
+    (cfg_dir / "accounts.yaml").write_text(
+        f"""version: 1
+selection: {selection}
+pool:
+  - id: max-primary
+    label: "Primary"
+    secret_ref: CLAUDE_OAUTH_max-primary
+    weight: 1
+    status: active
+  - id: max-secondary
+    label: "Secondary"
+    secret_ref: CLAUDE_OAUTH_max-secondary
+    weight: 2
+    status: active
+"""
+    )
+    return tmp
+
+
+# ──── tests ────
+
+
+def test_pool_load_yaml_two_accounts():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp))
+        pool = AccountPool.load(home)
+        assert len(pool.accounts) == 2, f"want 2, got {len(pool.accounts)}"
+        assert {a.id for a in pool.accounts} == {"max-primary", "max-secondary"}
+        assert pool.get("max-primary").label == "Primary"
+        assert pool.get("max-secondary").weight == 2
+
+
+def test_pool_falls_back_to_example_when_no_yaml():
+    # an OMEGA_HOME with only the example file should still load (the example
+    # is the template the repo ships with).
+    with tempfile.TemporaryDirectory() as tmp:
+        cfg = Path(tmp) / "Agentik_Providers" / "claude"
+        cfg.mkdir(parents=True)
+        (cfg / "accounts.example.yaml").write_text(
+            "version: 1\nselection: least-used\npool:\n  - id: max-x\n    "
+            "label: example\n    secret_ref: CLAUDE_OAUTH_max-x\n    "
+            "weight: 1\n    status: active\n"
+        )
+        pool = AccountPool.load(tmp)
+        assert len(pool.accounts) == 1
+        assert pool.accounts[0].id == "max-x"
+
+
+def test_pool_empty_when_no_config():
+    with tempfile.TemporaryDirectory() as tmp:
+        pool = AccountPool.load(tmp)
+        assert pool.accounts == []
+        # next() on an empty pool must raise — never silently return nothing
+        try:
+            pool.next()
+        except RuntimeError as exc:
+            assert "no active" in str(exc)
+        else:
+            raise AssertionError("expected RuntimeError on empty pool")
+
+
+def test_round_robin_alternates():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp), selection="round-robin")
+        pool = AccountPool.load(home)
+        picks = [pool.next().id for _ in range(4)]
+        # strict alternation regardless of usage counters
+        assert picks == [
+            "max-primary", "max-secondary", "max-primary", "max-secondary",
+        ], picks
+
+
+def test_least_used_picks_smaller_counter():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp), selection="least-used")
+        pool = AccountPool.load(home)
+        # bias primary to have more usage; secondary must win.
+        pool.usage_for("max-primary", 1000)
+        pick = pool.next()
+        assert pick.id == "max-secondary"
+        # ties: both at the same count → either is acceptable, but
+        # consistent picks one and the next call goes to the other.
+        pool.usage_for("max-secondary", 1000)
+        first = pool.next().id
+        second = pool.next().id
+        assert {first, second} <= {"max-primary", "max-secondary"}
+
+
+def test_set_status_rejects_bad_value():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp))
+        pool = AccountPool.load(home)
+        try:
+            pool.set_status("max-primary", "frozen")
+        except ValueError as exc:
+            assert "active" in str(exc)
+        else:
+            raise AssertionError("expected ValueError on bad status")
+        # but legal values work and persist through save/load
+        pool.set_status("max-primary", "resting")
+        pool.save(home)
+        reloaded = AccountPool.load(home)
+        assert reloaded.get("max-primary").status == "resting"
+
+
+def test_disabled_account_excluded_from_rotation():
+    with tempfile.TemporaryDirectory() as tmp:
+        home = _omega_home_with_pool(Path(tmp), selection="round-robin")
+        pool = AccountPool.load(home)
+        pool.set_status("max-secondary", "disabled")
+        picks = [pool.next().id for _ in range(3)]
+        assert all(p == "max-primary" for p in picks), picks
+
+
+def test_vault_round_trip_with_mode_600():
+    with tempfile.TemporaryDirectory() as tmp:
+        secret = "CLAUDE_OAUTH_max-primary"
+        path = write_token(tmp, secret, "sk-test-12345")
+        # exact file mode must be 0o600 — nothing else should be readable
+        mode = stat.S_IMODE(path.stat().st_mode)
+        assert mode == 0o600, f"want 0o600, got {oct(mode)}"
+        # round-trip
+        token = read_token(tmp, secret)
+        assert token == "sk-test-12345"
+        # vault_path resolves to the same path
+        assert vault_path(tmp, secret) == path
+
+
+def test_vault_refuses_path_traversal():
+    with tempfile.TemporaryDirectory() as tmp:
+        for bad in ("../escape", "a/b", "..", ""):
+            try:
+                vault_path(tmp, bad)
+            except ValueError:
+                continue
+            raise AssertionError(f"vault_path accepted bad ref: {bad!r}")
+
+
+def test_vault_read_missing_raises():
+    with tempfile.TemporaryDirectory() as tmp:
+        try:
+            read_token(tmp, "CLAUDE_OAUTH_nonexistent")
+        except FileNotFoundError:
+            return
+        raise AssertionError("expected FileNotFoundError on missing vault file")
+
+
+def test_billing_aggregates_per_account_from_events():
+    """Seed a real SQLite event store with task.* events that carry
+    `usage.account_id`, then verify BillingAggregator sums them correctly."""
+    with tempfile.TemporaryDirectory() as tmp:
+        db = Path(tmp) / "omega.db"
+        store = SQLiteStore(db)
+        # Two accounts billed across three task events. One event with no
+        # usage payload must be ignored (negative test in the same pass).
+        store.append(Event(
+            task_id="t-1", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-primary",
+                                "input_tokens": 100, "output_tokens": 200}},
+        ))
+        store.append(Event(
+            task_id="t-2", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-primary",
+                                "input_tokens": 50, "output_tokens": 75}},
+        ))
+        store.append(Event(
+            task_id="t-3", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-secondary",
+                                "input_tokens": 10, "output_tokens": 20}},
+        ))
+        # noise: an event with no usage payload — must NOT be counted
+        store.append(Event(
+            task_id="t-4", type=EventType.HEARTBEAT, payload={},
+        ))
+        # noise: an event with usage but zero tokens — must NOT be counted
+        store.append(Event(
+            task_id="t-5", type=EventType.COMPLETED,
+            payload={"usage": {"account_id": "max-secondary",
+                                "input_tokens": 0, "output_tokens": 0}},
+        ))
+        totals = BillingAggregator.from_event_log(store)
+        assert set(totals.keys()) == {"max-primary", "max-secondary"}, totals
+        assert totals["max-primary"] == {
+            "input_tokens": 150, "output_tokens": 275, "total_calls": 2,
+        }
+        assert totals["max-secondary"] == {
+            "input_tokens": 10, "output_tokens": 20, "total_calls": 1,
+        }
+        store.close()
+
+
+def test_billing_empty_log_returns_empty_dict():
+    with tempfile.TemporaryDirectory() as tmp:
+        db = Path(tmp) / "omega.db"
+        store = SQLiteStore(db)
+        totals = BillingAggregator.from_event_log(store)
+        assert totals == {}
+        store.close()
+
+
+def test_new_providers_import_without_api_key():
+    """The three new providers must be constructible with no env, no args.
+    Errors are deferred to `.run()` so the rest of the engine can keep
+    talking about them as objects."""
+    # cover the no-arg path explicitly
+    for cls in (GLMProvider, OpenAIProvider, DeepSeekProvider):
+        p = cls()
+        assert isinstance(p.id, str) and p.id
+        # default model is set
+        assert p._model, f"{cls.__name__} default model not set"
+        # custom args also accepted
+        p2 = cls(model="custom", api_key=None, base_url="https://example.test")
+        assert p2._model == "custom"
+        assert p2._base_url == "https://example.test"
+
+
+def _expect_runtime_error_on_run(provider, env_var: str) -> None:
+    """Calling .run() with no key in env AND no api_key arg must raise
+    a RuntimeError whose message names the missing env var."""
+    saved = os.environ.pop(env_var, None)
+    try:
+        provider.run(AgentRequest(role="worker", prompt="hi"))
+    except RuntimeError as exc:
+        msg = str(exc)
+        assert env_var in msg, f"{env_var} not in error message: {msg!r}"
+        return
+    finally:
+        if saved is not None:
+            os.environ[env_var] = saved
+    raise AssertionError(
+        f"{type(provider).__name__}.run() did not raise on missing key"
+    )
+
+
+def test_glm_provider_raises_without_key():
+    _expect_runtime_error_on_run(GLMProvider(), "GLM_API_KEY")
+
+
+def test_openai_provider_raises_without_key():
+    _expect_runtime_error_on_run(OpenAIProvider(), "OPENAI_API_KEY")
+
+
+def test_deepseek_provider_raises_without_key():
+    _expect_runtime_error_on_run(DeepSeekProvider(), "DEEPSEEK_API_KEY")
+
+
+def test_account_dataclass_round_trip_yaml():
+    """ClaudeAccount.to_yaml -> AccountPool.save -> AccountPool.load must
+    preserve all the fields we serialize."""
+    with tempfile.TemporaryDirectory() as tmp:
+        pool = AccountPool(
+            accounts=[
+                ClaudeAccount(
+                    id="max-new", label="round trip", secret_ref="CLAUDE_OAUTH_max-new",
+                    weight=3, status="resting", tokens_used=1234, last_used_at=1700.0,
+                ),
+            ],
+            selection="round-robin",
+        )
+        cfg = pool.save(tmp)
+        assert cfg.exists()
+        reloaded = AccountPool.load(tmp)
+        assert reloaded.selection == "round-robin"
+        a = reloaded.get("max-new")
+        assert a is not None
+        assert a.weight == 3
+        assert a.status == "resting"
+        assert a.tokens_used == 1234
+        assert a.last_used_at == 1700.0
+
+
+# ──── runner ────
+
+
+def _run_all() -> bool:
+    tests = [v for k, v in sorted(globals().items())
+             if k.startswith("test_") and callable(v)]
+    passed = 0
+    for t in tests:
+        t()
+        print(f"  PASS  {t.__name__}")
+        passed += 1
+    print(f"\n{passed}/{len(tests)} account tests passed")
+    return passed == len(tests)
+
+
+if __name__ == "__main__":
+    sys.exit(0 if _run_all() else 1)