nightshift-sdk 0.2.0__py3-none-any.whl

nightshift/cli/config.py

@@ -0,0 +1,52 @@
+"""CLI config — reads/writes ~/.nightshift/config.toml."""
+
+from __future__ import annotations
+
+import os
+import stat
+import tomllib
+
+import tomli_w
+
+
+CONFIG_DIR = os.path.expanduser("~/.nightshift")
+CONFIG_PATH = os.path.join(CONFIG_DIR, "config.toml")
+
+
+def load_config() -> dict:
+    """Load the CLI config file, returning {} if it doesn't exist."""
+    if not os.path.exists(CONFIG_PATH):
+        return {}
+    with open(CONFIG_PATH, "rb") as f:
+        return tomllib.load(f)
+
+
+def save_config(data: dict) -> None:
+    """Write the CLI config file with restricted permissions (0600)."""
+    os.makedirs(CONFIG_DIR, exist_ok=True)
+    with open(CONFIG_PATH, "wb") as f:
+        tomli_w.dump(data, f)
+    os.chmod(CONFIG_PATH, stat.S_IRUSR | stat.S_IWUSR)
+
+
+def get_url() -> str:
+    """Get the platform URL from config."""
+    cfg = load_config()
+    url = cfg.get("url", "")
+    if not url:
+        raise SystemExit("Not logged in. Run: nightshift login --url <URL> --api-key <KEY>")
+    return url
+
+
+def get_api_key() -> str:
+    """Get the API key from config."""
+    cfg = load_config()
+    key = cfg.get("api_key", "")
+    if not key:
+        raise SystemExit("Not logged in. Run: nightshift login --url <URL> --api-key <KEY>")
+    return key
+
+
+def get_auth_headers() -> dict[str, str]:
+    """Get Authorization headers for API requests."""
+    return {"Authorization": f"Bearer {get_api_key()}"}

nightshift/cli/main.py

@@ -0,0 +1,39 @@
+"""Nightshift CLI — deploy and manage agents on the platform."""
+
+from __future__ import annotations
+
+import click
+
+from nightshift.cli.commands.agents import agents
+from nightshift.cli.commands.api_key import api_key
+from nightshift.cli.commands.deploy import deploy
+from nightshift.cli.commands.login import login
+from nightshift.cli.commands.logs import logs
+from nightshift.cli.commands.run import run
+
+
+@click.group()
+def cli() -> None:
+    """Nightshift — autonomous agent orchestrator with Firecracker VMs."""
+    pass
+
+
+# Register subcommands
+cli.add_command(login)
+cli.add_command(deploy)
+cli.add_command(run)
+cli.add_command(logs)
+cli.add_command(agents)
+cli.add_command(api_key)
+
+
+@cli.command()
+@click.option("--host", default="0.0.0.0", help="Bind host")
+@click.option("--port", default=3000, type=int, help="Bind port")
+def serve(host: str, port: int) -> None:
+    """Start the Nightshift platform server."""
+    import asyncio
+
+    from nightshift.server import start_server
+
+    asyncio.run(start_server(host=host, port=port))

nightshift/config.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+@dataclass
+class NightshiftConfig:
+    workspace: str = ""
+    port: int = 3000
+    kernel_path: str = "/opt/nightshift/vmlinux"
+    base_rootfs_path: str = "/opt/nightshift/rootfs.ext4"
+    vm_timeout_seconds: int = 1800  # 30 minutes
+    vm_health_timeout_seconds: int = 60
+    vm_event_port: int = 8080
+    db_path: str = "/opt/nightshift/nightshift.db"
+    agents_storage_dir: str = "/opt/nightshift/agents"
+
+    @staticmethod
+    def from_env() -> NightshiftConfig:
+        return NightshiftConfig(
+            workspace=os.environ.get("NIGHTSHIFT_WORKSPACE", os.getcwd()),
+            port=int(os.environ.get("NIGHTSHIFT_PORT", "3000")),
+            kernel_path=os.environ.get("NIGHTSHIFT_KERNEL_PATH", "/opt/nightshift/vmlinux"),
+            base_rootfs_path=os.environ.get(
+                "NIGHTSHIFT_ROOTFS_PATH", "/opt/nightshift/rootfs.ext4"
+            ),
+            db_path=os.environ.get("NIGHTSHIFT_DB_PATH", "/opt/nightshift/nightshift.db"),
+            agents_storage_dir=os.environ.get(
+                "NIGHTSHIFT_AGENTS_DIR", "/opt/nightshift/agents"
+            ),
+        )
+
+    def env_vars_for_vm(self) -> dict[str, str]:
+        """Collect platform env vars to pass into the VM."""
+        env: dict[str, str] = {}
+        env["NIGHTSHIFT_WORKSPACE"] = "/workspace"
+        return env

nightshift/events.py

@@ -0,0 +1,124 @@
+"""Event types and in-memory event buffer.
+
+EventBuffer replaces the previous EventLog/EventStore/EventSink stack
+with a single class that stores events per run_id and supports async streaming.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from dataclasses import asdict, dataclass, field
+from typing import AsyncIterator, Literal, Union
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BaseEvent:
+    timestamp: float = field(default_factory=time.time)
+    run_id: str | None = None
+
+
+@dataclass
+class StartedEvent(BaseEvent):
+    type: Literal["nightshift.started"] = "nightshift.started"
+    workspace: str = ""
+
+
+@dataclass
+class CompletedEvent(BaseEvent):
+    type: Literal["nightshift.completed"] = "nightshift.completed"
+
+
+@dataclass
+class ErrorEvent(BaseEvent):
+    type: Literal["nightshift.error"] = "nightshift.error"
+    error: str = ""
+
+
+@dataclass
+class InterruptedEvent(BaseEvent):
+    type: Literal["nightshift.interrupted"] = "nightshift.interrupted"
+    reason: Literal["user_quit", "user_stop"] = "user_quit"
+
+
+NightshiftEvent = Union[
+    StartedEvent,
+    CompletedEvent,
+    ErrorEvent,
+    InterruptedEvent,
+]
+
+TERMINAL_EVENTS = {"nightshift.completed", "nightshift.error", "nightshift.interrupted"}
+
+
+class EventBuffer:
+    """Simple in-memory event buffer with async streaming.
+
+    Stores events per run_id and allows consumers to stream them
+    (replay + live-tail) via an asyncio.Condition.
+    """
+
+    def __init__(self) -> None:
+        self._runs: dict[str, list[tuple[str, dict]]] = {}
+        self._cond: asyncio.Condition = asyncio.Condition()
+        self._done: set[str] = set()
+
+    async def append(self, run_id: str, event_type: str, payload: dict) -> None:
+        """Append an event to a run's buffer and notify waiters."""
+        self._runs.setdefault(run_id, []).append((event_type, payload))
+        async with self._cond:
+            self._cond.notify_all()
+
+    async def stream(self, run_id: str, cursor: int = 0) -> AsyncIterator[tuple[str, dict]]:
+        """Yield events for a run, replaying from cursor then live-tailing.
+
+        Terminates when the run is marked done via cleanup().
+        """
+        while True:
+            events = self._runs.get(run_id, [])
+            while cursor < len(events):
+                yield events[cursor]
+                cursor += 1
+            if run_id in self._done:
+                return
+            async with self._cond:
+                await self._cond.wait()
+
+    async def cleanup(self, run_id: str) -> None:
+        """Mark a run as done and remove its events."""
+        self._runs.pop(run_id, None)
+        self._done.add(run_id)
+        async with self._cond:
+            self._cond.notify_all()
+
+    # ── Convenience methods (used by task.py, vm/manager.py, agent/entry.py) ──
+
+    async def publish(self, run_id: str, event: NightshiftEvent) -> None:
+        """Publish a typed NightshiftEvent."""
+        event.run_id = run_id
+        payload = asdict(event)
+        event_type = payload.pop("type")
+        await self.append(run_id, event_type, payload)
+
+    async def publish_raw(self, run_id: str, event_type: str, data: dict) -> None:
+        """Forward a raw event dict (e.g. from the guest SSE stream)."""
+        payload = {k: v for k, v in data.items() if k != "type"}
+        await self.append(run_id, event_type, payload)
+
+    async def stream_sse(self, run_id: str) -> AsyncIterator[dict]:
+        """Stream events formatted for SSE (event + data keys)."""
+        async for event_type, payload in self.stream(run_id):
+            yield {
+                "event": event_type,
+                "data": json.dumps({"type": event_type, **payload}),
+            }
+            if event_type in TERMINAL_EVENTS:
+                return
+
+
+# Backwards-compatible alias used by existing code
+EventLog = EventBuffer

nightshift/protocol/__init__.py

@@ -0,0 +1,4 @@
+from nightshift.protocol.events import serialize_message
+from nightshift.protocol.packaging import package_agent
+
+__all__ = ["serialize_message", "package_agent"]

nightshift/protocol/events.py

@@ -0,0 +1,36 @@
+"""Serialize Claude Agent SDK messages to JSON dicts for SSE transport."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import asdict, is_dataclass
+from typing import Any
+
+
+def serialize_message(message: Any) -> dict:
+    """Convert an agent message to a JSON-serializable dict.
+
+    Handles:
+    - dicts (passed through with defaults)
+    - dataclasses (converted via asdict)
+    - pydantic models (converted via model_dump)
+    - objects with a __dict__ attribute
+    """
+    if isinstance(message, dict):
+        data = dict(message)
+    elif is_dataclass(message) and not isinstance(message, type):
+        data = asdict(message)
+    elif hasattr(message, "model_dump"):
+        data = message.model_dump()
+    elif hasattr(message, "__dict__"):
+        data = dict(message.__dict__)
+    else:
+        data = {"data": str(message)}
+
+    if "type" not in data:
+        data["type"] = type(message).__name__
+
+    if "timestamp" not in data:
+        data["timestamp"] = time.time()
+
+    return data

nightshift/protocol/packaging.py

@@ -0,0 +1,71 @@
+"""Package agent source code for injection into a Firecracker VM."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import tempfile
+
+
+def _find_pyproject(start_path: str) -> str | None:
+    """Walk up from start_path to find the nearest pyproject.toml."""
+    directory = os.path.dirname(os.path.abspath(start_path))
+    while True:
+        candidate = os.path.join(directory, "pyproject.toml")
+        if os.path.isfile(candidate):
+            return candidate
+        parent = os.path.dirname(directory)
+        if parent == directory:
+            return None
+        directory = parent
+
+
+def package_agent(
+    module_path: str,
+    function_name: str,
+    prompt: str,
+) -> str:
+    """Package an agent's source file and a manifest for VM injection.
+
+    Creates a temp directory containing:
+    - The agent's source file (copied)
+    - manifest.json
+    - pyproject.toml (auto-detected from agent's project, if found)
+    - uv.lock (if present alongside pyproject.toml)
+
+    Returns the path to the temp directory.
+    """
+    pkg_dir = tempfile.mkdtemp(prefix="nightshift-pkg-")
+
+    # Copy the agent source file
+    source_filename = os.path.basename(module_path)
+    shutil.copy2(module_path, os.path.join(pkg_dir, source_filename))
+
+    # Auto-detect and copy pyproject.toml + uv.lock
+    pyproject_path = _find_pyproject(module_path)
+    has_pyproject = False
+    if pyproject_path:
+        shutil.copy2(pyproject_path, os.path.join(pkg_dir, "pyproject.toml"))
+        has_pyproject = True
+        project_dir = os.path.dirname(pyproject_path)
+        lock_path = os.path.join(project_dir, "uv.lock")
+        if os.path.isfile(lock_path):
+            shutil.copy2(lock_path, os.path.join(pkg_dir, "uv.lock"))
+
+    # Write manifest
+    manifest = {
+        "module": os.path.splitext(source_filename)[0],
+        "function": function_name,
+        "prompt": prompt,
+        "has_pyproject": has_pyproject,
+    }
+    with open(os.path.join(pkg_dir, "manifest.json"), "w") as f:
+        json.dump(manifest, f)
+
+    return pkg_dir
+
+
+def cleanup_package(pkg_dir: str) -> None:
+    """Remove a package directory created by package_agent."""
+    shutil.rmtree(pkg_dir, ignore_errors=True)

nightshift/registry.py

@@ -0,0 +1,267 @@
+"""SQLite-backed agent and run registry."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timezone
+
+import aiosqlite
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _uuid() -> str:
+    return str(uuid.uuid4())
+
+
+@dataclass
+class AgentRecord:
+    id: str
+    tenant_id: str
+    name: str
+    source_filename: str
+    function_name: str
+    config_json: str
+    storage_path: str
+    created_at: str
+    updated_at: str
+
+
+@dataclass
+class RunRecord:
+    id: str
+    agent_id: str
+    tenant_id: str
+    prompt: str
+    status: str
+    created_at: str
+    completed_at: str | None = None
+    error: str | None = None
+
+
+class AgentRegistry:
+    """Async SQLite registry for agents, runs, and API keys."""
+
+    def __init__(self, db_path: str) -> None:
+        self.db_path = db_path
+        self._db: aiosqlite.Connection | None = None
+
+    async def init_db(self) -> None:
+        """Open the database and create tables if needed."""
+        self._db = await aiosqlite.connect(self.db_path)
+        self._db.row_factory = aiosqlite.Row
+        await self._db.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS agents (
+                id TEXT PRIMARY KEY,
+                tenant_id TEXT NOT NULL,
+                name TEXT NOT NULL,
+                source_filename TEXT NOT NULL,
+                function_name TEXT NOT NULL,
+                config_json TEXT NOT NULL,
+                storage_path TEXT NOT NULL,
+                created_at TEXT NOT NULL,
+                updated_at TEXT NOT NULL,
+                UNIQUE(tenant_id, name)
+            );
+
+            CREATE TABLE IF NOT EXISTS runs (
+                id TEXT PRIMARY KEY,
+                agent_id TEXT NOT NULL,
+                tenant_id TEXT NOT NULL,
+                prompt TEXT NOT NULL,
+                status TEXT NOT NULL DEFAULT 'started',
+                created_at TEXT NOT NULL,
+                completed_at TEXT,
+                error TEXT
+            );
+
+            CREATE TABLE IF NOT EXISTS api_keys (
+                key_hash TEXT PRIMARY KEY,
+                tenant_id TEXT NOT NULL,
+                label TEXT DEFAULT '',
+                created_at TEXT NOT NULL
+            );
+            """
+        )
+        await self._db.commit()
+
+    async def close(self) -> None:
+        if self._db:
+            await self._db.close()
+
+    @property
+    def db(self) -> aiosqlite.Connection:
+        if self._db is None:
+            raise RuntimeError("Database not initialized — call init_db() first")
+        return self._db
+
+    # ── Agents ────────────────────────────────────────────────────
+
+    async def upsert_agent(
+        self,
+        tenant_id: str,
+        name: str,
+        source_filename: str,
+        function_name: str,
+        config_json: str,
+        storage_path: str,
+        agent_id: str | None = None,
+    ) -> AgentRecord:
+        """Insert or update an agent by (tenant_id, name)."""
+        now = _now()
+        if agent_id is None:
+            agent_id = _uuid()
+
+        # Check for existing agent with same (tenant_id, name)
+        row = await self.db.execute_fetchall(
+            "SELECT id, created_at FROM agents WHERE tenant_id = ? AND name = ?",
+            (tenant_id, name),
+        )
+        if row:
+            existing_id = row[0][0]
+            existing_created = row[0][1]
+            await self.db.execute(
+                """UPDATE agents
+                   SET source_filename = ?, function_name = ?, config_json = ?,
+                       storage_path = ?, updated_at = ?
+                   WHERE id = ?""",
+                (source_filename, function_name, config_json, storage_path, now, existing_id),
+            )
+            await self.db.commit()
+            return AgentRecord(
+                id=existing_id,
+                tenant_id=tenant_id,
+                name=name,
+                source_filename=source_filename,
+                function_name=function_name,
+                config_json=config_json,
+                storage_path=storage_path,
+                created_at=existing_created,
+                updated_at=now,
+            )
+
+        await self.db.execute(
+            """INSERT INTO agents (id, tenant_id, name, source_filename, function_name,
+                                   config_json, storage_path, created_at, updated_at)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            (agent_id, tenant_id, name, source_filename, function_name,
+             config_json, storage_path, now, now),
+        )
+        await self.db.commit()
+        return AgentRecord(
+            id=agent_id,
+            tenant_id=tenant_id,
+            name=name,
+            source_filename=source_filename,
+            function_name=function_name,
+            config_json=config_json,
+            storage_path=storage_path,
+            created_at=now,
+            updated_at=now,
+        )
+
+    async def get_agent(self, tenant_id: str, name: str) -> AgentRecord | None:
+        rows = await self.db.execute_fetchall(
+            "SELECT * FROM agents WHERE tenant_id = ? AND name = ?",
+            (tenant_id, name),
+        )
+        if not rows:
+            return None
+        r = rows[0]
+        return AgentRecord(
+            id=r[0], tenant_id=r[1], name=r[2], source_filename=r[3],
+            function_name=r[4], config_json=r[5], storage_path=r[6],
+            created_at=r[7], updated_at=r[8],
+        )
+
+    async def list_agents(self, tenant_id: str) -> list[AgentRecord]:
+        rows = await self.db.execute_fetchall(
+            "SELECT * FROM agents WHERE tenant_id = ? ORDER BY name",
+            (tenant_id,),
+        )
+        return [
+            AgentRecord(
+                id=r[0], tenant_id=r[1], name=r[2], source_filename=r[3],
+                function_name=r[4], config_json=r[5], storage_path=r[6],
+                created_at=r[7], updated_at=r[8],
+            )
+            for r in rows
+        ]
+
+    async def delete_agent(self, tenant_id: str, name: str) -> bool:
+        cursor = await self.db.execute(
+            "DELETE FROM agents WHERE tenant_id = ? AND name = ?",
+            (tenant_id, name),
+        )
+        await self.db.commit()
+        return cursor.rowcount > 0
+
+    # ── Runs ──────────────────────────────────────────────────────
+
+    async def create_run(
+        self,
+        agent_id: str,
+        tenant_id: str,
+        prompt: str,
+    ) -> RunRecord:
+        run_id = _uuid()
+        now = _now()
+        await self.db.execute(
+            """INSERT INTO runs (id, agent_id, tenant_id, prompt, status, created_at)
+               VALUES (?, ?, ?, ?, 'started', ?)""",
+            (run_id, agent_id, tenant_id, prompt, now),
+        )
+        await self.db.commit()
+        return RunRecord(
+            id=run_id,
+            agent_id=agent_id,
+            tenant_id=tenant_id,
+            prompt=prompt,
+            status="started",
+            created_at=now,
+        )
+
+    async def complete_run(self, run_id: str, error: str | None = None) -> None:
+        now = _now()
+        status = "error" if error else "completed"
+        await self.db.execute(
+            "UPDATE runs SET status = ?, completed_at = ?, error = ? WHERE id = ?",
+            (status, now, error, run_id),
+        )
+        await self.db.commit()
+
+    async def get_run(self, run_id: str) -> RunRecord | None:
+        rows = await self.db.execute_fetchall(
+            "SELECT * FROM runs WHERE id = ?", (run_id,),
+        )
+        if not rows:
+            return None
+        r = rows[0]
+        return RunRecord(
+            id=r[0], agent_id=r[1], tenant_id=r[2], prompt=r[3],
+            status=r[4], created_at=r[5], completed_at=r[6], error=r[7],
+        )
+
+    # ── API Keys ──────────────────────────────────────────────────
+
+    async def store_api_key(self, key_hash: str, tenant_id: str, label: str = "") -> None:
+        await self.db.execute(
+            """INSERT OR IGNORE INTO api_keys (key_hash, tenant_id, label, created_at)
+               VALUES (?, ?, ?, ?)""",
+            (key_hash, tenant_id, label, _now()),
+        )
+        await self.db.commit()
+
+    async def get_tenant_by_key_hash(self, key_hash: str) -> str | None:
+        rows = await self.db.execute_fetchall(
+            "SELECT tenant_id FROM api_keys WHERE key_hash = ?",
+            (key_hash,),
+        )
+        if not rows:
+            return None
+        return rows[0][0]

nightshift/sdk/__init__.py

@@ -0,0 +1,4 @@
+from nightshift.sdk.app import NightshiftApp
+from nightshift.sdk.config import AgentConfig
+
+__all__ = ["NightshiftApp", "AgentConfig"]

nightshift/sdk/app.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass, field
+from collections.abc import Callable
+from types import FunctionType
+from typing import Any
+
+from nightshift.sdk.config import AgentConfig
+
+
+@dataclass
+class RegisteredAgent:
+    """An agent function registered with NightshiftApp."""
+
+    name: str
+    fn: Callable[..., Any]
+    config: AgentConfig
+    module_path: str
+
+
+class NightshiftApp:
+    """Main application class for registering and serving agents."""
+
+    def __init__(self) -> None:
+        self._agents: dict[str, RegisteredAgent] = {}
+
+    def agent(self, config: AgentConfig | None = None, *, name: str | None = None):
+        """Decorator to register an async generator function as an agent.
+
+        Usage:
+            @app.agent(AgentConfig(vcpu_count=2))
+            async def my_agent(prompt: str):
+                yield message
+        """
+        if config is None:
+            config = AgentConfig()
+
+        def decorator(fn: FunctionType) -> FunctionType:
+            agent_name = name or fn.__name__
+            module_path = inspect.getfile(fn)
+            self._agents[agent_name] = RegisteredAgent(
+                name=agent_name,
+                fn=fn,
+                config=config,
+                module_path=module_path,
+            )
+            return fn
+
+        return decorator
+