feat: Toolset Context Managers

docs/sdk/agent_tools.mdx

@@ -21,6 +21,7 @@ Inheriting from this class provides:
 - Pydantic's declarative syntax for defining state (fields).
 - Automatic application of the `@configurable` decorator.
 - A `get_tools` method for discovering methods decorated with `@dreadnode.tool_method`.
+- Support for async context management, with automatic re-entrancy handling.
 
 ### name
 

docs/sdk/api.mdx

@@ -172,6 +172,7 @@ def create_project(
 ```python
 create_workspace(
     name: str,
+    key: str,
     organization_id: UUID,
     description: str | None = None,
 ) -> Workspace
@@ -198,6 +199,7 @@ Creates a new workspace.
 def create_workspace(
     self,
     name: str,
+    key: str,
     organization_id: UUID,
     description: str | None = None,
 ) -> Workspace:
@@ -214,6 +216,7 @@ def create_workspace(
 
     payload = {
         "name": name,
+        "key": key,
         "description": description,
         "org_id": str(organization_id),
     }
@@ -237,7 +240,7 @@ Deletes a specific workspace.
 
 * **`workspace_id`**
   (`str | UUID`)
-  –The workspace identifier.
+  –The workspace key.
 
 <Accordion title="Source code in dreadnode/api/client.py" icon="code">
 ```python
@@ -246,7 +249,7 @@ def delete_workspace(self, workspace_id: str | UUID) -> None:
     Deletes a specific workspace.
 
     Args:
-        workspace_id (str | UUID): The workspace identifier.
+        workspace_id (str | UUID): The workspace key.
     """
 
     self.request("DELETE", f"/workspaces/{workspace_id!s}")
@@ -707,17 +710,15 @@ def get_github_access_token(self, repos: list[str]) -> GithubTokenResponse:
 ### get\_organization
 
 ```python
-get_organization(
-    organization_id: str | UUID,
-) -> Organization
+get_organization(org_id_or_key: UUID | str) -> Organization
 ```
 
 Retrieves details of a specific organization.
 
 **Parameters:**
 
-* **`organization_id`**
-  (`str`)
+* **`org_id_or_key`**
+  (`str | UUID`)
   –The organization identifier.
 
 **Returns:**
@@ -727,17 +728,17 @@ Retrieves details of a specific organization.
 
 <Accordion title="Source code in dreadnode/api/client.py" icon="code">
 ```python
-def get_organization(self, organization_id: str | UUID) -> Organization:
+def get_organization(self, org_id_or_key: UUID | str) -> Organization:
     """
     Retrieves details of a specific organization.
 
     Args:
-        organization_id (str): The organization identifier.
+        org_id_or_key (str | UUID): The organization identifier.
 
     Returns:
         Organization: The Organization object.
     """
-    response = self.request("GET", f"/organizations/{organization_id!s}")
+    response = self.request("GET", f"/organizations/{org_id_or_key!s}")
     return Organization(**response.json())
 ```
 
@@ -1098,16 +1099,17 @@ def get_user_data_credentials(self) -> UserDataCredentials:
 
 ```python
 get_workspace(
-    workspace_id: str | UUID, org_id: UUID | None = None
+    workspace_id_or_key: UUID | str,
+    org_id: UUID | None = None,
 ) -> Workspace
 ```
 
 Retrieves details of a specific workspace.
 
 **Parameters:**
 
-* **`workspace_id`**
-  (`str`)
+* **`workspace_id_or_key`**
+  (`str | UUID`)
   –The workspace identifier.
 
 **Returns:**
@@ -1117,20 +1119,22 @@ Retrieves details of a specific workspace.
 
 <Accordion title="Source code in dreadnode/api/client.py" icon="code">
 ```python
-def get_workspace(self, workspace_id: str | UUID, org_id: UUID | None = None) -> Workspace:
+def get_workspace(
+    self, workspace_id_or_key: UUID | str, org_id: UUID | None = None
+) -> Workspace:
     """
     Retrieves details of a specific workspace.
 
     Args:
-        workspace_id (str): The workspace identifier.
+        workspace_id_or_key (str | UUID): The workspace identifier.
 
     Returns:
         Workspace: The Workspace object.
     """
     params: dict[str, str] = {}
     if org_id:
         params = {"org_id": str(org_id)}
-    response = self.request("GET", f"/workspaces/{workspace_id!s}", params=params)
+    response = self.request("GET", f"/workspaces/{workspace_id_or_key!s}", params=params)
     return Workspace(**response.json())
 ```
 
@@ -1793,6 +1797,14 @@ is_active: bool
 
 Is the organization active?
 
+### key
+
+```python
+key: str
+```
+
+URL-friendly identifier for the organization.
+
 ### max\_members
 
 ```python
@@ -1809,14 +1821,6 @@ name: str
 
 Name of the organization.
 
-### slug
-
-```python
-slug: str
-```
-
-URL-friendly slug for the organization.
-
 ### updated\_at
 
 ```python
@@ -2321,6 +2325,14 @@ is_default: bool
 
 Is the workspace the default one?
 
+### key
+
+```python
+key: str
+```
+
+Unique key for the workspace.
+
 ### name
 
 ```python
@@ -2353,14 +2365,6 @@ project_count: int | None
 
 Number of projects in the workspace.
 
-### slug
-
-```python
-slug: str
-```
-
-URL-friendly slug for the workspace.
-
 ### updated\_at
 
 ```python

docs/sdk/main.mdx

@@ -173,10 +173,13 @@ in the following order:
 1. Environment variables:
 2. `DREADNODE_SERVER_URL` or `DREADNODE_SERVER`
 3. `DREADNODE_API_TOKEN` or `DREADNODE_API_KEY`
-4. Dreadnode profile (from `dreadnode login`)
-5. Uses `profile` parameter if provided
-6. Falls back to `DREADNODE_PROFILE` environment variable
-7. Defaults to active profile
+4. `DREADNODE_ORGANIZATION`
+5. `DREADNODE_WORKSPACE`
+6. `DREADNODE_PROJECT`
+7. Dreadnode profile (from `dreadnode login`)
+8. Uses `profile` parameter if provided
+9. Falls back to `DREADNODE_PROFILE` environment variable
+10. Defaults to active profile
 
 **Parameters:**
 
@@ -214,7 +217,7 @@ in the following order:
   (`str | None`, default:
   `None`
   )
-  –The default project name to associate all runs with. This can also be in the format `org/workspace/project`.
+  –The default project name to associate all runs with. This can also be in the format `org/workspace/project` using the keys.
 * **`service_name`**
   (`str | None`, default:
   `None`
@@ -270,6 +273,10 @@ def configure(
     1. Environment variables:
        - `DREADNODE_SERVER_URL` or `DREADNODE_SERVER`
        - `DREADNODE_API_TOKEN` or `DREADNODE_API_KEY`
+       - `DREADNODE_ORGANIZATION`
+       - `DREADNODE_WORKSPACE`
+       - `DREADNODE_PROJECT`
+
     2. Dreadnode profile (from `dreadnode login`)
        - Uses `profile` parameter if provided
        - Falls back to `DREADNODE_PROFILE` environment variable
@@ -282,7 +289,7 @@ def configure(
         local_dir: The local directory to store data in.
         organization: The default organization name or ID to use.
         workspace: The default workspace name or ID to use.
-        project: The default project name to associate all runs with. This can also be in the format `org/workspace/project`.
+        project: The default project name to associate all runs with. This can also be in the format `org/workspace/project` using the keys.
         service_name: The service name to use for OpenTelemetry.
         service_version: The service version to use for OpenTelemetry.
         console: Log span information to the console (`DREADNODE_CONSOLE` or the default is True).
@@ -342,19 +349,8 @@ def configure(
         self.local_dir = local_dir
 
     _org, _workspace, _project = self._extract_project_components(project)
-
     self.organization = _org or organization or os.environ.get(ENV_ORGANIZATION)
-    with contextlib.suppress(ValueError):
-        self.organization = UUID(
-            str(self.organization)
-        )  # Now, it's a UUID if possible, else str (name/slug)
-
     self.workspace = _workspace or workspace or os.environ.get(ENV_WORKSPACE)
-    with contextlib.suppress(ValueError):
-        self.workspace = UUID(
-            str(self.workspace)
-        )  # Now, it's a UUID if possible, else str (name/slug)
-
     self.project = _project or project or os.environ.get(ENV_PROJECT)
 
     self.service_name = service_name

docs/sdk/task.mdx

@@ -683,6 +683,7 @@ async def run_always(self, *args: P.args, **kwargs: P.kwargs) -> TaskSpan[R]:  #
 
             # Log the output
 
+            output_object_hash = None
             if log_output and (
                 not isinstance(self.log_inputs, Inherited) or seems_useful_to_serialize(output)
             ):
@@ -691,13 +692,12 @@ async def run_always(self, *args: P.args, **kwargs: P.kwargs) -> TaskSpan[R]:  #
                     output,
                     attributes={"auto": True},
                 )
-            elif run is not None:
                 # Link the output to the inputs
-                for input_object_hash in input_object_hashes:
-                    run.link_objects(output_object_hash, input_object_hash)
-
-                if create_run:
-                    run.log_output("output", output, attributes={"auto": True})
+                if run is not None:
+                    for input_object_hash in input_object_hashes:
+                        run.link_objects(output_object_hash, input_object_hash)
+            elif run is not None and create_run:
+                run.log_output("output", output, attributes={"auto": True})
 
             # Score and check assertions
 

dreadnode/agent/agent.py

@@ -2,7 +2,7 @@
 import json
 import re
 import typing as t
-from contextlib import aclosing, asynccontextmanager
+from contextlib import AsyncExitStack, aclosing, asynccontextmanager
 from copy import deepcopy
 from textwrap import dedent
 
@@ -875,8 +875,16 @@ async def stream(
         commit: CommitBehavior = "always",
     ) -> t.AsyncIterator[t.AsyncGenerator[AgentEvent, None]]:
         thread = thread or self.thread
-        async with aclosing(self._stream_traced(thread, user_input, commit=commit)) as stream:
-            yield stream
+
+        async with AsyncExitStack() as stack:
+            # Ensure all tools are properly entered if they
+            # are context managers before we start using them
+            for tool_container in self.tools:
+                if hasattr(tool_container, "__aenter__") and hasattr(tool_container, "__aexit__"):
+                    await stack.enter_async_context(tool_container)
+
+            async with aclosing(self._stream_traced(thread, user_input, commit=commit)) as stream:
+                yield stream
 
     async def run(
         self,