dreadnode 2.0.23__tar.gz → 2.0.24__tar.gz

{dreadnode-2.0.23 → dreadnode-2.0.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 2.0.23
+Version: 2.0.24
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://dreadnode.io
 Project-URL: Documentation, https://docs.dreadnode.io

{dreadnode-2.0.23 → dreadnode-2.0.24}/dreadnode/airt/assessment.py

@@ -23,6 +23,7 @@ import contextlib
 import contextvars
 import os
 import signal
+import threading
 import typing as t
 from pathlib import Path
 
@@ -57,6 +58,26 @@ def _get_platform_context() -> tuple[ApiClient, Profile] | None:
     return None
 
 
+# Thread-local storage for injected platform context
+_thread_local = threading.local()
+
+
+def _set_platform_context(api: ApiClient, profile: Profile) -> None:
+    """Set platform context for the current thread (used by AIRT CLI)."""
+    _thread_local.api = api
+    _thread_local.profile = profile
+
+
+def _get_platform_context_with_fallback() -> tuple[ApiClient, Profile] | None:
+    """Get platform context with CLI injection fallback."""
+    # First try thread-local injected context (from AIRT CLI)
+    if hasattr(_thread_local, "api") and hasattr(_thread_local, "profile"):
+        return _thread_local.api, _thread_local.profile
+
+    # Fallback to global DEFAULT_INSTANCE
+    return _get_platform_context()
+
+
 class Assessment:
     """Orchestrates multi-attack assessments.
 
@@ -91,6 +112,7 @@ class Assessment:
         workflow_run_id: str | None = None,
         workflow_script: str | None = None,
         project_id: str | None = None,
+        runtime_id: str | None = None,
     ) -> None:
         self.name = name
         self.description = description
@@ -116,6 +138,7 @@ class Assessment:
         self._workflow_run_id = workflow_run_id
         self._workflow_script = workflow_script
         self._project_id = project_id
+        self._runtime_id = runtime_id
 
         self._session_id = session_id or os.environ.get("DREADNODE_SESSION_ID")
         if self._session_id is None:
@@ -177,7 +200,7 @@ class Assessment:
         if self._assessment_id is not None:
             return self._assessment_id
 
-        ctx = _get_platform_context()
+        ctx = _get_platform_context_with_fallback()
         if ctx is None:
             logger.debug("SDK offline -- assessment not registered with platform")
             return None
@@ -192,11 +215,28 @@ class Assessment:
             return None
 
         try:
+            # Get runtime_id - try to use a default runtime for the project
+            runtime_id = getattr(self, "_runtime_id", None)
+            if runtime_id is None:
+                # Auto-discover a runtime for this project (CLI users shouldn't need to specify this)
+                try:
+                    runtimes_resp = api.list_runtimes(profile.org_key, profile.workspace_key)
+                    runtime_items = runtimes_resp.get("items", [])
+                    matching_runtimes = [r for r in runtime_items if r["project_id"] == project_id]
+                    if matching_runtimes:
+                        runtime_id = matching_runtimes[0]["id"]
+                        logger.debug(f"Auto-selected runtime {runtime_id} for project {project_id}")
+                    else:
+                        logger.warning(f"No runtimes found for project {project_id}")
+                except Exception as e:
+                    logger.debug(f"Failed to auto-discover runtime for project: {e}")
+
             result = api.create_airt_assessment(
                 profile.org_key,
                 profile.workspace_key,
                 name=self.name,
                 project_id=project_id,
+                runtime_id=runtime_id,
                 description=self.description,
                 session_id=self._session_id,
                 target_model=self.target_model,
@@ -402,7 +442,7 @@ class Assessment:
             await self.register()
             # Notify the platform that execution has started (sets started_at server-side)
             if self._assessment_id is not None:
-                ctx = _get_platform_context()
+                ctx = _get_platform_context_with_fallback()
                 if ctx is not None:
                     api, session = ctx
                     try:
@@ -459,7 +499,7 @@ class Assessment:
                 provider.force_flush(timeout_millis=10_000)
 
         # Use the sync API client directly — no async gymnastics
-        ctx = _get_platform_context()
+        ctx = _get_platform_context_with_fallback()
         if ctx is None or self._assessment_id is None:
             return
 
@@ -499,7 +539,7 @@ class Assessment:
         Returns:
             True if successfully marked, False otherwise.
         """
-        ctx = _get_platform_context()
+        ctx = _get_platform_context_with_fallback()
         if ctx is None or self._assessment_id is None:
             return False
 
@@ -527,7 +567,7 @@ class Assessment:
         Returns:
             True if successfully marked, False otherwise.
         """
-        ctx = _get_platform_context()
+        ctx = _get_platform_context_with_fallback()
         if ctx is None or self._assessment_id is None:
             return False
 

{dreadnode-2.0.23 → dreadnode-2.0.24}/dreadnode/app/cli/airt.py

@@ -65,7 +65,7 @@ def _parse_json_object_option(value: str | None, *, field_name: str) -> dict[str
 
 
 def _resolve_project_id(api, profile, project):
-    """Resolve project name to UUID."""
+    """Resolve project name to UUID, creating the project if it doesn't exist."""
     from uuid import UUID
 
     candidate = project or getattr(profile, "project_key", None)
@@ -83,9 +83,25 @@ def _resolve_project_id(api, profile, project):
     try:
         resolved = api.get_project(profile.org_key, profile.workspace_key, str(candidate))
         return str(resolved.id)
-    except Exception as e:
-        print(f"[PROJECT DEBUG] Failed to resolve project '{candidate}': {e}")
-        return None
+    except Exception as get_error:
+        # Project doesn't exist, try to create it
+        try:
+            print(f"[PROJECT DEBUG] Project '{candidate}' not found, creating it...")
+            created = api.create_project(
+                profile.org_key,
+                profile.workspace_key,
+                name=str(candidate),
+                key=str(candidate),
+                description="Auto-created project for AIRT CLI run",
+            )
+            print(
+                f"[PROJECT DEBUG] Successfully created project '{candidate}' with ID {created.id}"
+            )
+            return str(created.id)
+        except Exception as create_error:
+            print(f"[PROJECT DEBUG] Failed to resolve project '{candidate}': {get_error}")
+            print(f"[PROJECT DEBUG] Failed to create project '{candidate}': {create_error}")
+            return None
 
 
 _SEVERITY_COLORS = {
@@ -1150,6 +1166,9 @@ def run(
             print_error(str(e))
             return
 
+    # Capture original project argument before connect() clears it
+    original_project = platform.project
+
     # Configure SDK and validate scope to populate project_id
     _api, profile = platform.connect()
 
@@ -1175,25 +1194,46 @@ def run(
     console.print()
 
     async def _run() -> None:
+        # Determine project for SDK configuration
+        # Use original CLI argument, fallback to profile, environment, then 'default'
+        import os
+
         from dreadnode.airt import Assessment
+        from dreadnode.airt.assessment import _set_platform_context
 
         # Configure the SDK using the validated profile
         from dreadnode.app.main import DEFAULT_INSTANCE
 
-        DEFAULT_INSTANCE.configure(
-            server=profile.url,
-            api_key=profile.api_key,
-            organization=profile.organization,
-            workspace=profile.workspace,
-            project=profile.project,
-        )
+        env_project = os.getenv("DREADNODE_PROJECT")
+        sdk_project = original_project or profile.project or env_project or "default"
+
+        # Try global SDK configuration first
+        try:
+            DEFAULT_INSTANCE.configure(
+                server=profile.url,
+                api_key=profile.api_key,
+                organization=profile.organization,
+                workspace=profile.workspace,
+                project=sdk_project,
+            )
+            from loguru import logger
+
+            logger.debug("SDK instance configured successfully for assessment registration")
+        except Exception as e:
+            from loguru import logger
+
+            logger.warning(f"Failed to configure SDK instance for assessment registration: {e}")
+            logger.info("Using fallback context injection for assessment registration")
+
+        # Always inject platform context as fallback for assessment registration
+        _set_platform_context(_api, profile)
 
         # Resolve project name to UUID
-        # Get project from environment if profile.project is None
+        # Use original CLI argument, fallback to profile, environment, then 'default'
         import os
 
         env_project = os.getenv("DREADNODE_PROJECT")
-        actual_project = profile.project or env_project
+        actual_project = original_project or profile.project or env_project or "default"
 
         resolved_project_id = _resolve_project_id(_api, profile, actual_project)
         final_project_id = resolved_project_id or profile.project_id or actual_project
@@ -1332,24 +1372,48 @@ def run_suite(
     console.print(f"  Config:  [dim]{file}[/dim]")
     console.print()
 
+    # Capture original project argument before connect() clears it
+    original_project = platform.project
+
     # Configure SDK
     _api, profile = platform.connect()
 
     target_fn = _build_target(model, max_tokens)
 
     async def _run_suite() -> list[dict[str, t.Any]]:
+        # Determine project for SDK configuration
+        # Use original CLI argument, fallback to profile, environment, then 'default'
+        import os
+
         from dreadnode.airt import Assessment
+        from dreadnode.airt.assessment import _set_platform_context
 
         # Configure the SDK using the validated profile
         from dreadnode.app.main import DEFAULT_INSTANCE
 
-        DEFAULT_INSTANCE.configure(
-            server=profile.url,
-            api_key=profile.api_key,
-            organization=profile.organization,
-            workspace=profile.workspace,
-            project=profile.project,
-        )
+        env_project = os.getenv("DREADNODE_PROJECT")
+        sdk_project = original_project or profile.project or env_project or "default"
+
+        # Try global SDK configuration first
+        try:
+            DEFAULT_INSTANCE.configure(
+                server=profile.url,
+                api_key=profile.api_key,
+                organization=profile.organization,
+                workspace=profile.workspace,
+                project=sdk_project,
+            )
+            from loguru import logger
+
+            logger.debug("SDK instance configured successfully for assessment registration")
+        except Exception as e:
+            from loguru import logger
+
+            logger.warning(f"Failed to configure SDK instance for assessment registration: {e}")
+            logger.info("Using fallback context injection for assessment registration")
+
+        # Always inject platform context as fallback for assessment registration
+        _set_platform_context(_api, profile)
 
         suite_results = []
         for i, goal_cfg in enumerate(goals, 1):
@@ -1362,11 +1426,11 @@ def run_suite(
             console.print(f"\n--- Goal {i}/{len(goals)}: {goal_text[:60]}...")
 
             # Resolve project name to UUID
-            # Get project from environment if profile.project is None
+            # Use original CLI argument, then environment variable as fallback
             import os
 
             env_project = os.getenv("DREADNODE_PROJECT")
-            actual_project = profile.project or env_project
+            actual_project = original_project or profile.project or env_project
 
             resolved_project_id = _resolve_project_id(_api, profile, actual_project)
             final_project_id = resolved_project_id or profile.project_id or actual_project

{dreadnode-2.0.23 → dreadnode-2.0.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreadnode"
-version = "2.0.23"
+version = "2.0.24"
 description = "Dreadnode SDK"
 authors = [{ name = "Dreadnode Team", email = "support@dreadnode.io" }]
 readme = "README.md"