agent-loop-guard-runtime 0.6.0a2__py3-none-any.whl

app/mcp/stdio.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import json
+import subprocess
+import sys
+import threading
+from pathlib import Path
+from typing import TextIO
+
+from app.core.config import AppConfig
+from app.db.repository import Repository
+from app.db.session import build_engine, build_session_factory, init_db
+from app.mcp.gateway import MCPGateway, jsonrpc_error, mock_response
+from app.mcp.policy import MCPPolicyEngine
+
+
+def _write_message(stream: TextIO, message: dict, lock: threading.Lock | None = None) -> None:
+    payload = json.dumps(message, ensure_ascii=False, separators=(",", ":")) + "\n"
+    if lock:
+        with lock:
+            stream.write(payload)
+            stream.flush()
+    else:
+        stream.write(payload)
+        stream.flush()
+
+
+def run_stdio_proxy(
+    command: list[str],
+    *,
+    policy_path: str | None = None,
+    server_id: str = "stdio",
+    config: AppConfig | None = None,
+) -> int:
+    if not command:
+        raise ValueError("An upstream command is required after --.")
+    config = config or AppConfig.from_env()
+    engine = build_engine(config)
+    init_db(engine, config)
+    session_factory = build_session_factory(engine)
+    process = subprocess.Popen(
+        command,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        encoding="utf-8",
+        bufsize=1,
+    )
+    assert process.stdin and process.stdout and process.stderr
+    output_lock = threading.Lock()
+    gateway_lock = threading.Lock()
+    pending: dict[str, str] = {}
+
+    db = session_factory()
+    repo = Repository(db)
+    repo.ensure_mcp_server(server_id, server_id, "stdio", " ".join(command))
+    policy = MCPPolicyEngine(policy_path or config.mcp_policy_path)
+    session = repo.start_mcp_session(server_id, client_name="stdio-client", mode=str(policy.data.get("mode", "enforce")))
+    gateway = MCPGateway(
+        repo,
+        policy,
+        server_id,
+        session_id=session.id,
+        approval_timeout_seconds=config.mcp_approval_timeout_seconds,
+    )
+
+    def upstream_reader() -> None:
+        for raw in process.stdout:
+            try:
+                response = json.loads(raw)
+            except json.JSONDecodeError:
+                print("[mcp upstream] ignored non-JSON stdout", file=sys.stderr)
+                continue
+            request_id = str(response.get("id")) if response.get("id") is not None else None
+            with gateway_lock:
+                if request_id and pending.pop(request_id, None) == "tools/list":
+                    response = gateway.filter_tools(response)
+            _write_message(sys.stdout, response, output_lock)
+
+    def stderr_reader() -> None:
+        for raw in process.stderr:
+            sys.stderr.write(raw)
+            sys.stderr.flush()
+
+    threads = [
+        threading.Thread(target=upstream_reader, daemon=True),
+        threading.Thread(target=stderr_reader, daemon=True),
+    ]
+    for thread in threads:
+        thread.start()
+
+    try:
+        for raw in sys.stdin:
+            try:
+                message = json.loads(raw)
+            except json.JSONDecodeError:
+                _write_message(sys.stdout, jsonrpc_error(None, -32700, "Parse error"), output_lock)
+                continue
+            with gateway_lock:
+                interception = gateway.intercept(message)
+            if not interception.forward:
+                if interception.response:
+                    _write_message(sys.stdout, interception.response, output_lock)
+                continue
+            request_id = interception.message.get("id")
+            if request_id is not None:
+                pending[str(request_id)] = str(interception.message.get("method") or "")
+            _write_message(process.stdin, interception.message)
+    finally:
+        if not process.stdin.closed:
+            process.stdin.close()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            process.terminate()
+            process.wait(timeout=5)
+        for thread in threads:
+            thread.join(timeout=2)
+        repo.end_mcp_session(session.id)
+        db.close()
+    return int(process.returncode or 0)
+
+
+def run_mock_stdio_server() -> int:
+    for raw in sys.stdin:
+        try:
+            message = json.loads(raw)
+        except json.JSONDecodeError:
+            _write_message(sys.stdout, jsonrpc_error(None, -32700, "Parse error"))
+            continue
+        response = mock_response(message)
+        if response is not None:
+            _write_message(sys.stdout, response)
+    return 0
+
+
+def validate_policy_file(path: str | Path) -> list[str]:
+    from app.mcp.policy import load_policy, validate_policy
+
+    data, _ = load_policy(path)
+    return validate_policy(data)

app/platform/__init__.py

@@ -0,0 +1 @@
+"""Shared platform services used by every Agent Loop Guard module."""

app/platform/alembic/__init__.py

@@ -0,0 +1,2 @@
+"""Packaged Alembic migration environment."""
+

app/platform/alembic/env.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+from app.db.models import Base
+
+config = context.config
+target_metadata = Base.metadata
+
+
+def run_migrations_offline() -> None:
+    context.configure(
+        url=config.get_main_option("sqlalchemy.url"),
+        target_metadata=target_metadata,
+        literal_binds=True,
+        compare_type=True,
+    )
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section) or {},
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+    with connectable.connect() as connection:
+        context.configure(connection=connection, target_metadata=target_metadata, compare_type=True)
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

app/platform/alembic/script.py.mako

@@ -0,0 +1,24 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}

app/platform/alembic/versions/0001_initial_schema.py

@@ -0,0 +1,18 @@
+"""Baseline the local Agent Loop Guard schema."""
+
+from alembic import op
+
+from app.db.models import Base
+
+revision = "0001_initial"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    Base.metadata.create_all(bind=op.get_bind())
+
+
+def downgrade() -> None:
+    Base.metadata.drop_all(bind=op.get_bind())

app/platform/alembic/versions/__init__.py

@@ -0,0 +1,2 @@
+"""Database schema revisions."""
+

app/platform/events.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import time
+import uuid
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from app.core.redaction import redact_value
+
+
+@dataclass(slots=True)
+class EventEnvelope:
+    """Stable, redacted event contract shared by all local modules."""
+
+    source: str
+    type: str
+    project_id: str = "default"
+    trace_id: str | None = None
+    span_id: str | None = None
+    severity: str = "info"
+    timestamp_ns: int = field(default_factory=time.time_ns)
+    attributes: dict[str, Any] = field(default_factory=dict)
+    event_id: str = field(default_factory=lambda: f"evt_{uuid.uuid4().hex[:24]}")
+    schema_version: str = "event.v1"
+
+    def to_dict(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload["attributes"] = redact_value(self.attributes)
+        return payload
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> EventEnvelope:
+        return cls(
+            event_id=str(payload.get("event_id") or f"evt_{uuid.uuid4().hex[:24]}"),
+            source=str(payload.get("source") or "unknown"),
+            type=str(payload.get("type") or "event"),
+            project_id=str(payload.get("project_id") or "default"),
+            trace_id=payload.get("trace_id"),
+            span_id=payload.get("span_id"),
+            severity=str(payload.get("severity") or "info"),
+            timestamp_ns=int(payload.get("timestamp_ns") or time.time_ns()),
+            attributes=dict(payload.get("attributes") or {}),
+            schema_version=str(payload.get("schema_version") or "event.v1"),
+        )

app/platform/maintenance.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import json
+import shutil
+import socket
+import sys
+import zipfile
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import httpx
+from sqlalchemy import delete, select, update
+from sqlalchemy.orm import Session
+
+from app.core.config import AppConfig
+from app.db.models import (
+    Event,
+    GuardSession,
+    RequestRecord,
+    ToolCall,
+    TraceArtifact,
+    TraceEvent,
+    TraceRun,
+    TraceSpan,
+)
+
+
+def sqlite_path(config: AppConfig) -> Path | None:
+    if not config.storage_url.startswith("sqlite:///"):
+        return None
+    raw = config.storage_url.removeprefix("sqlite:///")
+    if raw == ":memory:":
+        return None
+    path = Path(raw)
+    return path if path.is_absolute() else Path.cwd() / path
+
+
+def build_doctor_report(config: AppConfig) -> dict[str, Any]:
+    checks: list[dict[str, Any]] = []
+
+    def add(name: str, ok: bool, detail: str) -> None:
+        checks.append({"name": name, "ok": ok, "detail": detail})
+
+    add("python", sys.version_info >= (3, 11), sys.version.split()[0])
+    storage = sqlite_path(config)
+    if storage is None:
+        add("storage", True, config.storage_url)
+    else:
+        try:
+            storage.parent.mkdir(parents=True, exist_ok=True)
+            add("storage", True, str(storage))
+        except OSError as exc:
+            add("storage", False, str(exc))
+
+    try:
+        with socket.socket() as probe:
+            probe.settimeout(0.3)
+            available = probe.connect_ex((config.host, config.port)) != 0
+        if available:
+            add("port", True, f"{config.host}:{config.port} available")
+        else:
+            status = fetch_status(f"http://{config.host}:{config.port}")
+            add(
+                "port",
+                status["running"],
+                f"{config.host}:{config.port} "
+                + ("Agent Loop Guard is running" if status["running"] else "used by another process"),
+            )
+    except OSError as exc:
+        add("port", False, str(exc))
+
+    add("docker", shutil.which("docker") is not None, shutil.which("docker") or "not installed")
+    add("wsl", shutil.which("wsl") is not None, shutil.which("wsl") or "not installed")
+    return {"ok": all(item["ok"] for item in checks if item["name"] != "docker"), "checks": checks}
+
+
+def fetch_status(base_url: str) -> dict[str, Any]:
+    try:
+        response = httpx.get(base_url.rstrip("/") + "/api/health", timeout=2)
+        response.raise_for_status()
+        return {"running": True, "url": base_url, **response.json()}
+    except (httpx.HTTPError, ValueError) as exc:
+        return {"running": False, "url": base_url, "error": str(exc)}
+
+
+def create_backup(config: AppConfig, destination: Path, config_path: Path | None = None) -> Path:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    manifest = {
+        "schema_version": "backup.v1",
+        "created_at": datetime.now(UTC).isoformat(),
+        "storage_url": config.storage_url,
+    }
+    with zipfile.ZipFile(destination, "w", compression=zipfile.ZIP_DEFLATED) as archive:
+        archive.writestr("manifest.json", json.dumps(manifest, indent=2))
+        storage = sqlite_path(config)
+        if storage and storage.exists():
+            archive.write(storage, "data/agent_loop_guard.db")
+        if config_path and config_path.exists():
+            archive.write(config_path, "agent-loop-guard.yml")
+    return destination
+
+
+def restore_backup(config: AppConfig, source: Path, *, force: bool = False) -> Path:
+    storage = sqlite_path(config)
+    if storage is None:
+        raise ValueError("Restore currently supports file-backed SQLite storage only.")
+    if storage.exists() and not force:
+        raise FileExistsError(f"{storage} exists; pass --force to replace it.")
+    with zipfile.ZipFile(source) as archive:
+        if "data/agent_loop_guard.db" not in archive.namelist():
+            raise ValueError("Backup does not contain a SQLite database.")
+        storage.parent.mkdir(parents=True, exist_ok=True)
+        with archive.open("data/agent_loop_guard.db") as src, storage.open("wb") as dst:
+            shutil.copyfileobj(src, dst)
+    return storage
+
+
+def cleanup_old_data(db: Session, retention_days: int) -> dict[str, int]:
+    cutoff = datetime.now(UTC) - timedelta(days=max(1, retention_days))
+    old_trace_ids = select(TraceRun.id).where(TraceRun.updated_at < cutoff)
+    db.execute(delete(TraceArtifact).where(TraceArtifact.trace_id.in_(old_trace_ids)))
+    db.execute(delete(TraceEvent).where(TraceEvent.trace_id.in_(old_trace_ids)))
+    db.execute(delete(TraceSpan).where(TraceSpan.trace_id.in_(old_trace_ids)))
+    trace_result = db.execute(delete(TraceRun).where(TraceRun.updated_at < cutoff))
+
+    old_session_ids = select(GuardSession.id).where(GuardSession.updated_at < cutoff)
+    old_request_ids = select(RequestRecord.id).where(RequestRecord.session_id.in_(old_session_ids))
+    db.execute(update(TraceRun).where(TraceRun.source_session_id.in_(old_session_ids)).values(source_session_id=None))
+    db.execute(delete(ToolCall).where(ToolCall.request_id.in_(old_request_ids)))
+    db.execute(delete(Event).where(Event.session_id.in_(old_session_ids)))
+    db.execute(delete(RequestRecord).where(RequestRecord.session_id.in_(old_session_ids)))
+    session_result = db.execute(delete(GuardSession).where(GuardSession.updated_at < cutoff))
+    db.commit()
+    return {
+        "traces": int(trace_result.rowcount or 0),
+        "sessions": int(session_result.rowcount or 0),
+    }

app/platform/migrations.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from alembic import command
+from alembic.config import Config
+from sqlalchemy import inspect
+
+
+def alembic_config(storage_url: str) -> Config:
+    script_location = Path(__file__).with_name("alembic")
+    config = Config()
+    config.set_main_option("script_location", str(script_location))
+    config.set_main_option("sqlalchemy.url", storage_url.replace("%", "%%"))
+    return config
+
+
+def migrate_database(engine, storage_url: str) -> None:
+    config = alembic_config(storage_url)
+    tables = set(inspect(engine).get_table_names())
+    if tables and "alembic_version" not in tables:
+        # v0.1 databases were created from the same SQLAlchemy metadata.
+        command.stamp(config, "0001_initial")
+    command.upgrade(config, "head")

app/platform/setup.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import json
+import secrets
+from pathlib import Path
+
+import yaml
+
+from app.core.config import SAMPLE_CONFIG
+
+
+def render_profiles(base_url: str, gateway_key: str) -> dict[str, str]:
+    openai_url = base_url.rstrip("/") + "/v1"
+    return {
+        "codex.toml": (
+            'model = "demo-model"\n'
+            'model_provider = "agent_loop_guard"\n\n'
+            "[model_providers.agent_loop_guard]\n"
+            'name = "Agent Loop Guard"\n'
+            f'base_url = "{openai_url}"\n'
+            'env_key = "ALG_GATEWAY_KEY"\n'
+        ),
+        "claude.env": (
+            f"ANTHROPIC_BASE_URL={base_url.rstrip('/')}\n"
+            f"ANTHROPIC_AUTH_TOKEN={gateway_key}\n"
+        ),
+        "cline.txt": (
+            "API Provider: OpenAI Compatible\n"
+            f"Base URL: {openai_url}\n"
+            f"API Key: {gateway_key}\n"
+            "Model ID: demo-model\n"
+        ),
+        "opencode.json": json.dumps(
+            {
+                "provider": {
+                    "alg": {
+                        "npm": "@ai-sdk/openai-compatible",
+                        "name": "Agent Loop Guard",
+                        "options": {"baseURL": openai_url, "apiKey": "{env:ALG_GATEWAY_KEY}"},
+                        "models": {"demo-model": {"name": "Guarded model"}},
+                    }
+                },
+                "model": "alg/demo-model",
+            },
+            indent=2,
+        )
+        + "\n",
+    }
+
+
+def setup_workspace(
+    root: Path,
+    *,
+    host: str = "127.0.0.1",
+    port: int = 8787,
+    gateway_key: str | None = None,
+    force: bool = False,
+) -> dict[str, object]:
+    root.mkdir(parents=True, exist_ok=True)
+    config_path = root / "agent-loop-guard.yml"
+    profiles_dir = root / ".agent-loop-guard" / "profiles"
+    profiles_dir.mkdir(parents=True, exist_ok=True)
+
+    if config_path.exists() and not force:
+        config_written = False
+        existing = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {}
+        effective_gateway_key = str(existing.get("gateway_key") or gateway_key or "alg_demo_key")
+    else:
+        effective_gateway_key = gateway_key or f"alg_{secrets.token_urlsafe(24)}"
+        config_text = SAMPLE_CONFIG.replace(
+            "gateway_key: alg_demo_key", f"gateway_key: {effective_gateway_key}"
+        )
+        config_text = config_text.replace("host: 127.0.0.1", f"host: {host}", 1)
+        config_text = config_text.replace("port: 8787", f"port: {port}", 1)
+        config_path.write_text(config_text, encoding="utf-8")
+        config_written = True
+
+    base_url = f"http://{host}:{port}"
+    written_profiles = []
+    for name, content in render_profiles(base_url, effective_gateway_key).items():
+        path = profiles_dir / name
+        if force or not path.exists():
+            path.write_text(content, encoding="utf-8")
+        written_profiles.append(str(path))
+
+    return {
+        "config": str(config_path),
+        "config_written": config_written,
+        "profiles": written_profiles,
+        "base_url": base_url,
+        "gateway_key": effective_gateway_key if config_written else None,
+    }

app/providers/__init__.py

@@ -0,0 +1,5 @@
+from app.providers.base import ProviderResult, ProviderStream
+from app.providers.mock import MockProvider
+from app.providers.upstream import UpstreamProvider
+
+__all__ = ["MockProvider", "ProviderResult", "ProviderStream", "UpstreamProvider"]

app/providers/base.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(slots=True)
+class ProviderResult:
+    status_code: int
+    headers: dict[str, str]
+    content: bytes
+    json_body: Any | None = None
+    media_type: str | None = None
+
+
+@dataclass(slots=True)
+class ProviderStream:
+    status_code: int
+    headers: dict[str, str]
+    chunks: AsyncIterator[bytes]
+    media_type: str = "text/event-stream"
+