kubiya-control-plane-api 0.9.15__py3-none-any.whl

scripts/migrations.py

@@ -0,0 +1,39 @@
+"""Database migrations entry point."""
+import os
+import sys
+from pathlib import Path
+
+
+def upgrade():
+    """Run alembic upgrade head."""
+    from alembic.config import Config
+    from alembic import command
+
+    # Get the package directory where alembic.ini lives
+    package_dir = Path(__file__).parent.parent / "control_plane_api"
+    alembic_ini = package_dir / "alembic.ini"
+
+    if not alembic_ini.exists():
+        print(f"Error: alembic.ini not found at {alembic_ini}")
+        sys.exit(1)
+
+    # Create Alembic config
+    alembic_cfg = Config(str(alembic_ini))
+
+    # Load .env file from current directory if it exists
+    env_file = Path.cwd() / ".env"
+    if env_file.exists():
+        from dotenv import load_dotenv
+        load_dotenv(env_file)
+
+    # Get database URL from environment
+    database_url = os.environ.get("DATABASE_URL")
+    if not database_url:
+        print("Error: DATABASE_URL environment variable is not set")
+        sys.exit(1)
+
+    # Set the database URL in the config
+    alembic_cfg.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
+
+    # Run the upgrade
+    command.upgrade(alembic_cfg, "head")

scripts/seed_worker_queues.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""
+Seed the environments and worker_queues tables with default data.
+
+This script populates the database with environments and worker queues for Kubiya agent control plane.
+"""
+
+import sys
+import os
+from pathlib import Path
+
+# Load .env file from current directory if it exists
+env_file = Path.cwd() / ".env"
+if env_file.exists():
+    from dotenv import load_dotenv
+    load_dotenv(env_file)
+
+# Add parent directory to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from control_plane_api.app.database import get_engine
+from control_plane_api.app.models.worker import WorkerQueue
+from control_plane_api.app.models.environment import Environment
+from sqlalchemy.orm import Session
+import uuid
+
+
+# Default environment ID used by worker queues
+DEFAULT_ENVIRONMENT_ID = uuid.UUID("89441d4d-639a-4f2f-b204-30b42b65175f")
+
+
+def seed_environments(session: Session) -> bool:
+    """Seed the database with default environments. Returns True if seeded."""
+
+    environments_data = [
+        {
+            "id": DEFAULT_ENVIRONMENT_ID,
+            "organization_id": "kubiya-ai",
+            "name": "default",
+            "display_name": "Default Environment",
+            "description": "Default environment for all workers",
+            "tags": [],
+            "settings": {},
+            "status": "active",
+            "created_by": "cb629448-4b7e-46cb-b321-dedefe79c381",
+            "worker_token": uuid.UUID("4034ee2f-137c-456d-8616-10c2e5abe404"),
+            "execution_environment": {},
+            "policy_ids": [],
+        },
+    ]
+
+    # Check if environments already exist
+    existing_count = session.query(Environment).count()
+    if existing_count > 0:
+        print(f"⚠️  Database already has {existing_count} environments. Skipping environment seed.")
+        return False
+
+    print("🌱 Seeding environments...")
+
+    for env_data in environments_data:
+        env = Environment(**env_data)
+        session.add(env)
+        print(f"   ✅ Added {env_data['display_name']} ({env_data['id']})")
+
+    session.commit()
+    print(f"\n✨ Successfully seeded {len(environments_data)} environments!")
+    return True
+
+
+def seed_worker_queues(session: Session) -> bool:
+    """Seed the database with default worker queues. Returns True if seeded."""
+
+    queues_data = [
+        {
+            "id": uuid.UUID("e54ecc24-a818-40b6-b554-51005319b293"),
+            "organization_id": "kubiya-ai",
+            "environment_id": DEFAULT_ENVIRONMENT_ID,
+            "name": "my-default-queue",
+            "display_name": "My Default Queue",
+            "description": "",
+            "status": "active",
+            "max_workers": 10,
+            "heartbeat_interval": 60,
+            "tags": [],
+            "settings": {},
+            "created_by": "cb629448-4b7e-46cb-b321-dedefe79c381",
+        },
+    ]
+
+    # Check if queues already exist
+    existing_count = session.query(WorkerQueue).count()
+    if existing_count > 0:
+        print(f"⚠️  Database already has {existing_count} worker queues. Skipping worker queue seed.")
+        return False
+
+    print("🌱 Seeding worker queues...")
+
+    for queue_data in queues_data:
+        queue = WorkerQueue(**queue_data)
+        session.add(queue)
+        print(f"   ✅ Added {queue_data['name']} ({queue_data['id']})")
+
+    session.commit()
+    print(f"\n✨ Successfully seeded {len(queues_data)} worker queues!")
+    return True
+
+
+def seed_all():
+    """Seed all tables in the correct order"""
+
+    engine = get_engine()
+
+    with Session(engine) as session:
+        # Seed environments first (worker_queues depends on it)
+        seed_environments(session)
+
+        # Then seed worker queues
+        seed_worker_queues(session)
+
+
+if __name__ == "__main__":
+    try:
+        seed_all()
+    except Exception as e:
+        print(f"❌ Error seeding data: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)

scripts/setup_agent_runtime.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""
+Setup script for agent-runtime binary.
+Called by Go CLI or Python worker before starting worker.
+
+This script:
+1. Downloads agent-runtime binary from GitHub releases
+2. Verifies SHA256 checksum
+3. Makes binary executable
+4. Optionally starts the server
+5. Returns server info as JSON
+
+Usage:
+    python setup_agent_runtime.py [config_dir] [--start-server] [--version VERSION]
+"""
+
+import asyncio
+import sys
+import json
+import argparse
+from pathlib import Path
+
+
+async def ensure_agent_runtime(
+    config_dir: str,
+    version: str = "latest",
+    start_server: bool = False,
+    grpc_port: int = 50052,
+    http_port: int = 8082,
+) -> dict:
+    """
+    Ensure agent-runtime binary available and optionally start server.
+
+    Args:
+        config_dir: Directory for config and binaries
+        version: Version to download (default: "latest")
+        start_server: Whether to start the server
+        grpc_port: GRPC server port
+        http_port: HTTP health check port
+
+    Returns:
+        Dictionary with binary info and server info (if started)
+    """
+    from control_plane_api.worker.binary_manager import BinaryManager
+    from control_plane_api.worker.agent_runtime_server import AgentRuntimeServer, ServerConfig
+
+    binary_manager = BinaryManager(Path(config_dir))
+
+    # Download/verify binary
+    binary_path = await binary_manager.ensure_binary(version)
+
+    result = {
+        "binary_path": str(binary_path),
+        "version": binary_manager._get_current_version(),
+        "executable": binary_path.is_file() and binary_path.stat().st_mode & 0o111,
+    }
+
+    if start_server:
+        # Start server
+        server_config = ServerConfig(
+            grpc_port=grpc_port,
+            http_port=http_port,
+            config_dir=Path(config_dir),
+        )
+
+        server = AgentRuntimeServer(binary_path, server_config)
+        await server.start(wait_for_health=True, timeout=30)
+
+        result.update({
+            "server_started": True,
+            "grpc_address": server.grpc_address,
+            "http_port": http_port,
+            "pid": server.process.pid if server.process else None,
+            "config_file": str(server.config_file),
+        })
+    else:
+        result["server_started"] = False
+
+    return result
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Setup agent-runtime binary and optionally start server"
+    )
+    parser.add_argument(
+        "config_dir",
+        nargs="?",
+        default=str(Path.home() / ".kubiya"),
+        help="Config directory (default: ~/.kubiya)",
+    )
+    parser.add_argument(
+        "--start-server",
+        action="store_true",
+        help="Start agent-runtime server after download",
+    )
+    parser.add_argument(
+        "--version",
+        default="latest",
+        help="Version to download (default: latest)",
+    )
+    parser.add_argument(
+        "--grpc-port",
+        type=int,
+        default=50052,
+        help="GRPC server port (default: 50052)",
+    )
+    parser.add_argument(
+        "--http-port",
+        type=int,
+        default=8082,
+        help="HTTP health check port (default: 8082)",
+    )
+
+    args = parser.parse_args()
+
+    try:
+        result = asyncio.run(
+            ensure_agent_runtime(
+                config_dir=args.config_dir,
+                version=args.version,
+                start_server=args.start_server,
+                grpc_port=args.grpc_port,
+                http_port=args.http_port,
+            )
+        )
+
+        # Output JSON for Go CLI or other callers to parse
+        print(json.dumps(result, indent=2))
+        sys.exit(0)
+
+    except Exception as e:
+        error_result = {
+            "success": False,
+            "error": str(e),
+        }
+        print(json.dumps(error_result, indent=2), file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

worker_internal/__init__.py

@@ -0,0 +1 @@
+"""Worker Internal - Internal orchestration workers for Kubiya platform."""

worker_internal/planner/__init__.py

@@ -0,0 +1 @@
+"""Plan Orchestrator - Intelligent plan execution using Claude Code Agent SDK."""