nightshift-sdk 0.2.0__py3-none-any.whl

nightshift/__init__.py

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

nightshift/__main__.py

@@ -0,0 +1,8 @@
+"""Starts the Nightshift agent entry point"""
+
+from __future__ import annotations
+
+from nightshift.agent.entry import main
+
+if __name__ == "__main__":
+    main()

nightshift/agent/entry.py

@@ -0,0 +1,116 @@
+"""In-VM entry point.
+
+Runs inside the Firecracker VM on boot. Reads a manifest,
+dynamically imports the agent function, calls it with the prompt,
+and streams each yielded message out via SSE on :8080.
+
+Usage: python -m nightshift.agent
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import json
+import os
+import subprocess
+import sys
+
+from nightshift.events import (
+    CompletedEvent,
+    ErrorEvent,
+    EventLog,
+    StartedEvent,
+)
+from nightshift.protocol.events import serialize_message
+
+
+async def serve_events(log: EventLog, run_id: str = "vm", port: int = 8080) -> None:
+    """Serve an HTTP /events SSE endpoint for the host to subscribe to."""
+    import uvicorn
+    from fastapi import FastAPI
+    from sse_starlette.sse import EventSourceResponse
+
+    app = FastAPI()
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+    @app.get("/events")
+    async def events():
+        return EventSourceResponse(log.stream(run_id))
+
+    config = uvicorn.Config(app, host="0.0.0.0", port=port, log_level="warning")
+    server = uvicorn.Server(config)
+    await server.serve()
+
+
+async def run_agent() -> None:
+    """Main entry point for the in-VM agent."""
+    workspace = os.environ.get("NIGHTSHIFT_WORKSPACE", "/workspace")
+    agent_dir = os.environ.get("NIGHTSHIFT_AGENT_DIR", "/opt/nightshift/agent_pkg")
+    manifest_path = os.path.join(agent_dir, "manifest.json")
+
+    # Read manifest
+    with open(manifest_path) as f:
+        manifest = json.load(f)
+
+    module_name = manifest["module"]
+    function_name = manifest["function"]
+    prompt = manifest["prompt"]
+
+    log = EventLog()
+    run_id = "vm"
+
+    # Start event server in background
+    event_server_task = asyncio.create_task(serve_events(log, run_id=run_id, port=8080))
+    await asyncio.sleep(0.5)
+
+    try:
+        await log.publish(
+            run_id,
+            StartedEvent(workspace=workspace),
+        )
+
+        # Install agent dependencies from pyproject.toml if present
+        if manifest.get("has_pyproject"):
+            subprocess.run(
+                ["uv", "sync", "--project", agent_dir],
+                cwd=agent_dir,
+                check=True,
+                capture_output=True,
+            )
+
+        # Add agent dir to sys.path and import the agent module
+        if agent_dir not in sys.path:
+            sys.path.insert(0, agent_dir)
+        mod = importlib.import_module(module_name)
+        agent_fn = getattr(mod, function_name)
+
+        # Call the agent function (async generator) and stream messages
+        async for message in agent_fn(prompt):
+            data = serialize_message(message)
+            event_type = data.get("type", "agent.message")
+            await log.publish_raw(run_id, event_type, data)
+
+        await log.publish(run_id, CompletedEvent())
+
+    except Exception as e:
+        await log.publish(run_id, ErrorEvent(error=str(e)))
+
+    finally:
+        await asyncio.sleep(1)
+        event_server_task.cancel()
+        try:
+            await event_server_task
+        except asyncio.CancelledError:
+            pass
+
+
+def main() -> None:
+    asyncio.run(run_agent())
+
+
+if __name__ == "__main__":
+    main()

nightshift/auth.py

@@ -0,0 +1,57 @@
+"""API key authentication for the platform server."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import secrets
+
+from fastapi import Header, HTTPException
+
+from nightshift.registry import AgentRegistry
+
+
+def hash_api_key(key: str) -> str:
+    """SHA-256 hash of an API key."""
+    return hashlib.sha256(key.encode()).hexdigest()
+
+
+def generate_api_key() -> str:
+    """Generate a new API key in ns_<32 hex chars> format."""
+    return f"ns_{secrets.token_hex(16)}"
+
+
+async def bootstrap_api_key(registry: AgentRegistry) -> None:
+    """Seed the first API key from NIGHTSHIFT_API_KEY env var if set.
+
+    Idempotent — won't duplicate if the key already exists.
+    """
+    raw_key = os.environ.get("NIGHTSHIFT_API_KEY")
+    if not raw_key:
+        return
+
+    key_hash = hash_api_key(raw_key)
+    tenant_id = "default"
+    await registry.store_api_key(key_hash, tenant_id, label="bootstrap")
+
+
+async def get_tenant_id(
+    registry: AgentRegistry,
+    authorization: str = Header(),
+) -> str:
+    """Extract and verify the API key from the Authorization header.
+
+    Expected format: 'Bearer ns_...'
+    Returns the tenant_id if valid, raises 401 otherwise.
+    """
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid authorization header")
+
+    token = authorization[7:]  # strip "Bearer "
+    key_hash = hash_api_key(token)
+    tenant_id = await registry.get_tenant_by_key_hash(key_hash)
+
+    if tenant_id is None:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    return tenant_id

nightshift/cli/commands/agents.py

@@ -0,0 +1,56 @@
+"""nightshift agents — list and manage deployed agents."""
+
+from __future__ import annotations
+
+import click
+import httpx
+
+from nightshift.cli.config import get_auth_headers, get_url
+
+
+@click.group(invoke_without_command=True)
+@click.pass_context
+def agents(ctx) -> None:
+    """List deployed agents. Use 'agents rm' to delete one."""
+    if ctx.invoked_subcommand is not None:
+        return
+
+    url = get_url()
+    headers = get_auth_headers()
+
+    try:
+        r = httpx.get(f"{url}/api/agents", headers=headers, timeout=30)
+        r.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        raise click.ClickException(f"{e.response.status_code}: {e.response.text}")
+    except httpx.HTTPError as e:
+        raise click.ClickException(str(e))
+
+    agent_list = r.json()
+    if not agent_list:
+        click.echo("No agents deployed.")
+        return
+
+    # Table format
+    click.echo(f"{'NAME':<20} {'SOURCE':<25} {'UPDATED':<25}")
+    click.echo("-" * 70)
+    for a in agent_list:
+        click.echo(f"{a['name']:<20} {a['source_filename']:<25} {a['updated_at']:<25}")
+
+
+@agents.command()
+@click.argument("name")
+def rm(name: str) -> None:
+    """Delete a deployed agent."""
+    url = get_url()
+    headers = get_auth_headers()
+
+    try:
+        r = httpx.delete(f"{url}/api/agents/{name}", headers=headers, timeout=30)
+        r.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        raise click.ClickException(f"{e.response.status_code}: {e.response.text}")
+    except httpx.HTTPError as e:
+        raise click.ClickException(str(e))
+
+    click.echo(f"Deleted agent: {name}")

nightshift/cli/commands/api_key.py

@@ -0,0 +1,104 @@
+"""nightshift api-key — generate and manage API keys.
+
+These commands operate directly on the SQLite DB (no running server required).
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+import click
+
+from nightshift.auth import generate_api_key, hash_api_key
+from nightshift.config import NightshiftConfig
+from nightshift.registry import AgentRegistry
+
+
+async def _open_registry() -> AgentRegistry:
+    config = NightshiftConfig.from_env()
+    registry = AgentRegistry(config.db_path)
+    await registry.init_db()
+    return registry
+
+
+@click.group("api-key")
+def api_key() -> None:
+    """Generate and manage API keys (runs against the local DB)."""
+    pass
+
+
+@api_key.command()
+@click.option("--tenant", default="default", help="Tenant ID to associate with the key")
+@click.option("--label", default="", help="Human-readable label for the key")
+def generate(tenant: str, label: str) -> None:
+    """Generate a new API key and print it.
+
+    The raw key is only shown once — store it somewhere safe.
+    """
+
+    async def _run():
+        registry = await _open_registry()
+        try:
+            raw_key = generate_api_key()
+            key_hash = hash_api_key(raw_key)
+            await registry.store_api_key(key_hash, tenant, label=label)
+            return raw_key
+        finally:
+            await registry.close()
+
+    raw_key = asyncio.run(_run())
+    click.echo(raw_key)
+
+
+@api_key.command("list")
+def list_keys() -> None:
+    """List all API keys (shows hash prefix, tenant, label, created_at)."""
+
+    async def _run():
+        registry = await _open_registry()
+        try:
+            rows = await registry.db.execute_fetchall(
+                "SELECT key_hash, tenant_id, label, created_at FROM api_keys ORDER BY created_at"
+            )
+            return rows
+        finally:
+            await registry.close()
+
+    rows = asyncio.run(_run())
+    if not rows:
+        click.echo("No API keys found.")
+        return
+
+    click.echo(f"{'HASH PREFIX':<16} {'TENANT':<16} {'LABEL':<20} {'CREATED'}")
+    click.echo("-" * 76)
+    for r in rows:
+        click.echo(f"{r[0][:12]}...  {r[1]:<16} {r[2]:<20} {r[3]}")
+
+
+@api_key.command()
+@click.argument("hash_prefix")
+def revoke(hash_prefix: str) -> None:
+    """Revoke an API key by its hash prefix (from 'api-key list')."""
+
+    async def _run():
+        registry = await _open_registry()
+        try:
+            rows = await registry.db.execute_fetchall(
+                "SELECT key_hash FROM api_keys WHERE key_hash LIKE ?",
+                (f"{hash_prefix}%",),
+            )
+            if not rows:
+                raise click.ClickException(f"No key found matching prefix: {hash_prefix}")
+            if len(rows) > 1:
+                raise click.ClickException(
+                    f"Prefix '{hash_prefix}' matches {len(rows)} keys — be more specific"
+                )
+            key_hash = rows[0][0]
+            await registry.db.execute("DELETE FROM api_keys WHERE key_hash = ?", (key_hash,))
+            await registry.db.commit()
+            return key_hash
+        finally:
+            await registry.close()
+
+    key_hash = asyncio.run(_run())
+    click.echo(f"Revoked key {key_hash[:12]}...")

nightshift/cli/commands/deploy.py

@@ -0,0 +1,151 @@
+"""nightshift deploy — deploy an agent to the platform."""
+
+from __future__ import annotations
+
+import importlib.util
+import io
+import json
+import os
+import sys
+import tarfile
+
+import click
+import httpx
+
+from nightshift.cli.config import get_auth_headers, get_url
+
+# Patterns to exclude from the archive
+EXCLUDE_PATTERNS = {".git", "__pycache__", ".venv", ".env", "*.pyc", "node_modules", ".ruff_cache"}
+
+
+def _should_exclude(path: str) -> bool:
+    """Check if a path should be excluded from the archive."""
+    parts = path.split(os.sep)
+    for part in parts:
+        if part in EXCLUDE_PATTERNS:
+            return True
+        for pattern in EXCLUDE_PATTERNS:
+            if pattern.startswith("*") and part.endswith(pattern[1:]):
+                return True
+    return False
+
+
+def _make_archive(project_dir: str) -> bytes:
+    """Create a tar.gz archive of the project directory."""
+    buf = io.BytesIO()
+    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
+        for root, dirs, files in os.walk(project_dir):
+            # Filter out excluded directories in-place
+            dirs[:] = [d for d in dirs if not _should_exclude(d)]
+
+            for f in files:
+                if _should_exclude(f):
+                    continue
+                full_path = os.path.join(root, f)
+                arcname = os.path.relpath(full_path, project_dir)
+                tar.add(full_path, arcname=arcname)
+    buf.seek(0)
+    return buf.read()
+
+
+def _discover_agents(file_path: str) -> dict:
+    """Import a file and find the NightshiftApp instance to discover agents."""
+    file_path = os.path.abspath(file_path)
+    module_name = os.path.splitext(os.path.basename(file_path))[0]
+
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None or spec.loader is None:
+        raise click.ClickException(f"Cannot import {file_path}")
+
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = mod
+    spec.loader.exec_module(mod)
+
+    # Find the NightshiftApp instance
+    from nightshift.sdk.app import NightshiftApp
+
+    app = None
+    for attr_name in dir(mod):
+        obj = getattr(mod, attr_name)
+        if isinstance(obj, NightshiftApp):
+            app = obj
+            break
+
+    if app is None:
+        raise click.ClickException(f"No NightshiftApp instance found in {file_path}")
+
+    if not app._agents:
+        raise click.ClickException(f"No agents registered in {file_path}")
+
+    return {
+        name: {
+            "function_name": agent.fn.__name__,
+            "config": {
+                "workspace": agent.config.workspace,
+                "vcpu_count": agent.config.vcpu_count,
+                "mem_size_mib": agent.config.mem_size_mib,
+                "timeout_seconds": agent.config.timeout_seconds,
+                "forward_env": agent.config.forward_env,
+                "env": agent.config.env,
+            },
+        }
+        for name, agent in app._agents.items()
+    }
+
+
+@click.command()
+@click.argument("file", type=click.Path(exists=True))
+def deploy(file: str) -> None:
+    """Deploy agents from a Python file to the platform.
+
+    FILE is the path to a Python file containing a NightshiftApp with
+    registered agents (e.g. examples/basic_agent.py).
+    """
+    url = get_url()
+    headers = get_auth_headers()
+
+    click.echo(f"Discovering agents in {file}...")
+    agents = _discover_agents(file)
+
+    # Determine project directory (where pyproject.toml is, or the file's dir)
+    file_dir = os.path.dirname(os.path.abspath(file))
+    project_dir = file_dir
+
+    # Walk up to find pyproject.toml
+    check = file_dir
+    while check != os.path.dirname(check):
+        if os.path.exists(os.path.join(check, "pyproject.toml")):
+            project_dir = check
+            break
+        check = os.path.dirname(check)
+
+    source_filename = os.path.relpath(os.path.abspath(file), project_dir)
+
+    click.echo(f"Packaging {project_dir}...")
+    archive = _make_archive(project_dir)
+    click.echo(f"Archive size: {len(archive)} bytes")
+
+    for name, info in agents.items():
+        click.echo(f"Deploying {name}...")
+        try:
+            r = httpx.post(
+                f"{url}/api/agents",
+                data={
+                    "name": name,
+                    "source_filename": source_filename,
+                    "function_name": info["function_name"],
+                    "config_json": json.dumps(info["config"]),
+                },
+                files={"archive": ("archive.tar.gz", archive, "application/gzip")},
+                headers=headers,
+                timeout=120,
+            )
+            r.raise_for_status()
+            data = r.json()
+            click.echo(f"  {name}: {data['status']} (id: {data['id']})")
+        except httpx.HTTPStatusError as e:
+            click.echo(f"  {name}: FAILED — {e.response.status_code} {e.response.text}", err=True)
+        except httpx.HTTPError as e:
+            click.echo(f"  {name}: FAILED — {e}", err=True)
+
+    click.echo("Done.")

nightshift/cli/commands/login.py

@@ -0,0 +1,27 @@
+"""nightshift login — store platform credentials."""
+
+from __future__ import annotations
+
+import click
+import httpx
+
+from nightshift.cli.config import save_config
+
+
+@click.command()
+@click.option("--url", required=True, help="Platform server URL (e.g. http://localhost:3000)")
+@click.option("--api-key", required=True, help="API key (ns_...)")
+def login(url: str, api_key: str) -> None:
+    """Authenticate with a Nightshift platform server."""
+    # Strip trailing slash
+    url = url.rstrip("/")
+
+    # Verify the connection
+    try:
+        r = httpx.get(f"{url}/health", timeout=10)
+        r.raise_for_status()
+    except httpx.HTTPError as e:
+        raise click.ClickException(f"Cannot reach server at {url}: {e}")
+
+    save_config({"url": url, "api_key": api_key})
+    click.echo(f"Logged in to {url}")

nightshift/cli/commands/logs.py

@@ -0,0 +1,64 @@
+"""nightshift logs — stream events from a previous run."""
+
+from __future__ import annotations
+
+import json
+import sys
+
+import click
+import httpx
+from httpx_sse import connect_sse
+
+from nightshift.cli.config import get_auth_headers, get_url
+from nightshift.events import TERMINAL_EVENTS
+
+
+@click.command()
+@click.argument("run_id")
+def logs(run_id: str) -> None:
+    """Stream events from a run.
+
+    RUN_ID is the UUID returned by 'nightshift run'.
+    """
+    url = get_url()
+    headers = get_auth_headers()
+
+    try:
+        with httpx.Client(timeout=None, headers=headers) as client:
+            with connect_sse(client, "GET", f"{url}/api/runs/{run_id}/events") as sse:
+                for event in sse.iter_sse():
+                    if not event.data:
+                        continue
+                    try:
+                        data = json.loads(event.data)
+                    except json.JSONDecodeError:
+                        continue
+
+                    event_type = data.get("type", event.event)
+                    _print_event(event_type, data)
+
+                    if event_type in TERMINAL_EVENTS:
+                        return
+    except httpx.HTTPStatusError as e:
+        raise click.ClickException(f"{e.response.status_code}: {e.response.text}")
+    except httpx.HTTPError as e:
+        raise click.ClickException(f"Event stream error: {e}")
+
+
+def _print_event(event_type: str, data: dict) -> None:
+    """Pretty-print an SSE event to stdout."""
+    if event_type == "nightshift.started":
+        click.echo(f"[started] workspace={data.get('workspace', '')}")
+    elif event_type == "nightshift.completed":
+        click.echo("[completed]")
+    elif event_type == "nightshift.error":
+        click.echo(f"[error] {data.get('error', '')}", err=True)
+    elif event_type == "nightshift.interrupted":
+        click.echo(f"[interrupted] reason={data.get('reason', '')}")
+    else:
+        text = data.get("text", data.get("content", ""))
+        if text:
+            click.echo(text, nl=False)
+            sys.stdout.flush()
+        else:
+            click.echo(json.dumps(data, indent=2))

nightshift/cli/commands/run.py

@@ -0,0 +1,94 @@
+"""nightshift run — start an agent run and optionally follow events."""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+
+import click
+import httpx
+from httpx_sse import connect_sse
+
+from nightshift.cli.config import get_auth_headers, get_url
+from nightshift.events import TERMINAL_EVENTS
+
+
+@click.command()
+@click.argument("agent_name")
+@click.option("--prompt", "-p", required=True, help="Prompt to send to the agent")
+@click.option("--follow", "-f", is_flag=True, help="Follow the event stream (like tail -f)")
+def run(agent_name: str, prompt: str, follow: bool) -> None:
+    """Start a run for an agent on the platform."""
+    url = get_url()
+    headers = get_auth_headers()
+
+    # Forward ANTHROPIC_API_KEY from the local environment into the VM
+    runtime_env: dict[str, str] = {}
+    api_key = os.environ.get("ANTHROPIC_API_KEY")
+    if api_key:
+        runtime_env["ANTHROPIC_API_KEY"] = api_key
+    else:
+        click.echo("Warning: ANTHROPIC_API_KEY not set in local environment", err=True)
+
+    try:
+        r = httpx.post(
+            f"{url}/api/agents/{agent_name}/runs",
+            json={"prompt": prompt, "env": runtime_env},
+            headers=headers,
+            timeout=30,
+        )
+        r.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        raise click.ClickException(f"{e.response.status_code}: {e.response.text}")
+    except httpx.HTTPError as e:
+        raise click.ClickException(str(e))
+
+    data = r.json()
+    run_id = data["id"]
+    click.echo(f"Run started: {run_id}")
+
+    if follow:
+        _follow_events(url, headers, run_id)
+
+
+def _follow_events(url: str, headers: dict, run_id: str) -> None:
+    """Stream SSE events for a run to stdout."""
+    try:
+        with httpx.Client(timeout=None, headers=headers) as client:
+            with connect_sse(client, "GET", f"{url}/api/runs/{run_id}/events") as sse:
+                for event in sse.iter_sse():
+                    if not event.data:
+                        continue
+                    try:
+                        data = json.loads(event.data)
+                    except json.JSONDecodeError:
+                        continue
+
+                    event_type = data.get("type", event.event)
+                    _print_event(event_type, data)
+
+                    if event_type in TERMINAL_EVENTS:
+                        return
+    except httpx.HTTPError as e:
+        raise click.ClickException(f"Event stream error: {e}")
+
+
+def _print_event(event_type: str, data: dict) -> None:
+    """Pretty-print an SSE event to stdout."""
+    if event_type == "nightshift.started":
+        click.echo(f"[started] workspace={data.get('workspace', '')}")
+    elif event_type == "nightshift.completed":
+        click.echo("[completed]")
+    elif event_type == "nightshift.error":
+        click.echo(f"[error] {data.get('error', '')}", err=True)
+    elif event_type == "nightshift.interrupted":
+        click.echo(f"[interrupted] reason={data.get('reason', '')}")
+    else:
+        # Agent message — print the text content if available
+        text = data.get("text", data.get("content", ""))
+        if text:
+            click.echo(text, nl=False)
+            sys.stdout.flush()
+        else:
+            click.echo(json.dumps(data, indent=2))