dreadnode 2.0.35__tar.gz → 2.0.36__tar.gz

{dreadnode-2.0.35 → dreadnode-2.0.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 2.0.35
+Version: 2.0.36
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://dreadnode.io
 Project-URL: Documentation, https://docs.dreadnode.io
@@ -51,6 +51,9 @@ Requires-Dist: transformers>=4.57.1
 Requires-Dist: uvicorn>=0.34.0
 Requires-Dist: websockets>=14.0
 Requires-Dist: xmltodict>=1.0.2
+Provides-Extra: airt-ml
+Requires-Dist: scikit-learn<2.0.0,>=1.7.1; extra == 'airt-ml'
+Requires-Dist: torch>=2.9.0; extra == 'airt-ml'
 Provides-Extra: all
 Requires-Dist: confusables<2.0.0,>=1.2.0; extra == 'all'
 Requires-Dist: moviepy<3.0.0,>=2.1.2; extra == 'all'

{dreadnode-2.0.35 → dreadnode-2.0.36}/dreadnode/agents/agent.py

@@ -133,6 +133,31 @@ def _is_transient_api_error(error: BaseException) -> bool:
     return False
 
 
+def _replayable(messages: list[Message]) -> list[Message]:
+    """Drop assistant messages that carry nothing worth sending back.
+
+    A generation stopped by a content filter leaves an empty assistant message.
+    The trajectory keeps it — it is the record of what the provider did, and
+    the transcript renders a warning from it — but replaying it on the next
+    turn sends an empty text block, which some providers reject outright
+    (see ``CacheControlOnEmptyTextFixup``) and none can use (ENG-7585).
+
+    Reasoning lives in ``metadata``, not ``content``, so a thinking-only
+    message is still replayable and must survive this filter.
+    """
+
+    def carries_nothing(message: Message) -> bool:
+        return (
+            message.role == "assistant"
+            and not message.content
+            and not message.tool_calls
+            and not message.metadata.get("reasoning_content")
+            and not message.metadata.get("thinking_blocks")
+        )
+
+    return [message for message in messages if not carries_nothing(message)]
+
+
 class Agent(Executor[AgentEvent, Trajectory]):
     """
     Agent abstraction for applying tools, event logic, and message state to LLM generation.
@@ -776,7 +801,7 @@ class Agent(Executor[AgentEvent, Trajectory]):
         """
         traj = trajectory if trajectory is not None else self.trajectory
         messages = [
-            *deepcopy(traj.messages),
+            *_replayable(deepcopy(traj.messages)),
             Message(
                 "user",
                 str(self._current_input),
@@ -925,6 +950,15 @@ class Agent(Executor[AgentEvent, Trajectory]):
                         break
 
                     step_chat.generated[-1].metadata.update(step_chat.extra)
+                    # The provider's stop reason belongs to the assistant
+                    # message, not to a particular event stream — stamping it
+                    # here (before any event is dispatched) lets hooks,
+                    # transcript reloads and non-live consumers all recover the
+                    # same outcome live GenerationEnd observers see (ENG-7585).
+                    if step_chat.stop_reason:
+                        step_chat.generated[-1].metadata["generation_stop_reason"] = (
+                            step_chat.stop_reason
+                        )
                     for msg in step_chat.generated:
                         msg.metadata.setdefault("agent", self.name)
                         msg.metadata.setdefault("model", self.model_name)

{dreadnode-2.0.35 → dreadnode-2.0.36}/dreadnode/airt/__init__.py

@@ -54,6 +54,16 @@ Image adversarial attacks:
 
 Multimodal attacks:
 - multimodal_attack: Transform-based multimodal probing (vision, audio, text)
+
+Traditional-ML attacks (query a classifier predict API via PredictionTargetSpec):
+- equation_solving_extraction: linear-model weight recovery (Tramèr'16)
+- jacobian_extraction: Jacobian dataset augmentation (Papernot'17)
+- copycat_extraction: hard-label distillation (CopycatCNN)
+- knockoff_extraction: soft-label transfer-set extraction (Knockoff Nets)
+- threshold_membership: confidence/entropy/loss-threshold membership inference (Yeom'18)
+- label_only_membership: label-only membership inference (Choquette-Choo'21)
+- confidence_inversion: MI-Face confidence-maximizing model inversion (Fredrikson'15)
+- nes_inversion: NES-based confidence-maximizing model inversion
 """
 
 from dreadnode.airt.adversarial_reasoning import adversarial_reasoning_attack
@@ -72,6 +82,27 @@ from dreadnode.airt.crescendo import crescendo_attack
 from dreadnode.airt.deep_inception import deep_inception_attack
 from dreadnode.airt.drattack import drattack
 from dreadnode.airt.echo_chamber import echo_chamber_attack
+from dreadnode.airt.evasion import (
+    bae_evasion,
+    boundary_evasion,
+    deepwordbug_evasion,
+    hopskipjump_evasion,
+    pwws_evasion,
+    simba_evasion,
+    square_evasion,
+    text_evasion,
+    textbugger_evasion,
+    textfooler_evasion,
+    zoo_evasion,
+)
+from dreadnode.airt.extraction import (
+    activethief_extraction,
+    copycat_extraction,
+    distillation_extraction,
+    equation_solving_extraction,
+    jacobian_extraction,
+    knockoff_extraction,
+)
 from dreadnode.airt.genetic_persona import genetic_persona_attack
 from dreadnode.airt.goat import goat_attack
 from dreadnode.airt.goat_v2 import goat_v2_attack
@@ -83,11 +114,23 @@ from dreadnode.airt.image import (
     simba_attack,
     zoo_attack,
 )
+from dreadnode.airt.inversion import (
+    confidence_inversion,
+    nes_inversion,
+)
 from dreadnode.airt.j2_meta import j2_meta_attack
 from dreadnode.airt.jbdistill import jbdistill_attack
 from dreadnode.airt.jbfuzz import jbfuzz_attack
 from dreadnode.airt.lrm_autonomous import lrm_autonomous_attack
 from dreadnode.airt.mapf import mapf_attack
+from dreadnode.airt.membership import (
+    entropy_membership,
+    label_only_membership,
+    lira_membership,
+    loss_membership,
+    shadow_model_membership,
+    threshold_membership,
+)
 from dreadnode.airt.multimodal import multimodal_attack
 from dreadnode.airt.nexus import nexus_attack
 from dreadnode.airt.pair import pair_attack
@@ -104,8 +147,11 @@ from dreadnode.airt.siren import siren_attack
 from dreadnode.airt.tap import tap_attack
 from dreadnode.airt.target import extract_response_text, extract_tool_calls
 from dreadnode.airt.targets import (
+    Prediction,
+    PredictionTargetSpec,
     TargetAuth,
     TargetSpec,
+    build_prediction_target,
     build_target,
     nova_sonic_target,
 )
@@ -116,8 +162,11 @@ from dreadnode.airt.watermark_removal import watermark_removal_attack
 
 __all__ = [
     "Assessment",
+    "Prediction",
+    "PredictionTargetSpec",
     "TargetAuth",
     "TargetSpec",
+    "activethief_extraction",
     "adversarial_reasoning_attack",
     "advpromptier_attack",
     "alignment_faking_attack",
@@ -127,13 +176,22 @@ __all__ = [
     "attention_shifting_attack",
     "autodan_turbo_attack",
     "autoredteamer_attack",
+    "bae_evasion",
     "beast_attack",
+    "boundary_evasion",
+    "build_prediction_target",
     "build_target",
+    "confidence_inversion",
+    "copycat_extraction",
     "cot_jailbreak_attack",
     "crescendo_attack",
     "deep_inception_attack",
+    "deepwordbug_evasion",
+    "distillation_extraction",
     "drattack",
     "echo_chamber_attack",
+    "entropy_membership",
+    "equation_solving_extraction",
     "extract_response_text",
     "extract_tool_calls",
     "genetic_persona_attack",
@@ -141,19 +199,27 @@ __all__ = [
     "goat_v2_attack",
     "gptfuzzer_attack",
     "hopskipjump_attack",
+    "hopskipjump_evasion",
     "humor_bypass_attack",
     "j2_meta_attack",
+    "jacobian_extraction",
     "jbdistill_attack",
     "jbfuzz_attack",
+    "knockoff_extraction",
+    "label_only_membership",
+    "lira_membership",
+    "loss_membership",
     "lrm_autonomous_attack",
     "mapf_attack",
     "multimodal_attack",
     "nes_attack",
+    "nes_inversion",
     "nexus_attack",
     "nova_sonic_target",
     "pair_attack",
     "persona_hijack_attack",
     "prompt_attack",
+    "pwws_evasion",
     "quantization_safety_attack",
     "rainbow_attack",
     "refusal_aware_attack",
@@ -161,12 +227,20 @@ __all__ = [
     "reward_hacking_attack",
     "salami_slicing_attack",
     "self_persuasion_attack",
+    "shadow_model_membership",
     "simba_attack",
+    "simba_evasion",
     "siren_attack",
+    "square_evasion",
     "tap_attack",
     "templatefuzz_attack",
+    "text_evasion",
+    "textbugger_evasion",
+    "textfooler_evasion",
+    "threshold_membership",
     "tmap_trajectory_attack",
     "trojail_attack",
     "watermark_removal_attack",
     "zoo_attack",
+    "zoo_evasion",
 ]

dreadnode-2.0.36/dreadnode/airt/_base.py

@@ -0,0 +1,213 @@
+"""Shared base class for the black-box traditional-ML attacks.
+
+Extraction, membership inference, and evasion all probe a black-box predict API
+through the same :class:`PredictionTargetSpec` contract, count their queries,
+capture a sample of the (input -> prediction) pairs they send, emit a top-level
+``study`` span, and enforce a query budget. This module factors that shared
+plumbing into :class:`BlackBoxAttack` so each family only implements its own
+algorithm and its own result-to-span mapping.
+
+The base also provides ``_trace_step`` - a capped per-iteration span that records
+the intermediate adversarial state (distance, best score, candidate label /
+confidence, queries so far). Because we own the algorithm loop, every step of the
+attack trajectory is visible in the platform's Traces tab, which is exactly what
+off-the-shelf frameworks (ART, TextAttack) make hard to surface.
+"""
+
+import typing as t
+from contextlib import nullcontext, suppress
+
+import numpy as np
+
+from dreadnode.airt.targets.prediction import (
+    Prediction,
+    PredictionTargetSpec,
+    QueryInput,
+    predict_many,
+)
+
+
+@t.runtime_checkable
+class AttackResult(t.Protocol):
+    """The common surface every attack result exposes so :meth:`BlackBoxAttack.run`
+    can emit spans generically."""
+
+    query_samples: list[dict[str, t.Any]]
+
+    @property
+    def best_score(self) -> float: ...
+
+    @property
+    def metrics_detail(self) -> dict[str, t.Any]: ...
+
+
+def sample_input_preview(x: t.Any, n: int = 60) -> str:
+    """A compact, display-safe preview of a queried input (text snippet or a few
+    numeric features) for the finding's query sample table."""
+    if isinstance(x, str):
+        return x if len(x) <= n else x[:n] + "..."
+    try:
+        arr = np.asarray(x, dtype=np.float64).ravel()
+    except Exception:
+        return str(x)[:n]
+    head = ", ".join(f"{v:.3g}" for v in arr[:6])
+    return f"[{head}{', ...' if arr.size > 6 else ''}] ({arr.size}d)"
+
+
+class BlackBoxAttack:
+    """Base for a configured black-box attack against a predict API.
+
+    Subclasses set :attr:`attack_domain` (`model_extraction` /
+    `membership_inference` / `adversarial_ml`), a default goal, and implement
+    :meth:`_attack_name`, :meth:`_execute`, and :meth:`_span_attributes`. The
+    query counter, query-sample capture, phase/step tracing, and the top-level
+    ``study`` span are shared here.
+    """
+
+    #: One of "model_extraction", "membership_inference", "adversarial_ml".
+    attack_domain: str = ""
+    default_goal: str = ""
+    default_goal_category: str = ""
+
+    def __init__(
+        self,
+        *,
+        target: PredictionTargetSpec,
+        modality: str = "tabular",
+        max_queries: int = 500,
+        max_iterations: int = 30,
+        seed: int | None = None,
+        airt_assessment_id: str | None = None,
+        airt_target_model: str | None = None,
+        airt_goal: str | None = None,
+        airt_goal_category: str | None = None,
+    ) -> None:
+        self.target = target
+        self.modality = modality
+        self.max_queries = max_queries
+        self.max_iterations = max_iterations
+        self.rng = np.random.default_rng(seed)
+        self.airt_assessment_id = airt_assessment_id
+        self.airt_target_model = airt_target_model or target.name
+        self.airt_goal = airt_goal or self.default_goal
+        self.airt_goal_category = airt_goal_category or self.default_goal_category
+        self._query_count = 0
+        #: A sample of real (input -> prediction) pairs so the finding can show
+        #: the actual queries the attacker sent, not just the total count.
+        self._query_samples: list[dict[str, t.Any]] = []
+
+    # -- target I/O -------------------------------------------------------- #
+    async def _query(self, xs: t.Sequence[QueryInput]) -> list[Prediction]:
+        """Query the target, count the queries, and capture a sample."""
+        preds = await predict_many(self.target, xs)
+        self._query_count += len(xs)
+        self._capture_samples(xs, preds)
+        return preds
+
+    def _capture_samples(
+        self, xs: t.Sequence[QueryInput], preds: t.Sequence[Prediction], limit: int = 12
+    ) -> None:
+        for x, p in zip(xs, preds, strict=False):
+            if len(self._query_samples) >= limit:
+                return
+            vec = p.vector
+            conf = float(max(vec)) if vec is not None and len(vec) else None
+            self._query_samples.append(
+                {
+                    "input": sample_input_preview(x),
+                    "label": p.hard_label,
+                    "confidence": round(conf, 4) if conf is not None else None,
+                }
+            )
+
+    # -- budget ------------------------------------------------------------ #
+    def _budget_left(self) -> int:
+        return max(0, self.max_queries - self._query_count)
+
+    def _over_budget(self) -> bool:
+        return self._query_count >= self.max_queries
+
+    # -- tracing ----------------------------------------------------------- #
+    def _phase(self, name: str, **attrs: t.Any) -> t.Any:
+        """A coarse task-span phase (e.g. 'perturb + probe target')."""
+        try:
+            import dreadnode as dn
+
+            return dn.task_span(name, type="task", attributes={str(k): v for k, v in attrs.items()})
+        except Exception:
+            return nullcontext()
+
+    def _trace_step(
+        self,
+        iteration: int,
+        *,
+        input: t.Any = None,
+        output: t.Any = None,
+        metrics: dict[str, float] | None = None,
+        **attrs: t.Any,
+    ) -> None:
+        """Emit a point-in-time span for one algorithm iteration, recording the
+        intermediate adversarial state so the full attack trajectory is
+        inspectable in the Traces tab.
+
+        ``input`` is the current candidate (perturbed text / feature vector /
+        image preview) the attacker sent; ``output`` is the target's response at
+        this step (predicted label + confidence); ``metrics`` are the scalar
+        signals (distance, confidence, fidelity) plotted over the run. Capped
+        (first 20 iterations, then every 10th) so a long run does not explode the
+        trace tree."""
+        if iteration >= 20 and iteration % 10 != 0:
+            return
+        with self._phase(f"step {iteration}", iteration=iteration, **attrs) as span:
+            if span is None or not hasattr(span, "log_input"):
+                return
+            with suppress(Exception):
+                if input is not None:
+                    span.log_input("candidate", input)
+                if output is not None:
+                    span.log_output("prediction", output)
+                for key, value in (metrics or {}).items():
+                    span.log_metric(key, float(value), step=iteration)
+
+    # -- run scaffold ------------------------------------------------------ #
+    @property
+    def _attack_name(self) -> str:
+        raise NotImplementedError
+
+    def _study_kwargs(self) -> dict[str, t.Any]:
+        """Extra ``study_span`` kwargs beyond the common set (e.g. distance norm)."""
+        return {}
+
+    def _span_attributes(self, result: t.Any) -> dict[str, t.Any]:
+        """Map the result onto the ``AIRT_ATTRIBUTE_*`` span attributes to set.
+        ``None`` values are skipped."""
+        raise NotImplementedError
+
+    async def _execute(self) -> t.Any:
+        raise NotImplementedError
+
+    async def run(self) -> t.Any:
+        import json
+
+        from dreadnode.tracing.constants import AIRT_ATTRIBUTE_METRICS_DETAIL
+        from dreadnode.tracing.spans import study_span
+
+        with study_span(
+            name=self._attack_name,
+            airt_assessment_id=self.airt_assessment_id,
+            airt_attack_name=self._attack_name,
+            airt_attack_domain=self.attack_domain,
+            airt_target_model=self.airt_target_model,
+            airt_input_modality=self.modality,
+            airt_goal=self.airt_goal,
+            airt_goal_category=self.airt_goal_category,
+            **self._study_kwargs(),
+        ) as span:
+            result = await self._execute()
+            if hasattr(result, "query_samples"):
+                result.query_samples = self._query_samples
+            for key, value in self._span_attributes(result).items():
+                if value is not None:
+                    span.set_attribute(key, value)
+            span.set_attribute(AIRT_ATTRIBUTE_METRICS_DETAIL, json.dumps(result.metrics_detail))
+        return result

dreadnode-2.0.36/dreadnode/airt/art_engine.py

@@ -0,0 +1,183 @@
+"""Optional ART (Adversarial Robustness Toolbox) engine for the traditional-ML
+privacy attacks.
+
+When ``adversarial-robustness-toolbox`` is installed (the ``airt-ml`` extra), the
+extraction and membership attacks route their core step through ART's reference
+implementations - CopycatCNN / KnockoffNets for extraction, the black-box
+membership inference attacks for membership - instead of the native
+numpy/sklearn engine. Everything here is import-guarded: if ART (or torch) is
+missing, :func:`art_available` returns ``False`` and callers fall back to the
+native engine, so the base SDK never hard-depends on ART.
+
+The victim is always wrapped as a black-box: ART only ever sees a
+``predict_fn`` backed by predictions we already collected from the target's API,
+so no extra queries are issued and the attack stays honestly black-box.
+"""
+
+import typing as t
+
+import numpy as np
+from loguru import logger
+
+
+def art_available() -> bool:
+    """True when the Adversarial Robustness Toolbox and torch import cleanly.
+
+    We check for ``art.attacks.extraction`` specifically, not just ``import art``:
+    the top-level ``art`` name is also owned by the ASCII-art package the SDK
+    depends on, so a bare ``import art`` would false-positive on that library.
+    ART is not a hard dependency (it collides with the ASCII-art ``art``); this
+    returns True only in an environment where the adversarial toolbox is actually
+    installed, otherwise the native engine runs.
+    """
+    try:
+        import torch  # noqa: F401
+        from art.attacks.extraction import CopycatCNN  # noqa: F401
+        from art.estimators.classification import BlackBoxClassifier  # noqa: F401
+    except Exception:
+        return False
+    return True
+
+
+def _cached_victim(
+    x: np.ndarray,
+    proba: np.ndarray,
+    nb_classes: int,
+) -> t.Any:
+    """A BlackBoxClassifier whose predict_fn replays already-collected target
+    probabilities, keyed by feature-row bytes - ART "queries" it without issuing
+    a single extra network call."""
+    from art.estimators.classification import BlackBoxClassifier
+
+    lookup = {row.tobytes(): proba[i] for i, row in enumerate(np.ascontiguousarray(x))}
+    uniform = np.full(nb_classes, 1.0 / nb_classes)
+
+    def predict_fn(batch: np.ndarray) -> np.ndarray:
+        rows = np.ascontiguousarray(batch, dtype=x.dtype)
+        return np.asarray([lookup.get(r.tobytes(), uniform) for r in rows])
+
+    return BlackBoxClassifier(
+        predict_fn=predict_fn,
+        input_shape=(x.shape[1],),
+        nb_classes=nb_classes,
+    )
+
+
+def _torch_thieved(n_features: int, nb_classes: int) -> t.Any:
+    """A small torch MLP wrapped as an ART PyTorchClassifier - the surrogate that
+    ART trains to replicate the victim."""
+    import torch
+    from art.estimators.classification import PyTorchClassifier
+    from torch import nn
+
+    torch.manual_seed(0)
+    model = nn.Sequential(
+        nn.Linear(n_features, 128),
+        nn.ReLU(),
+        nn.Linear(128, 64),
+        nn.ReLU(),
+        nn.Linear(64, nb_classes),
+    )
+    return PyTorchClassifier(
+        model=model,
+        loss=nn.CrossEntropyLoss(),
+        optimizer=torch.optim.Adam(model.parameters(), lr=1e-3),
+        input_shape=(n_features,),
+        nb_classes=nb_classes,
+        clip_values=(-1e9, 1e9),
+        device_type="cpu",
+    )
+
+
+def art_extract_surrogate(
+    strategy: str,
+    x: np.ndarray,
+    proba: np.ndarray,
+    nb_classes: int,
+    *,
+    nb_epochs: int = 30,
+) -> t.Any | None:
+    """Train a thieved classifier with ART's CopycatCNN (copycat) or KnockoffNets
+    (knockoff) against a cached black-box victim.
+
+    Returns the trained ART classifier (its ``.predict`` yields probabilities), or
+    ``None`` if ART is unavailable or the run fails - the caller then falls back to
+    the native surrogate. Only ``copycat`` and ``knockoff`` map onto ART; the other
+    strategies stay native.
+    """
+    if strategy not in ("copycat", "knockoff") or not art_available():
+        return None
+    try:
+        from art.attacks.extraction import CopycatCNN, KnockoffNets
+
+        x = np.ascontiguousarray(x, dtype=np.float32)
+        victim = _cached_victim(x, proba, nb_classes)
+        thieved = _torch_thieved(x.shape[1], nb_classes)
+        n = len(x)
+        if strategy == "copycat":
+            attack: t.Any = CopycatCNN(
+                classifier=victim,
+                batch_size_fit=32,
+                batch_size_query=32,
+                nb_epochs=nb_epochs,
+                nb_stolen=n,
+                use_probability=False,
+            )
+        else:  # knockoff - soft-label transfer set
+            attack = KnockoffNets(
+                classifier=victim,
+                batch_size_fit=32,
+                batch_size_query=32,
+                nb_epochs=nb_epochs,
+                nb_stolen=n,
+                sampling_strategy="random",
+                reward="all",
+                verbose=False,
+                use_probability=True,
+            )
+        stolen = attack.extract(x=x, thieved_classifier=thieved)
+    except Exception as e:  # pragma: no cover - defensive fallback
+        logger.warning("ART extraction ({}) failed, falling back to native: {}", strategy, e)
+        return None
+    else:
+        logger.debug("ART {} extraction trained thieved classifier on {} samples", strategy, n)
+        return stolen
+
+
+def art_membership_scores(
+    method: str,
+    members_x: np.ndarray,
+    nonmembers_x: np.ndarray,
+    member_proba: np.ndarray,
+    nonmember_proba: np.ndarray,
+    member_labels: np.ndarray,
+    nonmember_labels: np.ndarray,
+    nb_classes: int,
+) -> np.ndarray | None:
+    """Per-record membership scores via ART's black-box rule-based membership
+    inference (Yeom et al.). Returns an array aligned to
+    ``[members..., nonmembers...]`` in ``[0, 1]``, or ``None`` to fall back to
+    native. Only the ``threshold`` method maps onto ART's black-box rule-based
+    attack; ``label_only`` stays native.
+    """
+    if method != "threshold" or not art_available():
+        return None
+    try:
+        from art.attacks.inference.membership_inference import (
+            MembershipInferenceBlackBoxRuleBased,
+        )
+
+        x = np.ascontiguousarray(np.vstack([members_x, nonmembers_x]), dtype=np.float32)
+        proba = np.vstack([member_proba, nonmember_proba])
+        labels = np.concatenate([member_labels, nonmember_labels]).astype(int)
+        victim = _cached_victim(x, proba, nb_classes)
+        attack = MembershipInferenceBlackBoxRuleBased(classifier=victim)
+        scores = np.asarray(attack.infer(x, labels, probabilities=True), dtype=np.float64)
+        # infer(probabilities=True) returns (N, 2) = [P(non-member), P(member)];
+        # the member column is the per-record membership score.
+        if scores.ndim == 2 and scores.shape[1] == 2:
+            scores = scores[:, 1]
+        return scores.ravel()
+    except Exception as e:  # pragma: no cover - defensive fallback
+        logger.warning("ART membership (threshold) failed, falling back to native: {}", e)
+        return None

{dreadnode-2.0.35 → dreadnode-2.0.36}/dreadnode/airt/assessment.py

@@ -181,6 +181,12 @@ class Assessment:
     async def __aenter__(self) -> Assessment:
         self._context_token = _current_assessment.set(self)
         self._tracing_enabled = True
+        # Register on enter so the assessment id is available to any attack run
+        # inside the block - including the traditional-ML privacy attacks that
+        # call their own .run() (not assessment.run()). study_span picks up the
+        # active assessment via the _current_assessment ContextVar, so the user
+        # never has to call register()/complete() or pass airt_assessment_id.
+        await self._ensure_started()
         return self
 
     async def __aexit__(
@@ -296,7 +302,11 @@ class Assessment:
 
     async def _finalize(self) -> None:
         """Flush pending OTEL spans and mark complete."""
-        if not self._attack_results:
+        # Complete when there is something to finalize: the generative path
+        # populates _attack_results, while the privacy attacks emit study spans
+        # under a registered assessment (id set on __aenter__). Skip only when
+        # both are absent (offline / never used).
+        if not self._attack_results and self._assessment_id is None:
             return
 
         # Flush pending OTEL spans BEFORE marking complete so CH has data

{dreadnode-2.0.35 → dreadnode-2.0.36}/dreadnode/airt/crescendo.py

@@ -15,6 +15,7 @@ from dreadnode.core.meta import TrialCandidate
 from dreadnode.core.scorer import Scorer
 from dreadnode.core.transforms import Transform, TransformsLike
 from dreadnode.generators.generator import Generator, get_generator
+from dreadnode.generators.proxy import resolve_dn_model_to_generator
 from dreadnode.optimization import Study
 from dreadnode.optimization.stopping import score_value
 from dreadnode.samplers import iterative_sampler
@@ -122,12 +123,15 @@ def crescendo_attack(
     # Replace {goal} placeholder in the template
     system_prompt_template = variant_system_prompt.format(goal=goal)
 
-    # Resolve generator once for reuse across turns
+    # Resolve generator once for reuse across turns. `dn/*` ids route through the
+    # platform gateway (get_generator can't resolve them); non-dn strings pass
+    # through unchanged.
+    resolved_attacker = resolve_dn_model_to_generator(attacker_model)
     attacker_generator: Generator
-    if isinstance(attacker_model, Generator):
-        attacker_generator = attacker_model
+    if isinstance(resolved_attacker, Generator):
+        attacker_generator = resolved_attacker
     else:
-        attacker_generator = get_generator(attacker_model)
+        attacker_generator = get_generator(resolved_attacker)
 
     # Create the refiner that generates the next prompt in the conversation
     async def crescendo_refiner(trials: "list[Trial[str]]") -> str: