dreadnode 2.0.6__tar.gz → 2.0.8__tar.gz

{dreadnode-2.0.6 → dreadnode-2.0.8}/PKG-INFO

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

{dreadnode-2.0.6 → dreadnode-2.0.8}/dreadnode/agents/__init__.py

@@ -2,7 +2,6 @@ import importlib
 import typing as t
 
 from dreadnode.agents import exceptions
-from dreadnode.agents.agent import Agent
 from dreadnode.agents.tools import (
     FunctionCall,
     FunctionDefinition,

{dreadnode-2.0.6 → dreadnode-2.0.8}/dreadnode/agents/agent.py

@@ -669,6 +669,18 @@ class Agent(Executor[AgentEvent, Trajectory]):
                         msg.metadata.setdefault("model", self.model_name)
                     messages.extend(step_chat.generated)
 
+                    # Emit content for TUI rendering BEFORE tools/stop-checks
+                    # so text is visible immediately (ENG-5879)
+                    last_msg = step_chat.generated[-1] if step_chat.generated else None
+                    if last_msg and last_msg.content:
+                        yield GenerationContent(
+                            agent_id=self.agent_id,
+                            agent_name=self.name,
+                            step=step_count,
+                            content=str(last_msg.content),
+                            extra=step_chat.extra,
+                        )
+
                     # Check stop conditions INSIDE span
                     if any(cond(self.trajectory.steps) for cond in self.stop_conditions):
                         logger.info("A stop condition was met. Ending run.")
@@ -794,19 +806,6 @@ class Agent(Executor[AgentEvent, Trajectory]):
                         finally:
                             await _event_queue.put(None)  # Signal done
 
-                    # Emit generation content BEFORE tools so TUI can
-                    # render text immediately (ENG-5879)
-                    last_msg = step_chat.generated[-1] if step_chat.generated else None
-                    if last_msg and last_msg.content:
-                        content_event = GenerationContent(
-                            agent_id=self.agent_id,
-                            agent_name=self.name,
-                            step=step_count,
-                            content=str(last_msg.content)[:4000],
-                            extra=step_chat.extra,
-                        )
-                        yield content_event
-
                     # Start tools in background
                     tools_task = asyncio.create_task(run_all_tools())
 

{dreadnode-2.0.6 → dreadnode-2.0.8}/dreadnode/agents/events.py

@@ -890,9 +890,26 @@ class GenerationError(AgentEvent):
     messages: list["Message"] = Field(default_factory=list)
 
     def _get_data(self) -> dict[str, t.Any]:
+        error_str = str(self.error)
+        # When litellm wraps provider exceptions with empty messages (e.g.
+        # "AnthropicException - ."), enrich with the underlying cause so the
+        # user gets something actionable instead of a cryptic dot.
+        if error_str.endswith(
+            ("Exception - .", "Exception - . Handle with `litellm.InternalServerError`.")
+        ):
+            cause = getattr(self.error, "__cause__", None) or getattr(
+                self.error, "__context__", None
+            )
+            if cause:
+                error_str = f"{error_str} (cause: {cause})"
+            else:
+                error_str = (
+                    f"{type(self.error).__name__}: The API returned an error with no details. "
+                    "This is usually a transient issue — try again."
+                )
         return {
             "model": self.generator.model if self.generator else None,
-            "error": str(self.error),
+            "error": error_str,
             "error_type": type(self.error).__name__,
             "step": self.step,
         }

{dreadnode-2.0.6 → dreadnode-2.0.8}/dreadnode/agents/mcp/client.py

@@ -144,6 +144,10 @@ class MCPClient:
     _exit_stack: AsyncExitStack
     _session: "ClientSession | None"
     _oauth_config: t.Any  # OAuthConfig | None
+    _owner_task: asyncio.Task[None] | None
+    _shutdown_event: asyncio.Event | None
+    _ready_future: asyncio.Future[None] | None
+    _lifecycle_lock: asyncio.Lock | None
 
     def __init__(
         self,
@@ -173,6 +177,10 @@ class MCPClient:
         self._oauth_config = oauth
         self._init_timeout = init_timeout
         self._stderr_capture: _StderrCapture | None = None
+        self._owner_task = None
+        self._shutdown_event = None
+        self._ready_future = None
+        self._lifecycle_lock = None
 
     @classmethod
     def from_config(cls, config: ServerConfig) -> "MCPClient":
@@ -222,6 +230,83 @@ class MCPClient:
 
         return execute_on_server
 
+    def _get_lifecycle_lock(self) -> asyncio.Lock:
+        if self._lifecycle_lock is None:
+            self._lifecycle_lock = asyncio.Lock()
+        return self._lifecycle_lock
+
+    @staticmethod
+    def _unwrap_exception(exc: BaseException) -> BaseException:
+        cause = exc
+        if isinstance(exc, BaseExceptionGroup):
+            leaf = exc.exceptions[0] if exc.exceptions else exc
+            while isinstance(leaf, BaseExceptionGroup) and leaf.exceptions:
+                leaf = leaf.exceptions[0]
+            cause = leaf
+        return cause
+
+    def _set_error_status(self, exc: BaseException) -> BaseException:
+        cause = self._unwrap_exception(exc)
+        error_msg = str(cause)
+
+        if self._stderr_capture and self._stderr_capture.last_lines:
+            stderr_tail = "\n".join(self._stderr_capture.last_lines[-10:])
+            error_msg = f"{error_msg}\n\nServer stderr:\n{stderr_tail}"
+
+        error_lower = error_msg.lower()
+        if "unauthorized" in error_lower or "401" in error_lower or "oauth" in error_lower:
+            self._status = MCPStatus.NEEDS_AUTH
+        else:
+            self._status = MCPStatus.FAILED
+
+        self._error = error_msg
+        return cause
+
+    def _reset_connection_state(self) -> None:
+        self._session = None
+        self.tools.clear()
+        if self._status == MCPStatus.CONNECTED:
+            self._status = MCPStatus.DISCONNECTED
+        self._exit_stack = AsyncExitStack()
+        self._stderr_capture = None
+
+    async def _run_connection(self, shutdown_event: asyncio.Event) -> None:
+        try:
+            if self.transport == "stdio":
+                self._session = await self._connect_via_stdio(
+                    t.cast("StdioConnection", self.connection),
+                )
+            elif self.transport == "streamable-http":
+                self._session = await self._connect_via_streamable_http(
+                    t.cast("dict[str, t.Any]", self.connection),
+                )
+            else:
+                msg = (
+                    f"Unsupported transport: {self.transport}. Must be 'stdio' or 'streamable-http'"
+                )
+                raise TypeError(msg)  # noqa: TRY301
+
+            await asyncio.wait_for(self.session.initialize(), timeout=self._init_timeout)
+            await asyncio.wait_for(self._load_tools(), timeout=self._init_timeout)
+
+            self._status = MCPStatus.CONNECTED
+            self._error = None
+
+            if self._ready_future is not None and not self._ready_future.done():
+                self._ready_future.set_result(None)
+
+            await shutdown_event.wait()
+        except BaseException as exc:
+            cause = self._set_error_status(exc)
+            if self._ready_future is not None and not self._ready_future.done():
+                if isinstance(cause, Exception):
+                    self._ready_future.set_exception(cause)
+                else:
+                    self._ready_future.set_exception(RuntimeError(str(cause)))
+        finally:
+            await self._exit_stack.aclose()
+            self._reset_connection_state()
+
     async def _load_tools(self) -> None:
         mcp_tool_result = await self.session.list_tools()
 
@@ -378,73 +463,43 @@ class MCPClient:
 
         Sets status to CONNECTED on success, FAILED or NEEDS_AUTH on error.
         """
-        try:
-            if self.transport == "stdio":
-                self._session = await self._connect_via_stdio(
-                    t.cast("StdioConnection", self.connection),
-                )
-            elif self.transport == "streamable-http":
-                self._session = await self._connect_via_streamable_http(
-                    t.cast("dict[str, t.Any]", self.connection),
-                )
+        async with self._get_lifecycle_lock():
+            if self._owner_task is not None and not self._owner_task.done():
+                ready_future = self._ready_future
             else:
-                raise TypeError(  # noqa: TRY301 - intentional routing through shared connection failure handling
-                    f"Unsupported transport: {self.transport}. "
-                    "Must be 'stdio' or 'streamable-http'",
+                self._reset_connection_state()
+                self._shutdown_event = asyncio.Event()
+                self._ready_future = asyncio.get_running_loop().create_future()
+                self._owner_task = asyncio.create_task(
+                    self._run_connection(self._shutdown_event),
+                    name=f"mcp-client:{self.transport}",
                 )
+                ready_future = self._ready_future
 
-            await asyncio.wait_for(self.session.initialize(), timeout=self._init_timeout)
-            await asyncio.wait_for(self._load_tools(), timeout=self._init_timeout)
-
-            self._status = MCPStatus.CONNECTED
-            self._error = None
-
-        except BaseException as e:
-            await self._shutdown()
-
-            # anyio TaskGroups wrap failures in BaseExceptionGroup — unwrap
-            # to surface the actual root cause (e.g. ConnectionRefusedError)
-            # instead of the opaque "unhandled errors in a TaskGroup" message.
-            cause = e
-            if isinstance(e, BaseExceptionGroup):
-                leaf = e.exceptions[0] if e.exceptions else e
-                while isinstance(leaf, BaseExceptionGroup) and leaf.exceptions:
-                    leaf = leaf.exceptions[0]
-                cause = leaf
-
-            error_msg = str(cause)
-
-            # Append captured subprocess stderr for stdio failures —
-            # this surfaces the actual crash output (tracebacks, missing
-            # modules, config errors) instead of just "Connection closed".
-            if self._stderr_capture and self._stderr_capture.last_lines:
-                stderr_tail = "\n".join(self._stderr_capture.last_lines[-10:])
-                error_msg = f"{error_msg}\n\nServer stderr:\n{stderr_tail}"
-
-            # Check for auth-related failures
-            error_lower = error_msg.lower()
-            if "unauthorized" in error_lower or "401" in error_lower or "oauth" in error_lower:
-                self._status = MCPStatus.NEEDS_AUTH
-            else:
-                self._status = MCPStatus.FAILED
-
-            self._error = error_msg
-
-            # Re-raise the unwrapped cause so callers can catch with `except Exception`
-            if cause is not e and isinstance(cause, Exception):
-                raise cause from e
-            raise
+        assert ready_future is not None
+        await ready_future
 
     async def disconnect(self) -> None:
         """Disconnect from the MCP server."""
         await self._shutdown()
 
     async def _shutdown(self) -> None:
-        await self._exit_stack.aclose()
-        self._session = None
-        self.tools.clear()
-        if self._status == MCPStatus.CONNECTED:
-            self._status = MCPStatus.DISCONNECTED
+        async with self._get_lifecycle_lock():
+            owner_task = self._owner_task
+            shutdown_event = self._shutdown_event
+            if owner_task is None:
+                self._reset_connection_state()
+                return
+            if shutdown_event is not None:
+                shutdown_event.set()
+
+        await owner_task
+
+        async with self._get_lifecycle_lock():
+            if self._owner_task is owner_task:
+                self._owner_task = None
+                self._shutdown_event = None
+                self._ready_future = None
 
     async def __aenter__(self) -> "MCPClient":
         await self.connect()

{dreadnode-2.0.6 → dreadnode-2.0.8}/dreadnode/airt/assessment.py

@@ -36,7 +36,7 @@ if t.TYPE_CHECKING:
 
     from dreadnode.airt.analytics.types import GoalCategory
     from dreadnode.app.api.client import ApiClient
-    from dreadnode.app.server.session import Session
+    from dreadnode.app.config import Profile
     from dreadnode.core.task import Task
     from dreadnode.optimization.study import Study
 
@@ -45,13 +45,13 @@ _current_assessment: contextvars.ContextVar[Assessment | None] = contextvars.Con
 )
 
 
-def _get_platform_context() -> tuple[ApiClient, Session] | None:
-    """Get the API client and session if connected to the platform."""
+def _get_platform_context() -> tuple[ApiClient, Profile] | None:
+    """Get the API client and profile if connected to the platform."""
     try:
         from dreadnode.app.main import DEFAULT_INSTANCE
 
         if DEFAULT_INSTANCE.can_sync:
-            return DEFAULT_INSTANCE.api, DEFAULT_INSTANCE.session
+            return DEFAULT_INSTANCE.api, DEFAULT_INSTANCE.profile
     except Exception:
         logger.debug("Unable to resolve platform context for AIRT assessment")
     return None
@@ -82,6 +82,9 @@ class Assessment:
         attack_defaults: dict[str, t.Any] | None = None,
         description: str | None = None,
         session_id: str | None = None,
+        target_model: str | None = None,
+        attacker_model: str | None = None,
+        judge_model: str | None = None,
         target_config: dict[str, t.Any] | None = None,
         attacker_config: dict[str, t.Any] | None = None,
         attack_manifest: list[dict[str, t.Any]] | None = None,
@@ -97,6 +100,11 @@ class Assessment:
         self.goal_category = goal_category
         self._attack_defaults = attack_defaults or {}
 
+        # First-class model identifiers for platform UI
+        self.target_model = target_model or model
+        self.attacker_model = attacker_model or model
+        self.judge_model = judge_model or model
+
         if target_config is None and model is not None:
             target_config = {"model": model}
         if attacker_config is None and model is not None:
@@ -143,15 +151,18 @@ class Assessment:
         Returns:
             The platform assessment ID, or None if offline.
         """
+        if self._assessment_id is not None:
+            return self._assessment_id
+
         ctx = _get_platform_context()
         if ctx is None:
             logger.debug("SDK offline -- assessment not registered with platform")
             return None
 
-        api, session = ctx
+        api, profile = ctx
         project_id = self._project_id
-        if project_id is None and session.project is not None:
-            project_id = session.project.id
+        if project_id is None and profile.project_id is not None:
+            project_id = profile.project_id
 
         if project_id is None:
             logger.warning("No project_id available -- cannot register assessment")
@@ -159,12 +170,15 @@ class Assessment:
 
         try:
             result = api.create_airt_assessment(
-                session.org_key,
-                session.workspace_key,
+                profile.org_key,
+                profile.workspace_key,
                 name=self.name,
                 project_id=project_id,
                 description=self.description,
                 session_id=self._session_id,
+                target_model=self.target_model,
+                attacker_model=self.attacker_model,
+                judge_model=self.judge_model,
                 target_config=self.target_config,
                 attacker_config=self.attacker_config,
                 attack_manifest=self._attack_manifest,
@@ -319,7 +333,7 @@ class Assessment:
 
         merged = {**self._attack_defaults, **kwargs}
         merged.setdefault("airt_goal_category", self.goal_category)
-        merged.setdefault("airt_target_model", model)
+        merged.setdefault("airt_target_model", self.target_model or model)
 
         return factory(goal, target, model, model, **merged)
 
@@ -383,9 +397,9 @@ class Assessment:
         if ctx is None or self._assessment_id is None:
             return None
 
-        api, session = ctx
-        org = session.org_key
-        ws = session.workspace_key
+        api, profile = ctx
+        org = profile.org_key
+        ws = profile.workspace_key
 
         try:
             trace_id = uuid.uuid4().hex
@@ -439,11 +453,11 @@ class Assessment:
         if ctx is None or self._assessment_id is None:
             return False
 
-        api, session = ctx
+        api, profile = ctx
         try:
             api.update_airt_assessment(
-                session.org_key,
-                session.workspace_key,
+                profile.org_key,
+                profile.workspace_key,
                 self._assessment_id,
                 status="completed",
             )
@@ -467,14 +481,14 @@ class Assessment:
         if ctx is None or self._assessment_id is None:
             return False
 
-        api, session = ctx
+        api, profile = ctx
         try:
             kwargs: dict[str, t.Any] = {"status": "failed"}
             if reason:
                 kwargs["description"] = f"{self.description or ''}\n\nFailure: {reason}".strip()
             api.update_airt_assessment(
-                session.org_key,
-                session.workspace_key,
+                profile.org_key,
+                profile.workspace_key,
                 self._assessment_id,
                 **kwargs,
             )