amina-cli 0.3.0__tar.gz → 0.4.1__tar.gz

{amina_cli-0.3.0 → amina_cli-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amina-cli
-Version: 0.3.0
+Version: 0.4.1
 Summary: CLI for AminoAnalytica protein engineering platform
 Project-URL: Homepage, https://aminoanalytica.com
 Project-URL: Documentation, https://docs.aminoanalytica.com

{amina_cli-0.3.0 → amina_cli-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "amina-cli"
-version = "0.3.0"
+version = "0.4.1"
 description = "CLI for AminoAnalytica protein engineering platform"
 readme = "README.md"
 license = {text = "Apache-2.0"}

{amina_cli-0.3.0 → amina_cli-0.4.1}/src/amina_cli/__init__.py

@@ -9,4 +9,4 @@ Quick start:
     amina run esmfold --sequence "MKFLILLFNILCLFPVLAADNH"
 """
 
-__version__ = "0.3.0"
+__version__ = "0.4.1"

amina_cli-0.4.1/src/amina_cli/auth.py

@@ -0,0 +1,558 @@
+"""
+Local credential storage for Amina CLI.
+
+This module handles storing and retrieving API keys locally.
+Validation happens server-side when the key is used.
+
+Credentials are stored in ~/.amina/credentials.json
+
+This is the same pattern used by:
+- OpenAI CLI: stores key in ~/.openai
+- AWS CLI: stores in ~/.aws/credentials
+- Stripe CLI: stores in ~/.config/stripe
+"""
+
+import contextlib
+import fcntl
+import json
+import os
+from pathlib import Path
+from typing import Iterator, Optional
+
+# Configuration directory and files
+CONFIG_DIR = Path.home() / ".amina"
+CREDENTIALS_FILE = CONFIG_DIR / "credentials.json"
+
+# API key format validation
+API_KEY_PREFIX = "ami_"
+API_KEY_MIN_LENGTH = 36  # ami_ + 32 hex chars
+
+
+class AuthError(Exception):
+    """Raised when authentication fails or credentials are missing."""
+
+    pass
+
+
+def validate_key_format(key: str) -> bool:
+    """
+    Validate that a key has the correct format.
+
+    Only checks format, NOT validity. Server validates when key is used.
+
+    Args:
+        key: The API key to validate
+
+    Returns:
+        True if format is valid
+    """
+    if not key:
+        return False
+    if not key.startswith(API_KEY_PREFIX):
+        return False
+    if len(key) < API_KEY_MIN_LENGTH:
+        return False
+    return True
+
+
+def set_api_key(key: str) -> None:
+    """
+    Store an API key locally.
+
+    Format is validated but the key is NOT verified against the server.
+    Verification happens when the key is used for the first time.
+
+    Args:
+        key: The API key to store
+
+    Raises:
+        ValueError: If key format is invalid
+    """
+    if not validate_key_format(key):
+        raise ValueError(
+            f"Invalid API key format. Must start with '{API_KEY_PREFIX}' "
+            f"and be at least {API_KEY_MIN_LENGTH} characters."
+        )
+
+    save_credentials({"api_key": key})
+
+
+def save_credentials(data: dict) -> None:
+    """
+    Save credentials to the config file.
+
+    Follows the same security pattern as OpenAI CLI:
+    - Directory: 0o700 (owner only)
+    - File: 0o600 (owner read/write only)
+
+    Args:
+        data: Dictionary of credentials to save
+    """
+    import os
+
+    # Create config directory with restrictive permissions (owner only)
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True, mode=0o700)
+
+    # Write credentials file atomically with secure permissions
+    # This prevents a race condition where the file briefly exists with wrong permissions
+    content = json.dumps(data, indent=2)
+    flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC
+    with os.fdopen(os.open(CREDENTIALS_FILE, flags, 0o600), "w") as f:
+        f.write(content)
+
+
+def get_credentials() -> dict:
+    """
+    Load stored credentials.
+
+    Returns:
+        Dictionary with 'api_key' and any other stored data
+
+    Raises:
+        AuthError: If not authenticated
+    """
+    if not CREDENTIALS_FILE.exists():
+        raise AuthError(
+            "Not authenticated. Run: amina auth set-key <your_api_key>\n"
+            "Get an API key at: https://app.aminoanalytica.com/settings/api"
+        )
+
+    try:
+        data = json.loads(CREDENTIALS_FILE.read_text())
+        if not data.get("api_key"):
+            raise AuthError("Credentials file exists but no API key found.")
+        return data
+    except json.JSONDecodeError:
+        raise AuthError("Credentials file is corrupted. Run: amina auth set-key <key>")
+
+
+def get_api_key() -> str:
+    """
+    Get the stored API key.
+
+    Convenience wrapper around get_credentials().
+
+    Returns:
+        The stored API key
+
+    Raises:
+        AuthError: If not authenticated
+    """
+    return get_credentials()["api_key"]
+
+
+def clear_credentials() -> None:
+    """
+    Remove all stored credentials (logout).
+
+    Idempotent - safe to call even if not authenticated.
+    """
+    if CREDENTIALS_FILE.exists():
+        CREDENTIALS_FILE.unlink()
+
+
+def get_key_display(key: str) -> str:
+    """
+    Get a safe display version of an API key.
+
+    Shows prefix and first few chars, hides the rest.
+
+    Args:
+        key: The full API key
+
+    Returns:
+        Safe display string like "ami_7f3k..."
+    """
+    if len(key) > 12:
+        return f"{key[:12]}..."
+    return key
+
+
+def is_authenticated() -> bool:
+    """
+    Check if credentials are stored.
+
+    Returns:
+        True if credentials file exists with an API key
+    """
+    try:
+        get_credentials()
+        return True
+    except AuthError:
+        return False
+
+
+def get_session_id() -> Optional[str]:
+    """
+    Get the linked session/conversation ID if set.
+
+    When a session ID is set, CLI outputs will be linked to that
+    web conversation for visibility in the dashboard.
+
+    Returns:
+        Session ID or None if not linked
+    """
+    try:
+        data = get_credentials()
+        return data.get("session_id")
+    except AuthError:
+        return None
+
+
+def set_session_id(session_id: Optional[str]) -> None:
+    """
+    Link CLI outputs to a web conversation.
+
+    Args:
+        session_id: The conversation ID to link, or None to unlink
+    """
+    try:
+        data = get_credentials()
+    except AuthError:
+        raise AuthError("Must be authenticated before linking a session.")
+
+    if session_id:
+        data["session_id"] = session_id
+    else:
+        data.pop("session_id", None)
+
+    save_credentials(data)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# JOB HISTORY STORAGE
+# ═══════════════════════════════════════════════════════════════════════════════
+
+JOBS_FILE = CONFIG_DIR / "jobs.json"
+# Lock path is derived from JOBS_FILE inside _with_jobs_lock() so tests that
+# patch JOBS_FILE don't have to patch a second constant.
+MAX_JOB_HISTORY = 10_000  # Keep last N terminal jobs; non-terminal entries are
+# never evicted regardless of cap (see save_job).
+
+# Statuses that must NEVER be evicted by trim. Losing a queued/running job from
+# the local cache means losing the only client-side handle to wait on it.
+NON_TERMINAL_STATUSES = frozenset({"submitted", "queued", "running", "starting", "pending", "completing"})
+
+
+@contextlib.contextmanager
+def _with_jobs_lock() -> Iterator[None]:
+    """Hold an exclusive advisory lock on jobs.json mutations.
+
+    Without this, concurrent CLI processes (parallel `--background` submits,
+    parallel `jobs list` refreshes) race read-modify-write on jobs.json and
+    silently lose entries. The lock lives on a sibling `.lock` file so reads
+    can still proceed without blocking.
+
+    Both the parent dir and lock path are derived from JOBS_FILE on each call
+    so that tests patching JOBS_FILE (and prod paths via env in future) do not
+    have to also patch a separate LOCK_FILE constant.
+    """
+    JOBS_FILE.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+    lock_path = JOBS_FILE.with_name("jobs.lock")
+    fd = os.open(lock_path, os.O_WRONLY | os.O_CREAT, 0o600)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(fd, fcntl.LOCK_UN)
+    finally:
+        os.close(fd)
+
+
+def _atomic_write_jobs(content: str) -> None:
+    """Write jobs.json via tmp + rename so concurrent readers never see a
+    partially-truncated file (the prior O_TRUNC + write was non-atomic)."""
+    tmp_path = JOBS_FILE.with_suffix(JOBS_FILE.suffix + ".tmp")
+    fd = os.open(tmp_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w") as f:
+        f.write(content)
+    os.replace(tmp_path, JOBS_FILE)
+
+
+def _trim_with_non_terminal_protection(jobs: list[dict]) -> list[dict]:
+    """Trim the cache to MAX_JOB_HISTORY entries, never evicting non-terminal
+    jobs. The cap is a soft cap for terminal entries only; an active job that
+    happens to be old must still be reachable for `jobs wait` / `jobs status`.
+    """
+    if len(jobs) <= MAX_JOB_HISTORY:
+        return jobs
+    head = jobs[:MAX_JOB_HISTORY]
+    overflow = jobs[MAX_JOB_HISTORY:]
+    rescued = [j for j in overflow if j.get("status") in NON_TERMINAL_STATUSES]
+    return head + rescued
+
+
+def save_job(job_info: dict) -> None:
+    """
+    Save a job to the local history.
+
+    Used by background job submission to track jobs for later status checks.
+
+    Args:
+        job_info: Dict with job_id, call_id, user_id, tool_name, etc.
+    """
+    from datetime import datetime
+
+    with _with_jobs_lock():
+        jobs = []
+        if JOBS_FILE.exists():
+            try:
+                jobs = json.loads(JOBS_FILE.read_text())
+            except (json.JSONDecodeError, IOError):
+                jobs = []
+
+        job_info["submitted_at"] = datetime.now().isoformat()
+        # Honour a caller-supplied status (e.g. "queued") instead of always
+        # stamping "submitted" — otherwise a queued job is silently mis-labelled
+        # in the local cache until the next server refresh.
+        job_info.setdefault("status", "submitted")
+
+        # Most-recent first.
+        jobs.insert(0, job_info)
+        jobs = _trim_with_non_terminal_protection(jobs)
+
+        _atomic_write_jobs(json.dumps(jobs, indent=2))
+
+
+def get_job_history(limit: int = 20) -> list[dict]:
+    """
+    Get recent jobs from local history.
+
+    Args:
+        limit: Maximum number of jobs to return
+
+    Returns:
+        List of job info dicts, most recent first
+    """
+    if not JOBS_FILE.exists():
+        return []
+
+    try:
+        jobs = json.loads(JOBS_FILE.read_text())
+        return jobs[:limit]
+    except (json.JSONDecodeError, IOError):
+        return []
+
+
+def get_job_info(job_id: str) -> Optional[dict]:
+    """
+    Look up a job by ID (full or partial match).
+
+    Args:
+        job_id: Full job ID or prefix to match
+
+    Returns:
+        Job info dict or None if not found
+    """
+    jobs = get_job_history(limit=MAX_JOB_HISTORY)
+
+    # Try exact match first
+    for job in jobs:
+        if job.get("job_id") == job_id:
+            return job
+
+    # Try prefix match
+    for job in jobs:
+        if job.get("job_id", "").startswith(job_id):
+            return job
+
+    return None
+
+
+def update_job_status(job_id: str, status: str) -> None:
+    """
+    Update the status of a job in local history.
+
+    Args:
+        job_id: Job ID to update
+        status: New status (e.g., "completed", "failed")
+    """
+    with _with_jobs_lock():
+        if not JOBS_FILE.exists():
+            return
+
+        try:
+            jobs = json.loads(JOBS_FILE.read_text())
+        except (json.JSONDecodeError, IOError):
+            return
+
+        for job in jobs:
+            if job.get("job_id") == job_id:
+                job["status"] = status
+                break
+
+        _atomic_write_jobs(json.dumps(jobs, indent=2))
+
+
+def update_job_info(job_id: str, updates: dict) -> None:
+    """
+    Merge fields into a saved job entry.
+
+    Used to persist a resolved call_id so subsequent polls skip the queue check.
+
+    Args:
+        job_id: Job ID to update
+        updates: Dict of fields to merge (e.g., {"call_id": "fc-xxx"})
+    """
+    with _with_jobs_lock():
+        if not JOBS_FILE.exists():
+            return
+
+        try:
+            jobs = json.loads(JOBS_FILE.read_text())
+        except (json.JSONDecodeError, IOError):
+            return
+
+        for job in jobs:
+            if job.get("job_id") == job_id:
+                job.update(updates)
+                break
+
+        _atomic_write_jobs(json.dumps(jobs, indent=2))
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# SERVER-SIDE REFRESH (writeback path for `amina jobs list`)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+# Fields copied from a server job_tracker row into the local cache.
+# Local-only bookkeeping (reserved_cost, compute_type, submitted_at) is preserved.
+_SERVER_REFRESH_FIELDS = (
+    "status",
+    "message",
+    "runtime_seconds",
+    "cost_usd",
+    "completed_at",
+    "error_message",
+    "results",
+)
+
+
+def _apply_server_fields_to(target: dict, server_row: dict, *, only_if_missing_call_id: bool) -> None:
+    """Apply server-row fields to a local cache entry.
+
+    Single helper for both the insert and update paths so they don't drift.
+    The only behavioural difference between the two callsites is whether
+    ``call_id`` should be overwritten (insert: always; update: only fill if
+    the local entry hasn't learned one yet).
+    """
+    metadata = server_row.get("metadata") or {}
+    server_call_id = metadata.get("call_id")
+    if server_call_id and (not only_if_missing_call_id or not target.get("call_id")):
+        target["call_id"] = server_call_id
+
+    request_params = server_row.get("request_params") or {}
+    if isinstance(request_params, dict):
+        if request_params.get("job_name"):
+            target["job_name"] = request_params["job_name"]
+        if request_params:
+            target["request_params"] = request_params
+
+    for field in _SERVER_REFRESH_FIELDS:
+        if server_row.get(field) is None:
+            continue
+        # Cancelled is sticky on refresh too. If the gateway hadn't yet
+        # applied its own cancel-stickiness guard before this row was read,
+        # don't let a stale server "failed" overwrite the local "cancelled".
+        if field == "status" and target.get("status") == "cancelled" and server_row[field] == "failed":
+            continue
+        target[field] = server_row[field]
+
+
+def _server_row_to_local_entry(server_row: dict) -> dict:
+    """Project a job_tracker row into the local-cache shape."""
+    from datetime import datetime
+
+    entry: dict = {
+        "job_id": server_row.get("job_id", ""),
+        "user_id": server_row.get("user_id", ""),
+        "tool_name": server_row.get("tool_name", ""),
+        "compute_type": server_row.get("gpu_type") or "cpu",
+        "submitted_at": server_row.get("created_at") or datetime.now().isoformat(),
+        "status": server_row.get("status", "submitted"),
+    }
+    _apply_server_fields_to(entry, server_row, only_if_missing_call_id=False)
+    return entry
+
+
+def _meta_path() -> Path:
+    """Sibling file for jobs.json metadata (last_refresh_at, etc.).
+
+    Kept out of jobs.json itself so the file shape stays a pure list and any
+    downstream `jq` consumer keeps working unchanged.
+    """
+    return JOBS_FILE.with_name("jobs.meta.json")
+
+
+def _write_meta(meta: dict) -> None:
+    """Atomic write for the meta sidecar."""
+    p = _meta_path()
+    tmp = p.with_suffix(p.suffix + ".tmp")
+    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w") as f:
+        f.write(json.dumps(meta, indent=2, default=str))
+    os.replace(tmp, p)
+
+
+def read_jobs_meta() -> dict:
+    """Read the jobs metadata sidecar (last_refresh_at, total_unfiltered_count, ...).
+
+    Returns ``{}`` if the file is missing or unreadable.
+    """
+    p = _meta_path()
+    if not p.exists():
+        return {}
+    try:
+        return json.loads(p.read_text())
+    except (json.JSONDecodeError, IOError):
+        return {}
+
+
+def merge_server_jobs(server_jobs: list[dict], total_unfiltered_count: Optional[int] = None) -> None:
+    """
+    Merge authoritative server rows into the local job cache.
+
+    For each server row: update the matching local entry's status/metadata
+    fields in place, or insert a new local entry if the job is unknown locally
+    (e.g., evicted by the FIFO, or submitted from another machine). After the
+    merge, writes ``last_refresh_at`` (and optional ``total_unfiltered_count``)
+    to the meta sidecar so future telemetry / TTL throttling can act on it.
+
+    This is the writeback path for ``amina jobs list`` auto-refresh.
+    """
+    from datetime import datetime
+
+    with _with_jobs_lock():
+        existing: list[dict] = []
+        if JOBS_FILE.exists():
+            try:
+                existing = json.loads(JOBS_FILE.read_text())
+            except (json.JSONDecodeError, IOError):
+                existing = []
+
+        by_job_id = {job.get("job_id"): job for job in existing if job.get("job_id")}
+
+        for row in server_jobs:
+            job_id = row.get("job_id")
+            if not job_id:
+                continue
+            local = by_job_id.get(job_id)
+            if local is None:
+                new_entry = _server_row_to_local_entry(row)
+                existing.append(new_entry)
+                by_job_id[job_id] = new_entry
+                continue
+            _apply_server_fields_to(local, row, only_if_missing_call_id=True)
+
+        # Re-sort newest-first by submitted_at so freshly-merged server rows and
+        # any pre-existing local entries interleave correctly. Without this the
+        # display order drifts to whatever insertion order produced.
+        existing.sort(key=lambda e: e.get("submitted_at") or "", reverse=True)
+
+        _atomic_write_jobs(json.dumps(existing, indent=2, default=str))
+
+        meta = {"last_refresh_at": datetime.now().isoformat()}
+        if total_unfiltered_count is not None:
+            meta["total_unfiltered_count"] = total_unfiltered_count
+        _write_meta(meta)