snowglobe 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

snowglobe/client/src/app.py

@@ -20,8 +20,8 @@ from apscheduler import AsyncScheduler
 from apscheduler.triggers.interval import IntervalTrigger
 from fastapi import FastAPI, HTTPException, Request
 
-from .cli_utils import info
-from .config import config
+from .cli_utils import info, shutdown_manager
+from .config import config, get_api_key_or_raise
 from .models import CompletionFunctionOutputs, CompletionRequest, RiskEvaluationRequest
 from .stats import initialize_stats, track_batch_completion
 from .utils import fetch_experiments, fetch_messages
@@ -190,7 +190,7 @@ async def process_application_heartbeat(app_id):
         connection_test_response = await client.post(
             connection_test_url,
             json=connection_test_payload,
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
 
     if not connection_test_response.is_success:
@@ -240,7 +240,7 @@ async def process_risk_evaluation(test, risk_name):
                 "risk_type": risk_name,
                 "risk_triggered": risk_evaluation.triggered,
             },
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
     LOGGER.debug(f"Time taken: {time.time() - start} seconds")
     if not risk_evaluation_response.is_success:
@@ -271,7 +271,7 @@ async def process_test(test, completion_fn, app_id):
                 "response": completionOutput.response,
                 "persona": test["persona"],
             },
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
     LOGGER.debug(f"Time taken: {time.time() - start} seconds")
     # remove test['id'] from queued_tests
@@ -282,186 +282,382 @@ async def process_test(test, completion_fn, app_id):
 
 # The task to run
 async def poll_for_completions():
+    # Exit early if shutdown requested
+    if await shutdown_manager.is_async_shutdown_requested():
+        return
+
+    job_id = "poll_for_completions"
+
     if len(apps) == 0:
         LOGGER.warning("No applications found. Skipping completions check.")
         return
-    experiments = await fetch_experiments()
 
-    async with httpx.AsyncClient() as client:
-        LOGGER.info(f"Polling {len(experiments)} experiments for completions...")
-        for experiment in experiments:
-            LOGGER.debug(
-                f"Checking experiment: id={experiment.get('id')}, name={experiment.get('name', 'unknown')}"
-            )
-        for experiment in experiments:
-            app_id = experiment.get("app_id")
-            if app_id not in apps:
+    # Register this job as active
+    shutdown_manager.register_active_job(job_id)
+
+    try:
+        experiments = await fetch_experiments()
+    except Exception as e:
+        # Handle shutdown-related connection errors gracefully
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Connection error during shutdown (expected): {e}")
+            return
+        else:
+            LOGGER.error(f"Error fetching experiments: {e}")
+            raise
+    finally:
+        shutdown_manager.unregister_active_job(job_id)
+
+    try:
+        async with httpx.AsyncClient() as client:
+            LOGGER.info(f"Polling {len(experiments)} experiments for completions...")
+            for experiment in experiments:
                 LOGGER.debug(
-                    f"Skipping experiment as we do not have a completion function for app_id: {app_id}"
-                )
-                continue
-            experiment_id = experiment["id"]
-            experiment_request = await client.get(
-                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}?appId={app_id}",
-                headers={"x-api-key": config.API_KEY},
-            )
-            if not experiment_request.is_success:
-                LOGGER.error(
-                    f"Error fetching experiment {experiment_id}: {experiment_request.text}"
+                    f"Checking experiment: id={experiment.get('id')}, name={experiment.get('name', 'unknown')}"
                 )
-                continue
-            experiment = experiment_request.json()
 
-            limit = 10
-            LOGGER.debug(f"Checking for tests for experiment {experiment_id}")
-            tests_response = await client.get(
-                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}/tests?appId={app_id}&include-risk-evaluations=false&limit={limit}&unprocessed-only=true",
-                headers={"x-api-key": config.API_KEY},
-            )
+            for experiment in experiments:
+                # Check for shutdown between each experiment
+                if await shutdown_manager.is_async_shutdown_requested():
+                    LOGGER.debug("Shutdown requested, stopping completions polling")
+                    return
 
-            if not tests_response.is_success:
-                LOGGER.error(f"Error fetching tests: {tests_response.text}")
-                continue
+                app_id = experiment.get("app_id")
+                if app_id not in apps:
+                    LOGGER.debug(
+                        f"Skipping experiment as we do not have a completion function for app_id: {app_id}"
+                    )
+                    continue
+                experiment_id = experiment["id"]
 
-            tests = tests_response.json()
-            completion_count = 0
-            for test in tests:
-                LOGGER.debug(f"Found test {test['id']} for experiment {experiment_id}")
-                if not test["response"] and test["id"] not in queued_tests:
-                    completion_request = await httpx.AsyncClient().post(
-                        f"{config.SNOWGLOBE_CLIENT_URL}/completion",
-                        json={"test": test, "app_id": app_id},
-                        timeout=30,
+                try:
+                    experiment_request = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}?appId={app_id}",
+                        headers={"x-api-key": get_api_key_or_raise()},
                     )
-                    # if 429 raise and exception and stop this batch
-                    if (
-                        not completion_request.is_success
-                        and completion_request.status_code == 429
-                    ):
-                        LOGGER.warning(
-                            f"Completion Rate limit exceeded for test {test['id']}: {completion_request.text}"
-                        )
-                        raise ValueError(
-                            status_code=429,
-                            detail=f"Rate limit exceeded for test {test['id']}",
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(
+                            f"Connection error fetching experiment {experiment_id}: {e}"
                         )
-                    completion_count += 1
-                    queued_tests[test["id"]] = True
-            if completion_count > 0:
-                experiment_name = experiment.get("name", "unknown")
-                if LOGGER.level <= logging.INFO:  # Verbose mode
-                    LOGGER.info(
-                        f"Processed {completion_count} completions for experiment {experiment_name} ({experiment_id})"
+                        continue
+
+                if not experiment_request.is_success:
+                    LOGGER.error(
+                        f"Error fetching experiment {experiment_id}: {experiment_request.text}"
                     )
-                else:  # Clean UI mode
-                    timestamp = datetime.datetime.now().strftime("%H:%M:%S")
-                    info(
-                        f"[{timestamp}] ✓ Batch complete: {completion_count} responses sent ({experiment_name})"
+                    continue
+                experiment = experiment_request.json()
+
+                limit = 10
+                LOGGER.debug(f"Checking for tests for experiment {experiment_id}")
+
+                try:
+                    tests_response = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment_id}/tests?appId={app_id}&include-risk-evaluations=false&limit={limit}&unprocessed-only=true",
+                        headers={"x-api-key": get_api_key_or_raise()},
                     )
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(f"Connection error fetching tests: {e}")
+                        continue
+
+                if not tests_response.is_success:
+                    LOGGER.error(f"Error fetching tests: {tests_response.text}")
+                    continue
 
-                    # Track batch completion
-                    track_batch_completion(experiment_name, completion_count)
+                tests = tests_response.json()
+                completion_count = 0
+                for test in tests:
+                    # Check for shutdown before processing each test
+                    if await shutdown_manager.is_async_shutdown_requested():
+                        LOGGER.debug("Shutdown requested, stopping test processing")
+                        return
+
+                    LOGGER.debug(
+                        f"Found test {test['id']} for experiment {experiment_id}"
+                    )
+                    if not test["response"] and test["id"] not in queued_tests:
+                        try:
+                            completion_request = await httpx.AsyncClient().post(
+                                f"{config.SNOWGLOBE_CLIENT_URL}/completion",
+                                json={"test": test, "app_id": app_id},
+                                timeout=30,
+                            )
+                        except (httpx.ConnectError, httpx.TimeoutException) as e:
+                            if shutdown_manager.is_shutdown_requested():
+                                LOGGER.debug(
+                                    f"HTTP error during shutdown (expected): {e}"
+                                )
+                                return
+                            else:
+                                LOGGER.error(
+                                    f"Connection error posting completion: {e}"
+                                )
+                                continue
+
+                        # if 429 raise and exception and stop this batch
+                        if (
+                            not completion_request.is_success
+                            and completion_request.status_code == 429
+                        ):
+                            LOGGER.warning(
+                                f"Completion Rate limit exceeded for test {test['id']}: {completion_request.text}"
+                            )
+                            raise ValueError(
+                                status_code=429,
+                                detail=f"Rate limit exceeded for test {test['id']}",
+                            )
+                        completion_count += 1
+                        queued_tests[test["id"]] = True
+
+                if completion_count > 0:
+                    experiment_name = experiment.get("name", "unknown")
+                    if LOGGER.level <= logging.INFO:  # Verbose mode
+                        LOGGER.info(
+                            f"Processed {completion_count} completions for experiment {experiment_name} ({experiment_id})"
+                        )
+                    else:  # Clean UI mode
+                        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
+                        info(
+                            f"[{timestamp}] ✓ Batch complete: {completion_count} responses sent ({experiment_name})"
+                        )
+
+                        # Track batch completion
+                        track_batch_completion(experiment_name, completion_count)
+
+    except Exception as e:
+        # Handle any other errors
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Error during shutdown (expected): {e}")
+        else:
+            LOGGER.error(f"Unexpected error in poll_for_completions: {e}")
+            raise
 
 
 async def process_application_heartbeats():
-    connection_test_count = 0
-    LOGGER.info("Processing application heartbeats...")
-    for app_id, app_info in apps.items():
-        connection_test_request = await httpx.AsyncClient().post(
-            f"{config.SNOWGLOBE_CLIENT_URL}/heartbeat",
-            json={"app_id": app_id},
-            timeout=30,
-        )
-        if not connection_test_request.is_success:
-            LOGGER.error(
-                f"Error sending heartbeat for application {app_id}: {connection_test_request.text}"
-            )
-            continue
-        connection_test_count += 1
+    # Exit early if shutdown requested
+    if await shutdown_manager.is_async_shutdown_requested():
+        return
 
-    LOGGER.info(f"Processed {connection_test_count} heartbeats for applications.")
+    job_id = "process_application_heartbeats"
+    shutdown_manager.register_active_job(job_id)
 
+    try:
+        connection_test_count = 0
+        LOGGER.info("Processing application heartbeats...")
+        for app_id, app_info in apps.items():
+            # Check for shutdown between each heartbeat
+            if await shutdown_manager.is_async_shutdown_requested():
+                LOGGER.debug("Shutdown requested, stopping heartbeats")
+                return
 
-async def poll_for_risk_evaluations():
-    """Poll for risk evaluations and process them."""
-    experiments = await fetch_experiments()
-    LOGGER.info("Checking for pending risk evaluations...")
-    LOGGER.debug(f"Found {len(experiments)} experiments with validation in progress")
-    # experiments = [{"id": "123"}]
-    async with httpx.AsyncClient() as client:
-        for experiment in experiments:
-            experiment_request = await client.get(
-                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}",
-                headers={"x-api-key": config.API_KEY},
-            )
-            if not experiment_request.is_success:
+            try:
+                connection_test_request = await httpx.AsyncClient().post(
+                    f"{config.SNOWGLOBE_CLIENT_URL}/heartbeat",
+                    json={"app_id": app_id},
+                    timeout=30,
+                )
+            except (httpx.ConnectError, httpx.TimeoutException) as e:
+                if shutdown_manager.is_shutdown_requested():
+                    LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                    return
+                else:
+                    LOGGER.error(
+                        f"Connection error sending heartbeat for {app_id}: {e}"
+                    )
+                    continue
+
+            if not connection_test_request.is_success:
                 LOGGER.error(
-                    f"Error fetching experiment {experiment['id']}: {experiment_request.text}"
+                    f"Error sending heartbeat for application {app_id}: {connection_test_request.text}"
                 )
                 continue
-            experiment = experiment_request.json()
-            risk_eval_count = 0
-            for risk_name in risks.keys():
+            connection_test_count += 1
+
+        LOGGER.info(f"Processed {connection_test_count} heartbeats for applications.")
+    finally:
+        shutdown_manager.unregister_active_job(job_id)
+
+
+async def poll_for_risk_evaluations():
+    """Poll for risk evaluations and process them."""
+    # Exit early if shutdown requested
+    if await shutdown_manager.is_async_shutdown_requested():
+        return
+
+    job_id = "poll_for_risk_evaluations"
+    shutdown_manager.register_active_job(job_id)
+
+    try:
+        experiments = await fetch_experiments()
+    except Exception as e:
+        # Handle shutdown-related connection errors gracefully
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Connection error during shutdown (expected): {e}")
+            return
+        else:
+            LOGGER.error(f"Error fetching experiments: {e}")
+            raise
+    finally:
+        shutdown_manager.unregister_active_job(job_id)
+
+    try:
+        LOGGER.info("Checking for pending risk evaluations...")
+        LOGGER.debug(
+            f"Found {len(experiments)} experiments with validation in progress"
+        )
+
+        async with httpx.AsyncClient() as client:
+            for experiment in experiments:
+                # Check for shutdown between experiments
+                if await shutdown_manager.is_async_shutdown_requested():
+                    LOGGER.debug("Shutdown requested, stopping risk evaluations")
+                    return
+
                 try:
-                    if (
-                        risk_name
-                        not in experiment.get("source_data", {})
-                        .get("evaluation_configuration", {})
-                        .keys()
-                    ):
-                        LOGGER.debug(
-                            f"Skipping experiment {experiment['id']} as it does not have risk {risk_name}"
+                    experiment_request = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}",
+                        headers={"x-api-key": get_api_key_or_raise()},
+                    )
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(
+                            f"Connection error fetching experiment {experiment['id']}: {e}"
                         )
                         continue
-                    LOGGER.debug(
-                        f"checking for tests for experiment {experiment['id']}"
-                    )
-                    tests_response = await client.get(
-                        f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}/tests?unevaluated-risk={quote_plus(risk_name)}&include-risk-evaluations=true",
-                        headers={"x-api-key": config.API_KEY},
+
+                if not experiment_request.is_success:
+                    LOGGER.error(
+                        f"Error fetching experiment {experiment['id']}: {experiment_request.text}"
                     )
+                    continue
+                experiment = experiment_request.json()
+                risk_eval_count = 0
 
-                    if not tests_response.is_success:
-                        message = (
-                            tests_response.json().get("message") or tests_response.text
-                        )
-                        raise ValueError(
-                            status_code=tests_response.status_code,
-                            message=message,
-                        )
-                    tests = tests_response.json()
+                for risk_name in risks.keys():
+                    # Check for shutdown before processing each risk
+                    if await shutdown_manager.is_async_shutdown_requested():
+                        LOGGER.debug("Shutdown requested, stopping risk processing")
+                        return
 
-                    for test in tests:
-                        test_id = test["id"]
+                    try:
                         if (
-                            test_id not in queued_evaluations
-                            and test.get("response") is not None
+                            risk_name
+                            not in experiment.get("source_data", {})
+                            .get("evaluation_configuration", {})
+                            .keys()
                         ):
-                            risk_eval_response = await httpx.AsyncClient().post(
-                                f"{config.SNOWGLOBE_CLIENT_URL}/risk-evaluation",
-                                json={"test": test, "risk_name": risk_name},
-                                timeout=30,
+                            LOGGER.debug(
+                                f"Skipping experiment {experiment['id']} as it does not have risk {risk_name}"
                             )
-                            # if risk evaltion response is 429 raise and exception and bail on this batch
-                            if (
-                                not risk_eval_response.is_success
-                                and risk_eval_response.status_code == 429
-                            ):
-                                LOGGER.error(
-                                    f"Rate limit exceeded for risk evaluation {test['id']}: {risk_eval_response.text}"
+                            continue
+                        LOGGER.debug(
+                            f"checking for tests for experiment {experiment['id']}"
+                        )
+
+                        try:
+                            tests_response = await client.get(
+                                f"{config.CONTROL_PLANE_URL}/api/experiments/{experiment['id']}/tests?unevaluated-risk={quote_plus(risk_name)}&include-risk-evaluations=true",
+                                headers={"x-api-key": get_api_key_or_raise()},
+                            )
+                        except (httpx.ConnectError, httpx.TimeoutException) as e:
+                            if shutdown_manager.is_shutdown_requested():
+                                LOGGER.debug(
+                                    f"HTTP error during shutdown (expected): {e}"
                                 )
-                                raise ValueError(
-                                    status_code=429,
-                                    detail=f"Rate limit exceeded for risk evaluation {test['id']}",
+                                return
+                            else:
+                                LOGGER.error(f"Connection error fetching tests: {e}")
+                                continue
+
+                        if not tests_response.is_success:
+                            message = (
+                                tests_response.json().get("message")
+                                or tests_response.text
+                            )
+                            raise ValueError(
+                                status_code=tests_response.status_code,
+                                message=message,
+                            )
+                        tests = tests_response.json()
+
+                        for test in tests:
+                            # Check for shutdown before processing each test
+                            if await shutdown_manager.is_async_shutdown_requested():
+                                LOGGER.debug(
+                                    "Shutdown requested, stopping test evaluation"
                                 )
-                            queued_evaluations[test_id] = True
-                            risk_eval_count += 1
-                except Exception as e:
-                    LOGGER.error(f"Error fetching tests: {e}")
-            if risk_eval_count > 0:
-                LOGGER.info(
-                    f"Processed {risk_eval_count} risk evaluations for experiment {experiment.get('name', 'unknown')} ({experiment['id']})"
-                )
+                                return
+
+                            test_id = test["id"]
+                            if (
+                                test_id not in queued_evaluations
+                                and test.get("response") is not None
+                            ):
+                                try:
+                                    risk_eval_response = await httpx.AsyncClient().post(
+                                        f"{config.SNOWGLOBE_CLIENT_URL}/risk-evaluation",
+                                        json={"test": test, "risk_name": risk_name},
+                                        timeout=30,
+                                    )
+                                except (
+                                    httpx.ConnectError,
+                                    httpx.TimeoutException,
+                                ) as e:
+                                    if shutdown_manager.is_shutdown_requested():
+                                        LOGGER.debug(
+                                            f"HTTP error during shutdown (expected): {e}"
+                                        )
+                                        return
+                                    else:
+                                        LOGGER.error(
+                                            f"Connection error posting risk evaluation: {e}"
+                                        )
+                                        continue
+
+                                # if risk evaltion response is 429 raise and exception and bail on this batch
+                                if (
+                                    not risk_eval_response.is_success
+                                    and risk_eval_response.status_code == 429
+                                ):
+                                    LOGGER.error(
+                                        f"Rate limit exceeded for risk evaluation {test['id']}: {risk_eval_response.text}"
+                                    )
+                                    raise ValueError(
+                                        status_code=429,
+                                        detail=f"Rate limit exceeded for risk evaluation {test['id']}",
+                                    )
+                                queued_evaluations[test_id] = True
+                                risk_eval_count += 1
+                    except Exception as e:
+                        if shutdown_manager.is_shutdown_requested():
+                            LOGGER.debug(f"Error during shutdown (expected): {e}")
+                            return
+                        else:
+                            LOGGER.error(f"Error fetching tests: {e}")
+
+                if risk_eval_count > 0:
+                    LOGGER.info(
+                        f"Processed {risk_eval_count} risk evaluations for experiment {experiment.get('name', 'unknown')} ({experiment['id']})"
+                    )
+
+    except Exception as e:
+        # Handle any other errors
+        if shutdown_manager.is_shutdown_requested():
+            LOGGER.debug(f"Error during shutdown (expected): {e}")
+        else:
+            LOGGER.error(f"Unexpected error in poll_for_risk_evaluations: {e}")
+            raise
 
 
 # Ensure the scheduler shuts down properly on application exit.
@@ -599,7 +795,7 @@ async def lifespan(app: FastAPI):
     for app_id in apps.keys():
         try:
             connection_test_url = (
-            f"{config.CONTROL_PLANE_URL}/api/failed-code-connection-tests"
+                f"{config.CONTROL_PLANE_URL}/api/failed-code-connection-tests"
             )
             connection_test_payload = {
                 "appId": app_id,

snowglobe/client/src/cli.py

@@ -7,6 +7,7 @@ import sys
 import threading
 import time
 import webbrowser
+from importlib.metadata import version
 from typing import Optional, Tuple
 
 import typer
@@ -49,6 +50,7 @@ cli_app = typer.Typer(
     help="❄️  Snowglobe CLI - Connect your applications to Snowglobe experiments",
     add_completion=False,
     rich_markup_mode="rich",
+    no_args_is_help=False,
 )
 
 
@@ -64,7 +66,7 @@ def setup_global_options(
     cli_state.json_output = json_output
 
 
-@cli_app.callback()
+@cli_app.callback(invoke_without_command=True)
 def main(
     ctx: typer.Context,
     verbose: bool = typer.Option(
@@ -72,10 +74,26 @@ def main(
     ),
     quiet: bool = typer.Option(False, "--quiet", "-q", help="Minimize output"),
     json_output: bool = typer.Option(False, "--json", help="Output in JSON format"),
+    version_flag: bool = typer.Option(
+        False, "--version", help="Show version and exit", is_flag=True
+    ),
 ):
     """
     ❄️  Snowglobe CLI - Connect your applications to Snowglobe experiments
     """
+    if version_flag:
+        try:
+            pkg_version = version("snowglobe-connect")
+            typer.echo(f"snowglobe-connect {pkg_version}")
+        except Exception:
+            typer.echo("snowglobe-connect (version unknown)")
+        raise typer.Exit()
+
+    # Show help if no command provided and no version flag
+    if not ctx.invoked_subcommand and not version_flag:
+        typer.echo(ctx.get_help())
+        raise typer.Exit()
+
     setup_global_options(ctx, verbose, quiet, json_output)
 
 
@@ -185,16 +203,16 @@ def test(
 
 @cli_app.command()
 def init(
-    name: Optional[str] = typer.Option(
-        None, "--name", help="Custom filename for the agent wrapper"
-    ),
-    skip_template: bool = typer.Option(
-        False, "--skip-template", help="Skip creating template file"
+    file: Optional[str] = typer.Option(
+        None,
+        "--file",
+        "-f",
+        help="Path or filename (within project) for the agent wrapper",
     ),
     # option for stateful agent
     stateful: bool = typer.Option(
         False, "--stateful", help="Initialize a stateful agent template"
-    )
+    ),
 ):
     """
     Initialize a new Snowglobe agent in the current directory
@@ -255,15 +273,37 @@ def init(
         pm.ensure_project_structure()
 
     # Determine filename
-    if name:
-        # User provided custom name
-        filename = name if name.endswith(".py") else f"{name}.py"
+    if file:
+        # User provided explicit file path or name
+        from pathlib import Path
+
+        provided_path = Path(os.path.expanduser(file))
+
+        # Ensure .py suffix
+        if provided_path.suffix != ".py":
+            provided_path = provided_path.with_suffix(".py")
+
+        # Normalize to a path within the project root
+        if provided_path.is_absolute():
+            try:
+                relative_path = provided_path.relative_to(pm.project_root)
+            except ValueError:
+                error("--file must be within the current project directory")
+                info(f"Project root: {pm.project_root}")
+                raise typer.Exit(1)
+        else:
+            relative_path = provided_path
+
+        # Store mapping as a POSIX-style relative path string
+        filename = relative_path.as_posix()
     else:
         # Generate from app name
         filename = pm.sanitize_filename(app_name)
 
-    # Find available filename
-    filename = pm.find_available_filename(filename)
+    # If filename was auto-generated from app name, find an available one
+    if not file:
+        filename = pm.find_available_filename(filename)
+
     file_path = pm.project_root / filename
 
     success(f"Using filename: {filename}")
@@ -272,34 +312,33 @@ def init(
         snowglobe_connect_template = stateful_snowglobe_connect_template
     else:
         snowglobe_connect_template = stateless_snowglobe_connect_template
+
+    # Ensure parent directories exist if a subpath was provided
+    os.makedirs(file_path.parent, exist_ok=True)
+
     # Create agent wrapper file
-    if not skip_template:
-        with spinner("Creating agent wrapper"):
-            with open(file_path, "w") as f:
-                f.write(snowglobe_connect_template)
-        success(f"Created agent wrapper: {filename}")
+    with spinner("Creating agent wrapper"):
+        with open(file_path, "w") as f:
+            f.write(snowglobe_connect_template)
+    success(f"Created agent wrapper: {filename}")
 
-        # Add to mapping
-        pm.add_agent_mapping(filename, app_id, app_name)
-        success("Added mapping to .snowglobe/agents.json")
+    # Add to mapping
+    pm.add_agent_mapping(filename, app_id, app_name)
+    success("Added mapping to .snowglobe/agents.json")
 
-        console.print("\n[dim]Project structure:[/dim]")
-        console.print("[dim]  .snowglobe/agents.json\t- UUID mappings[/dim]")
-        console.print(f"[dim]  {filename}\t- Your agent wrapper[/dim]")
+    console.print("\n[dim]Project structure:[/dim]")
+    console.print("[dim]  .snowglobe/agents.json\t- UUID mappings[/dim]")
+    console.print(f"[dim]  {filename}\t- Your agent wrapper[/dim]")
 
-        console.print()
-        info("Next steps:")
-        console.print("1. Edit the process_scenario function in your agent file:")
-        console.print(f"   [bold cyan]{filename}[/bold cyan]")
-        console.print("2. Implement your application logic")
-        console.print("3. Test your agent:")
-        console.print("   [bold green]snowglobe-connect test[/bold green]")
-        console.print("4. Start the client:")
-        console.print("   [bold green]snowglobe-connect start[/bold green]")
-    else:
-        info(f"Skipped template creation. You'll need to create {filename} manually")
-        # Still add the mapping even if we skip template
-        pm.add_agent_mapping(filename, app_id, app_name)
+    console.print()
+    info("Next steps:")
+    console.print("1. Edit the process_scenario function in your agent file:")
+    console.print(f"   [bold cyan]{filename}[/bold cyan]")
+    console.print("2. Implement your application logic")
+    console.print("3. Test your agent:")
+    console.print("   [bold green]snowglobe-connect test[/bold green]")
+    console.print("4. Start the client:")
+    console.print("   [bold green]snowglobe-connect start[/bold green]")
 
     console.print()
     success(f"Agent '{app_name}' initialized successfully! 🎉")
@@ -494,6 +533,7 @@ async def completion_fn(request: CompletionRequest) -> CompletionFunctionOutputs
     return CompletionFunctionOutputs(response=response_data.messages[0]["content"])
 """
 
+
 def _save_api_key_to_rc(api_key: str, rc_path: str) -> None:
     """Save API key to config file"""
     # Ensure directory exists
@@ -706,15 +746,48 @@ def start(
         "[bold green]🚀 Agent server is live! Processing scenarios...[/bold green]\n"
     )
 
-    # Handle Ctrl+C gracefully
-    def signal_handler(sig, frame):
-        graceful_shutdown()
-
-    signal.signal(signal.SIGINT, signal_handler)
+    # Handle Ctrl+C gracefully (legacy handler delegates to smart shutdown internally)
+    signal.signal(signal.SIGINT, graceful_shutdown)
 
     if not verbose:
         console.print("[dim]Press Ctrl+C to stop[/dim]\n")
 
+    # Show helpful links to open the application(s) in the Snowglobe UI
+    try:
+        pm = get_project_manager()
+        agents = pm.list_agents()
+
+        if agents:
+            if len(agents) == 1:
+                _, agent_info = agents[0]
+                app_id = agent_info.get("uuid")
+                app_name = agent_info.get("name", "Application")
+                if app_id:
+                    app_url = f"{UI_URL}/applications/{app_id}"
+                    console.print(
+                        "🔗 Agent connected! Go to your Snowglobe agent and kick off a simulation:"
+                    )
+                    console.print(f"   [link={app_url}]{app_name}[/link]\n")
+            else:
+                console.print(
+                    "🔗 Agents connected! Go to your Snowglobe agents and kick off simulations:"
+                )
+                for _, agent_info in agents[:5]:
+                    app_id = agent_info.get("uuid")
+                    app_name = agent_info.get("name", "Application")
+                    if app_id:
+                        app_url = f"{UI_URL}/applications/{app_id}"
+                        console.print(
+                            f"   - {app_name}: [link={app_url}]{app_url}[/link]"
+                        )
+                if len(agents) > 5:
+                    console.print("   (and more...)\n")
+                else:
+                    console.print()
+    except Exception:
+        # Do not block startup if we cannot load agents mapping
+        pass
+
     # Import start_client here to avoid config initialization at module import time
     from .app import start_client
 

snowglobe/client/src/cli_utils.py

@@ -1,7 +1,9 @@
+import asyncio
 import contextlib
 import math
 import os
 import sys
+import threading
 import time
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -25,7 +27,67 @@ class CliState:
         self.json_output = False
 
 
+class ShutdownManager:
+    """Coordinate graceful shutdown between signal handler and background jobs"""
+
+    def __init__(self):
+        self.shutdown_event = threading.Event()  # For sync code
+        self.async_shutdown_event = None  # Created when event loop available
+        self.force_shutdown = False
+        self.shutdown_initiated = False
+        self.active_jobs = set()  # Track active background jobs
+        self.active_jobs_lock = threading.Lock()
+
+    def register_active_job(self, job_id: str):
+        """Register a job as actively running"""
+        with self.active_jobs_lock:
+            self.active_jobs.add(job_id)
+
+    def unregister_active_job(self, job_id: str):
+        """Unregister a job when it completes"""
+        with self.active_jobs_lock:
+            self.active_jobs.discard(job_id)
+
+    def has_active_jobs(self) -> bool:
+        """Check if there are any active jobs running"""
+        with self.active_jobs_lock:
+            return len(self.active_jobs) > 0
+
+    def initiate_shutdown(self) -> bool:
+        """Initiate shutdown. Returns True if this is a force shutdown (second Ctrl+C)"""
+        if self.shutdown_initiated:
+            # Second Ctrl+C - force shutdown
+            self.force_shutdown = True
+            return True
+        else:
+            # First Ctrl+C - start graceful shutdown
+            self.shutdown_initiated = True
+            self.shutdown_event.set()
+
+            # Set async event if event loop exists
+            try:
+                loop = asyncio.get_running_loop()
+                if self.async_shutdown_event is None:
+                    self.async_shutdown_event = asyncio.Event()
+                loop.call_soon_threadsafe(self.async_shutdown_event.set)
+            except RuntimeError:
+                pass
+
+            return False
+
+    def is_shutdown_requested(self) -> bool:
+        """Check if shutdown has been requested"""
+        return self.shutdown_event.is_set()
+
+    async def is_async_shutdown_requested(self) -> bool:
+        """Async version of shutdown check"""
+        if self.async_shutdown_event is None:
+            return self.shutdown_event.is_set()
+        return self.async_shutdown_event.is_set() or self.shutdown_event.is_set()
+
+
 cli_state = CliState()
+shutdown_manager = ShutdownManager()
 
 
 def get_api_key() -> Optional[str]:
@@ -104,13 +166,8 @@ def docs_link(message: str, url: str = "https://www.snowglobe.so/docs") -> None:
         console.print(f"📖 {message}: {url}", style="cyan")
 
 
-def graceful_shutdown():
-    """Handle graceful shutdown with session summary"""
-    console.print("\n")
-    warning("🛑 Shutting down gracefully...")
-    success("Completing current scenarios")
-    success("Connection closed")
-
+def show_session_summary():
+    """Display session summary statistics"""
     stats = get_shutdown_stats()
 
     if stats and stats["total_messages"] > 0:
@@ -130,6 +187,90 @@ def graceful_shutdown():
     else:
         success("No scenarios processed during this session")
 
+
+def begin_graceful_shutdown():
+    """Coordinate graceful shutdown - runs in background thread"""
+    TIMEOUT_SECONDS = 8
+
+    console.print("\n")
+    warning("🛑 Shutting down gracefully...")
+    success("Completing current scenarios")
+
+    # Wait for active jobs to finish (with timeout)
+    start_time = time.time()
+    while time.time() - start_time < TIMEOUT_SECONDS:
+        if not shutdown_manager.has_active_jobs():
+            break
+        time.sleep(0.1)
+
+        # Check if user pressed Ctrl+C again (force quit)
+        if shutdown_manager.force_shutdown:
+            return  # Signal handler will handle force quit
+
+    # Jobs finished or timeout reached
+    elapsed = time.time() - start_time
+    if elapsed >= TIMEOUT_SECONDS:
+        warning(
+            f"Shutdown timeout reached after {TIMEOUT_SECONDS}s, completing shutdown..."
+        )
+
+    success("Connection closed")
+
+    # Show session summary
+    show_session_summary()
+
+    console.print()
+    success("Agent disconnected successfully")
+
+    # Exit cleanly
+    sys.exit(0)
+
+
+def smart_signal_handler(sig, frame):
+    """Smart two-stage signal handler with immediate exit when possible"""
+    force_quit = shutdown_manager.initiate_shutdown()
+
+    if force_quit:
+        # Second Ctrl+C - force quit immediately
+        console.print("\n[bold red]🚨 Force shutdown![/bold red]")
+        sys.exit(1)
+    else:
+        # First Ctrl+C - check if we have active work
+        if shutdown_manager.has_active_jobs():
+            # Have active jobs - start graceful shutdown with user feedback
+            console.print(
+                "\n[bold yellow]🛑 Graceful shutdown initiated...[/bold yellow]"
+            )
+            console.print("[dim]Press Ctrl+C again to force quit[/dim]")
+
+            # Start graceful shutdown in background thread
+            threading.Thread(target=begin_graceful_shutdown, daemon=True).start()
+        else:
+            # No active jobs - shutdown immediately
+            console.print("\n[bold blue]🛑 Shutting down...[/bold blue]")
+
+            # Show session summary and exit
+            show_session_summary()
+            console.print()
+            success("Agent disconnected successfully")
+            sys.exit(0)
+
+
+def graceful_shutdown():
+    """Handle graceful shutdown with session summary (idempotent, always exits 0).
+
+    This legacy entrypoint avoids toggling internal shutdown state so repeated
+    calls behave consistently in tests and scripts.
+    """
+    # If there are active jobs, allow them a moment to finish similar to begin_graceful_shutdown,
+    # but without changing the shutdown state in a way that triggers force-exit on reentry.
+    if shutdown_manager.has_active_jobs():
+        console.print("\n[bold yellow]🛑 Graceful shutdown initiated...[/bold yellow]")
+        console.print("[dim]Press Ctrl+C again to force quit[/dim]")
+
+    # Show session summary and exit cleanly
+    console.print("\n[bold blue]🛑 Shutting down...[/bold blue]")
+    show_session_summary()
     console.print()
     success("Agent disconnected successfully")
     sys.exit(0)
@@ -174,12 +315,17 @@ def select_stateful_interactive(
     if cli_state.json_output:
         # For JSON mode, just return the default stateful value
         return stateful
-    info("Some stateful agents such as ones that maintain communication over a websocket or convo specific completion endpoint require stateful integration.")
-    info("If your agent takes messages and completions on a single completion endpoint regardless of context, you can answer no to the following question.")
+    info(
+        "Some stateful agents such as ones that maintain communication over a websocket or convo specific completion endpoint require stateful integration."
+    )
+    info(
+        "If your agent takes messages and completions on a single completion endpoint regardless of context, you can answer no to the following question."
+    )
     if Confirm.ask("Would you like to create a new application?"):
         return True
     return False
 
+
 def select_application_interactive(
     applications: List[Dict[str, Any]],
 ) -> Optional[Dict[str, Any]]:

snowglobe/client/src/config.py

@@ -42,7 +42,7 @@ def migrate_rc_file_if_needed():
 
 
 class Config:
-    def __init__(self):
+    def __init__(self, require_api_key=True):
         # Migrate legacy .snowgloberc if needed
         migrate_rc_file_if_needed()
 
@@ -52,8 +52,16 @@ class Config:
         self.CONCURRENT_HEARTBEATS_PER_INTERVAL = 120
         self.CONCURRENT_HEARTBEATS_INTERVAL_SECONDS = 60
 
-        # Initialize API key - just raise exception if missing, let caller handle it
-        self.API_KEY = self.get_api_key()
+        # Initialize API key - optionally raise exception if missing
+        if require_api_key:
+            self.API_KEY = self.get_api_key()
+        else:
+            # For auth command, don't require API key during initialization
+            try:
+                self.API_KEY = self.get_api_key()
+            except ValueError:
+                self.API_KEY = None
+
         self.APPLICATION_ID = self.get_application_id()
         self.CONTROL_PLANE_URL = self.get_control_plane_url()
         self.SNOWGLOBE_CLIENT_URL = self.get_snowglobe_client_url()
@@ -210,4 +218,28 @@ class Config:
         return application_id
 
 
-config = Config()
+config = Config(require_api_key=False)
+
+
+def get_api_key_or_raise() -> str:
+    """Get API key, raising ValueError if not found - for commands that require it"""
+    api_key = config.API_KEY
+    if not api_key:
+        # Try to get it fresh in case it was just set
+        try:
+            api_key = config.get_api_key()
+            config.API_KEY = api_key  # Cache it
+        except ValueError:
+            pass
+
+    if not api_key:
+        raise ValueError(
+            "API key is required either passed as an argument or set as an environment variable. \n"
+            "You can set the API key as an environment variable by running: \n"
+            "export SNOWGLOBE_API_KEY=<your_api_key> \n"
+            "or \n"
+            "export GUARDRAILS_API_KEY=<your_api_key> \n"
+            "Or you can create a .snowglobe/config.rc file in the current directory with the line: \n"
+            "SNOWGLOBE_API_KEY=<your_api_key> \n"
+        )
+    return api_key

snowglobe/client/src/utils.py

@@ -2,7 +2,7 @@ from logging import getLogger
 
 import httpx
 
-from .config import config
+from .config import config, get_api_key_or_raise
 from .models import SnowglobeData, SnowglobeMessage
 
 LOGGER = getLogger(__name__)
@@ -21,7 +21,7 @@ async def fetch_experiments(app_id: str = None) -> list[dict]:
             experiments_url += f"&appId={config.APPLICATION_ID}"
         experiments_response = await client.get(
             experiments_url,
-            headers={"x-api-key": config.API_KEY},
+            headers={"x-api-key": get_api_key_or_raise()},
         )
 
     if not experiments_response.status_code == 200:
@@ -78,7 +78,7 @@ async def fetch_messages(*, test) -> list[SnowglobeMessage]:
             # get parent test
             parent_test_response = await client.get(
                 f"{config.CONTROL_PLANE_URL}/api/experiments/{test['experiment_id']}/tests/{parent_id}",
-                headers={"x-api-key": config.API_KEY},
+                headers={"x-api-key": get_api_key_or_raise()},
             )
 
             if not parent_test_response.status_code == 200:

{snowglobe-0.4.0.dist-info → snowglobe-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowglobe
-Version: 0.4.0
+Version: 0.4.2
 Summary: client server for usage with snowglobe experiments
 Author-email: Guardrails AI <contact@guardrailsai.com>
 License: MIT License
@@ -49,21 +49,14 @@ The Snowglobe Connect SDK helps you connect your AI agents to Snowglobe. It send
 
 ## Installation
 
-```bash
-# Paste Guardrails Client or extract from guardrailsrc
-export GUARDRAILS_TOKEN=$(cat ~/.guardrailsrc| awk -F 'token=' '{print $2}' | awk '{print $1}' | tr -d '\n')
-```
-
 ```
 # Install client
-pip install -U --index-url="https://__token__:$GUARDRAILS_TOKEN@pypi.guardrailsai.com/simple" \
-    --extra-index-url="https://pypi.org/simple" snowglobe-connect
+pip install snowglobe-connect
 ```
 
 If using uv, set the `--prerelease=allow` flag
 ```
-pip install -U --index-url="https://__token__:$GUARDRAILS_TOKEN@pypi.guardrailsai.com/simple" \
-    --extra-index-url="https://pypi.org/simple" --prerelease=allow snowglobe-connect
+uv pip install --prerelease=allow snowglobe-connect
 ```
 
 

snowglobe-0.4.2.dist-info/RECORD

@@ -0,0 +1,15 @@
+snowglobe/client/__init__.py,sha256=kzp9wPUUYBXqDSKZbfmD4vrAQvrWSW5HOvtpFlEJWfs,353
+snowglobe/client/src/app.py,sha256=MRV0mzrH4iqkCk1H5pCQZgMj4jkIRNiUx0NkZpzK6tc,39035
+snowglobe/client/src/cli.py,sha256=E3KZWEHiF38BejVIVgSRMINxKRKOS7p_aemhX4xt90E,27846
+snowglobe/client/src/cli_utils.py,sha256=NZ018IISxt6hwaxbAraWGoRDMWMsnAeLBhdr3N1ZWDw,17104
+snowglobe/client/src/config.py,sha256=YRx_AQEZoHaAqk6guTxynIEGV_iJ3wNNGtMmaKsYMbc,10488
+snowglobe/client/src/models.py,sha256=BX310WrDN9Fd8v68me3XGL_ic1ulvjCrZyIT2ND1eUo,866
+snowglobe/client/src/project_manager.py,sha256=Ze-qs4dQI2kIV-PmtWZ1b67hMUfsnsMHus90aT8HOow,9970
+snowglobe/client/src/stats.py,sha256=IdaXroOZBmvLVa_p9pDE6hsxsc7-fBEDnLf8O6Ch0GA,1596
+snowglobe/client/src/utils.py,sha256=hHOht0hc8fv3OuPTz2Tqs639CzSAF34JTZs5ifKV6YI,3708
+snowglobe-0.4.2.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
+snowglobe-0.4.2.dist-info/METADATA,sha256=_0jWIHcgJTgXb84y8G8IOW_5HsO6dVoqn25Avs_uqh4,4483
+snowglobe-0.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+snowglobe-0.4.2.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
+snowglobe-0.4.2.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
+snowglobe-0.4.2.dist-info/RECORD,,

snowglobe-0.4.0.dist-info/RECORD

@@ -1,15 +0,0 @@
-snowglobe/client/__init__.py,sha256=kzp9wPUUYBXqDSKZbfmD4vrAQvrWSW5HOvtpFlEJWfs,353
-snowglobe/client/src/app.py,sha256=PZCTwF3n_2Bi825dOflGeN8GSjuf5ebKLSxV7S90TxA,30289
-snowglobe/client/src/cli.py,sha256=f8AKB0mAcMK-63aZczs4LxnB33RqoxKPEnv2w1lgu1c,25159
-snowglobe/client/src/cli_utils.py,sha256=PRWpbrALfTc3fpAXl32pfyZWkLYi_3N03csTQR1wnjc,11911
-snowglobe/client/src/config.py,sha256=HAJD7RkO6IB_mFkmGLpy9Ma3sB0KpZE7zmGsa9K__iE,9284
-snowglobe/client/src/models.py,sha256=BX310WrDN9Fd8v68me3XGL_ic1ulvjCrZyIT2ND1eUo,866
-snowglobe/client/src/project_manager.py,sha256=Ze-qs4dQI2kIV-PmtWZ1b67hMUfsnsMHus90aT8HOow,9970
-snowglobe/client/src/stats.py,sha256=IdaXroOZBmvLVa_p9pDE6hsxsc7-fBEDnLf8O6Ch0GA,1596
-snowglobe/client/src/utils.py,sha256=U0nQjTjO28OghG7lV6BuI_2MkcsOhu31Nz00nCJU4sM,3670
-snowglobe-0.4.0.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
-snowglobe-0.4.0.dist-info/METADATA,sha256=OEuNnqFkNHhu3z0cgORX17kyKGOhoAvo7IxUfkopMTM,4921
-snowglobe-0.4.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-snowglobe-0.4.0.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
-snowglobe-0.4.0.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
-snowglobe-0.4.0.dist-info/RECORD,,