agento-core 0.1.4__py3-none-any.whl

agento/framework/agent_config_writer.py

@@ -0,0 +1,162 @@
+"""Generate agent CLI config files from resolved scoped config.
+
+Before each worker run, generates native config files that agent CLIs expect:
+  - .claude.json (Claude Code project config)
+  - .claude/settings.json (Claude Code user settings)
+  - .mcp.json (MCP server configuration)
+  - .codex/config.toml (Codex CLI config)
+
+Config field paths follow the convention:
+  agent/claude/model          -> model for Claude CLI
+  agent/claude/personality    -> system prompt / personality
+  agent/mcp/servers           -> MCP server definitions (JSON)
+  agent/codex/model           -> model for Codex CLI
+"""
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Config path prefix for agent CLI settings
+AGENT_CONFIG_PREFIX = "agent/"
+
+
+def _get_agent_config(resolved_config: dict[str, tuple[str, bool]]) -> dict[str, str]:
+    """Extract agent/* paths from resolved DB overrides into a flat dict.
+
+    Returns {relative_path: value}, e.g. {"claude/model": "opus-4"}.
+    """
+    result = {}
+    for path, (value, _encrypted) in resolved_config.items():
+        if path.startswith(AGENT_CONFIG_PREFIX) and value is not None:
+            relative = path[len(AGENT_CONFIG_PREFIX):]
+            result[relative] = value
+    return result
+
+
+def generate_claude_config(working_dir: Path, agent_config: dict[str, str]) -> None:
+    """Generate .claude.json and .claude/settings.json in the working directory."""
+    # .claude.json — project-level config
+    claude_json: dict[str, Any] = {}
+
+    model = agent_config.get("claude/model")
+    if model:
+        claude_json["model"] = model
+
+    personality = agent_config.get("claude/personality")
+    if personality:
+        claude_json["systemPrompt"] = personality
+
+    permissions = agent_config.get("claude/permissions")
+    if permissions:
+        try:
+            claude_json["permissions"] = json.loads(permissions)
+        except (json.JSONDecodeError, TypeError):
+            logger.warning("Invalid JSON in agent/claude/permissions, skipping")
+
+    if claude_json:
+        config_path = working_dir / ".claude.json"
+        config_path.write_text(json.dumps(claude_json, indent=2) + "\n")
+        logger.debug("Generated %s", config_path)
+
+    # .claude/settings.json — user-level settings
+    settings: dict[str, Any] = {}
+
+    trust_level = agent_config.get("claude/trust_level")
+    if trust_level:
+        settings["permissions"] = {"dangerouslySkipPermissions": trust_level == "full"}
+
+    if settings:
+        settings_dir = working_dir / ".claude"
+        settings_dir.mkdir(parents=True, exist_ok=True)
+        settings_path = settings_dir / "settings.json"
+        settings_path.write_text(json.dumps(settings, indent=2) + "\n")
+        logger.debug("Generated %s", settings_path)
+
+
+def _inject_agent_view_id(servers: dict, agent_view_id: int) -> dict:
+    """Append ?agent_view_id=N to toolbox SSE/MCP URLs in mcpServers config."""
+    for server_cfg in servers.values():
+        url = server_cfg.get("url", "")
+        if "/sse" in url or "/mcp" in url:
+            sep = "&" if "?" in url else "?"
+            server_cfg["url"] = f"{url}{sep}agent_view_id={agent_view_id}"
+    return servers
+
+
+def generate_mcp_config(
+    working_dir: Path,
+    agent_config: dict[str, str],
+    *,
+    agent_view_id: int | None = None,
+) -> None:
+    """Generate .mcp.json in the working directory."""
+    servers_raw = agent_config.get("mcp/servers")
+    if not servers_raw:
+        return
+
+    try:
+        servers = json.loads(servers_raw)
+    except (json.JSONDecodeError, TypeError):
+        logger.warning("Invalid JSON in agent/mcp/servers, skipping .mcp.json generation")
+        return
+
+    if agent_view_id is not None:
+        servers = _inject_agent_view_id(servers, agent_view_id)
+
+    mcp_config = {"mcpServers": servers}
+    config_path = working_dir / ".mcp.json"
+    config_path.write_text(json.dumps(mcp_config, indent=2) + "\n")
+    logger.debug("Generated %s", config_path)
+
+
+def generate_codex_config(working_dir: Path, agent_config: dict[str, str]) -> None:
+    """Generate .codex/config.toml in the working directory."""
+    lines: list[str] = []
+
+    model = agent_config.get("codex/model")
+    if model:
+        lines.append(f'model = "{model}"')
+
+    approval_mode = agent_config.get("codex/approval_mode")
+    if approval_mode:
+        lines.append(f'approval_mode = "{approval_mode}"')
+
+    if not lines:
+        return
+
+    codex_dir = working_dir / ".codex"
+    codex_dir.mkdir(parents=True, exist_ok=True)
+    config_path = codex_dir / "config.toml"
+    config_path.write_text("\n".join(lines) + "\n")
+    logger.debug("Generated %s", config_path)
+
+
+def populate_agent_configs(
+    working_dir: str | Path,
+    scoped_overrides: dict[str, tuple[str, bool]],
+    *,
+    agent_view_id: int | None = None,
+) -> None:
+    """Generate all agent CLI config files from scoped DB overrides.
+
+    Called before each worker run with the merged (agent_view -> workspace -> global) overrides.
+    """
+    wd = Path(working_dir)
+    wd.mkdir(parents=True, exist_ok=True)
+
+    agent_config = _get_agent_config(scoped_overrides)
+
+    if not agent_config:
+        logger.debug("No agent/* config paths found, skipping config file generation")
+        return
+
+    generate_claude_config(wd, agent_config)
+    generate_mcp_config(wd, agent_config, agent_view_id=agent_view_id)
+    generate_codex_config(wd, agent_config)
+
+    logger.info("Populated agent config files in %s", wd)

agento/framework/agent_manager/__init__.py

@@ -0,0 +1,36 @@
+"""Agent Manager — multi-token orchestration for LLM agent providers."""
+
+from .active import read_credentials, resolve_active_token, update_active_token
+from .auth import AuthenticationError, AuthResult, authenticate_interactive, save_credentials
+from .config import AgentManagerConfig
+from .models import AgentProvider, RotationResult, Token, UsageSummary
+from .rotator import rotate_all, rotate_tokens, select_best_token
+from .token_store import deregister_token, get_token, get_token_by_path, list_tokens, register_token, set_primary_token
+from .usage_store import get_usage_summaries, get_usage_summary, record_usage
+
+__all__ = [
+    "AgentManagerConfig",
+    "AgentProvider",
+    "AuthResult",
+    "AuthenticationError",
+    "RotationResult",
+    "Token",
+    "UsageSummary",
+    "authenticate_interactive",
+    "deregister_token",
+    "get_token",
+    "get_token_by_path",
+    "get_usage_summaries",
+    "get_usage_summary",
+    "list_tokens",
+    "read_credentials",
+    "record_usage",
+    "register_token",
+    "resolve_active_token",
+    "rotate_all",
+    "rotate_tokens",
+    "save_credentials",
+    "select_best_token",
+    "set_primary_token",
+    "update_active_token",
+]

agento/framework/agent_manager/active.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import tempfile
+from pathlib import Path
+
+from .config import AgentManagerConfig
+from .models import AgentProvider, Token
+
+
+def resolve_active_token(
+    config: AgentManagerConfig,
+    agent_type: AgentProvider,
+) -> str | None:
+    """Read the active symlink for an agent type. Returns the target path, or None."""
+    link = Path(config.active_dir) / agent_type.value
+    if not link.is_symlink():
+        return None
+    target = link.resolve()
+    if not target.is_file():
+        return None
+    return str(target)
+
+
+def update_active_token(
+    config: AgentManagerConfig,
+    agent_type: AgentProvider,
+    token: Token,
+    logger: logging.Logger | None = None,
+) -> None:
+    """Atomically update the active symlink for an agent type.
+
+    Strategy: create temp symlink in active_dir, then os.rename() over the
+    real one. os.rename() is atomic on POSIX when src and dst are on the
+    same filesystem.
+    """
+    active_dir = Path(config.active_dir)
+    active_dir.mkdir(parents=True, exist_ok=True)
+
+    link_path = active_dir / agent_type.value
+    target = token.credentials_path
+
+    # mkstemp creates a file; remove it so we can create a symlink at that path.
+    fd, tmp_path = tempfile.mkstemp(dir=str(active_dir), prefix=f".{agent_type.value}.")
+    os.close(fd)
+    os.unlink(tmp_path)
+    os.symlink(target, tmp_path)
+    os.rename(tmp_path, str(link_path))
+
+    if logger:
+        logger.info(
+            f"Active token updated: agent_type={agent_type.value} "
+            f"label={token.label} target={target}"
+        )
+
+
+def read_credentials(credentials_path: str) -> dict:
+    """Read opaque credentials JSON from a token file."""
+    with open(credentials_path) as f:
+        return json.load(f)

agento/framework/agent_manager/auth.py

@@ -0,0 +1,131 @@
+"""Interactive OAuth authentication for agent CLI tools.
+
+Launches ``claude`` or ``codex`` CLI in an isolated temporary HOME directory
+to perform OAuth, then extracts and normalises credentials to the internal
+JSON format used by token runners.
+
+The isolation prevents the auth flow from overwriting the main active
+credentials at ``/workspace/.claude`` or ``/workspace/.codex``.
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+import shutil
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol, runtime_checkable
+
+from .models import AgentProvider
+
+
+class AuthenticationError(RuntimeError):
+    """Raised when the interactive auth flow fails or is cancelled."""
+
+
+@dataclass
+class AuthResult:
+    """Normalised credentials extracted after successful OAuth."""
+
+    subscription_key: str
+    refresh_token: str | None
+    expires_at: int | None
+    subscription_type: str | None
+    id_token: str | None = None
+    raw_auth: dict | None = None  # Native auth.json for agents that need it (Codex)
+
+
+@runtime_checkable
+class AuthStrategy(Protocol):
+    """Protocol for provider-specific authentication flows."""
+
+    def authenticate(self, tmp_home: str, logger: logging.Logger) -> AuthResult: ...
+
+
+# ---------------------------------------------------------------------------
+# Auth strategy registry
+# ---------------------------------------------------------------------------
+
+_STRATEGIES: dict[AgentProvider, AuthStrategy] = {}
+
+
+def register_auth_strategy(provider: AgentProvider, strategy: AuthStrategy) -> None:
+    _STRATEGIES[provider] = strategy
+
+
+def get_auth_strategy(provider: AgentProvider) -> AuthStrategy | None:
+    return _STRATEGIES.get(provider)
+
+
+def clear_auth_strategies() -> None:
+    _STRATEGIES.clear()
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def authenticate_interactive(
+    agent_type: AgentProvider,
+    logger: logging.Logger | None = None,
+) -> AuthResult:
+    """Run interactive OAuth for the given agent type.
+
+    Creates an isolated temp HOME directory so the auth flow does NOT
+    touch ``~/.claude`` or ``~/.codex`` (symlinked to ``/workspace/``).
+
+    Raises :class:`AuthenticationError` on failure or user cancellation.
+    """
+    _log = logger or logging.getLogger(__name__)
+
+    strategy = _STRATEGIES.get(agent_type)
+    if strategy is None:
+        raise ValueError(f"No auth strategy registered for: {agent_type.value}")
+
+    tmp_home = tempfile.mkdtemp(prefix=f"auth_{agent_type.value}_")
+    _log.info(f"Using isolated HOME: {tmp_home}")
+
+    try:
+        return strategy.authenticate(tmp_home, _log)
+    finally:
+        shutil.rmtree(tmp_home, ignore_errors=True)
+        _log.debug(f"Cleaned up temp HOME: {tmp_home}")
+
+
+def save_credentials(auth_result: AuthResult, output_path: str) -> None:
+    """Save normalised credentials to a JSON file.
+
+    Creates parent directories if needed.
+    """
+    data = {
+        "subscription_key": auth_result.subscription_key,
+        "refresh_token": auth_result.refresh_token,
+        "expires_at": auth_result.expires_at,
+        "subscription_type": auth_result.subscription_type,
+        "id_token": auth_result.id_token,
+        "raw_auth": auth_result.raw_auth,
+    }
+    path = Path(output_path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(data, indent=2))
+
+
+# ---------------------------------------------------------------------------
+# Shared CLI helper (used by strategy implementations in modules)
+# ---------------------------------------------------------------------------
+
+def _run_cli(cmd: list[str], tmp_home: str, name: str) -> None:
+    """Run a CLI command with isolated HOME. Raises on failure."""
+    env = {**os.environ, "HOME": tmp_home}
+    try:
+        proc = subprocess.run(cmd, env=env)
+    except FileNotFoundError as exc:
+        raise AuthenticationError(f"{name} CLI not found. Is it installed?") from exc
+
+    if proc.returncode != 0:
+        raise AuthenticationError(
+            f"{name} login failed with exit code {proc.returncode}"
+        )

agento/framework/agent_manager/config.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class AgentManagerConfig:
+    tokens_dir: str = "/etc/tokens"
+    active_dir: str = "/etc/tokens/active"
+    usage_window_hours: int = 24
+    rotation_interval_hours: int = 1
+
+    @classmethod
+    def from_env(cls) -> AgentManagerConfig:
+        """Build from env vars only."""
+        return cls(
+            tokens_dir=os.environ.get("AGENT_TOKENS_DIR", "/etc/tokens"),
+            active_dir=os.environ.get("AGENT_ACTIVE_DIR", "/etc/tokens/active"),
+            usage_window_hours=int(os.environ.get("AGENT_USAGE_WINDOW_HOURS", "24")),
+            rotation_interval_hours=int(os.environ.get("AGENT_ROTATION_INTERVAL_HOURS", "1")),
+        )

agento/framework/agent_manager/models.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+
+
+class AgentProvider(Enum):
+    CLAUDE = "claude"
+    CODEX = "codex"
+
+
+@dataclass
+class Token:
+    id: int
+    agent_type: AgentProvider
+    label: str
+    credentials_path: str
+    model: str | None
+    is_primary: bool
+    token_limit: int
+    enabled: bool
+    created_at: datetime
+    updated_at: datetime
+
+    @classmethod
+    def from_row(cls, row: dict) -> Token:
+        return cls(
+            id=row["id"],
+            agent_type=AgentProvider(row["agent_type"]),
+            label=row["label"],
+            credentials_path=row["credentials_path"],
+            model=row.get("model"),
+            is_primary=bool(row.get("is_primary", False)),
+            token_limit=row["token_limit"],
+            enabled=bool(row["enabled"]),
+            created_at=row["created_at"],
+            updated_at=row["updated_at"],
+        )
+
+
+@dataclass
+class UsageSummary:
+    token_id: int
+    total_tokens: int
+    call_count: int
+
+
+@dataclass
+class RotationResult:
+    agent_type: AgentProvider
+    previous_token_id: int | None
+    new_token_id: int
+    reason: str
+    timestamp: datetime

agento/framework/agent_manager/rotator.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+from datetime import UTC, datetime
+
+import pymysql
+
+from .active import resolve_active_token, update_active_token
+from .config import AgentManagerConfig
+from .models import AgentProvider, RotationResult, Token, UsageSummary
+from .token_store import get_token_by_path, list_tokens
+from .usage_store import get_usage_summaries
+
+
+def select_best_token(
+    tokens: list[Token],
+    usage_map: dict[int, UsageSummary],
+) -> Token | None:
+    """Pick the token with the most remaining capacity.
+
+    Algorithm:
+    0. If any token has is_primary=True, return it immediately (sticky selection).
+    1. For each token, compute remaining = token_limit - total_tokens_used.
+       If token_limit == 0 (unlimited), remaining is infinity.
+    2. Return the token with the highest remaining capacity.
+    3. Tie-break: prefer the token with fewer total calls.
+    4. Returns None if no tokens available.
+    """
+    if not tokens:
+        return None
+
+    # Respect manually-set primary token
+    for t in tokens:
+        if t.is_primary:
+            return t
+
+    def _sort_key(t: Token) -> tuple[float, int]:
+        summary = usage_map.get(t.id)
+        used = summary.total_tokens if summary else 0
+        calls = summary.call_count if summary else 0
+        remaining = float("inf") if t.token_limit == 0 else (t.token_limit - used)
+        # Negate remaining so highest is first; use calls as tie-break (fewer = better)
+        return (-remaining, calls)
+
+    return min(tokens, key=_sort_key)
+
+
+def rotate_tokens(
+    conn: pymysql.Connection,
+    config: AgentManagerConfig,
+    agent_type: AgentProvider,
+    logger: logging.Logger | None = None,
+) -> RotationResult | None:
+    """Perform rotation for a single agent type."""
+    _log = logger or logging.getLogger(__name__)
+
+    tokens = list_tokens(conn, agent_type=agent_type, enabled_only=True)
+    if not tokens:
+        _log.warning(f"No enabled tokens for agent_type={agent_type.value}")
+        return None
+
+    # Current active
+    active_path = resolve_active_token(config, agent_type)
+    previous_token = get_token_by_path(conn, active_path) if active_path else None
+
+    # Usage summaries
+    summaries = get_usage_summaries(conn, agent_type.value, config.usage_window_hours)
+    usage_map = {s.token_id: s for s in summaries}
+
+    best = select_best_token(tokens, usage_map)
+    if best is None:
+        return None
+
+    # Update symlink if changed (or if no active token yet)
+    if previous_token is None or previous_token.id != best.id:
+        update_active_token(config, agent_type, best, logger)
+        reason = "initial" if previous_token is None else "rotation"
+    else:
+        reason = "unchanged"
+
+    return RotationResult(
+        agent_type=agent_type,
+        previous_token_id=previous_token.id if previous_token else None,
+        new_token_id=best.id,
+        reason=reason,
+        timestamp=datetime.now(UTC),
+    )
+
+
+def rotate_all(
+    conn: pymysql.Connection,
+    config: AgentManagerConfig,
+    logger: logging.Logger | None = None,
+) -> list[RotationResult]:
+    """Rotate tokens for all known agent types."""
+    results = []
+    for agent_type in AgentProvider:
+        result = rotate_tokens(conn, config, agent_type, logger)
+        if result:
+            results.append(result)
+    return results