Skip to content

FEAT: Add reset_conversation_async hook to PromptTarget - #2320

Open
Mahdi Alhakim (mahdi-al-hakim) wants to merge 1 commit into
microsoft:mainfrom
mahdi-al-hakim:feat/reset-conversation-hook
Open

FEAT: Add reset_conversation_async hook to PromptTarget#2320
Mahdi Alhakim (mahdi-al-hakim) wants to merge 1 commit into
microsoft:mainfrom
mahdi-al-hakim:feat/reset-conversation-hook

Conversation

@mahdi-al-hakim

Copy link
Copy Markdown

Description

Implements item 3 from #1247, which Roman Lutz (@romanlutz) scoped on 2026-07-03 as:

a standardized conversation-reset hook: a no-op reset_conversation_async(*, conversation_id=...) on base PromptTarget, invoked from attack _teardown_async, that targets holding external state (Playwright page, HTTP/websocket conversation_id) override.

Today a target that holds state per conversation has no standard place to release it. RealtimeTarget grew its own cleanup_conversation_async, nothing calls it, and the issue reporter had to monkey-patch _teardown_async to get it invoked.

What this adds

  • PromptTarget.reset_conversation_async(*, conversation_id), a no-op on the base so stateless targets are unaffected.
  • A base implementation of AttackStrategy._teardown_async that hands the target every objective-target conversation the run used.
  • RealtimeTarget implements the hook. cleanup_conversation_async now delegates to it and warns.

Every conversation, not just the last one

The obvious version of this resets one id and leaks the rest. Three things in-tree rotate the objective-target conversation mid-run:

  • PromptSendingAttack mints a fresh id per retry (prompt_sending.py:154), so max_attempts_on_failure=3 leaves three behind.
  • CrescendoAttack mints one per backtrack (crescendo.py:817).
  • MultiTurnAttackStrategy._rotate_conversation_for_single_turn_target mints one per turn against single-turn targets, so an 8-turn run leaves seven.

In all three the superseded ids are already recorded on context.related_conversations as ConversationType.PRUNED, so _get_objective_conversation_ids returns the live conversation plus those.

TAP is the exception. TAPAttackContext inherits session but never uses it, and keeps one objective conversation per tree node plus best_conversation_id. The base lookup would have reset a session.conversation_id that is never sent anywhere while leaking every node, so TreeOfAttacksWithPruningAttack overrides the lookup.

Teardown must not swallow the real error

_teardown_async runs in the finally of _execution_context_async, so an override that raises replaces whatever error the attack was already reporting. Closing a websocket during shutdown is exactly the kind of thing that raises, so the reset is wrapped and logged. A target that only duck-types PromptTarget hits the same guard, so it logs once per conversation instead of breaking the run.

Removed the no-op teardown overrides

Eight attacks had a _teardown_async that did nothing, five saying so in a comment and three in the docstring. Now that the base does something, each would have needed await super()._teardown_async(...), and forgetting that line would silently disable the reset with a green test suite. They are deleted so the base is the single source of truth. TreeOfAttacksWithPruningAttack keeps a real override of the id lookup rather than of teardown. BargeInAttack's override noted that its session closes its own connection in run_async; the base reset is still right there because it is keyed by conversation id and an unknown id is a no-op.

On the naming reconciliation

#1247 mentions cleanup_conversation_async and cleanup_target_async together, with naming to be reconciled.

cleanup_conversation_async is the same concept as the new hook, so it is deprecated with removed_in="1.3.0" (current version plus two minors, per the style guide). It has no callers anywhere in pyrit/, doc/, or the scenarios, so the warning cannot fire during a normal run, and the signature is unchanged so any external caller still works.

cleanup_target_async is left alone. Closing the whole target is a different lifetime from releasing one conversation, and it is not declared on PromptTarget at all today, only on RealtimeTarget and as a stub on TextTarget. Hoisting it onto the base is a real cleanup but it is a separate change, and I did not want to widen this one. Happy to do it next if you want them fully unified.

One case this does not cover

A TAP node re-mints its own objective_target_conversation_id per turn against a single-turn target (tree_of_attacks.py:562 and :623) without recording the abandoned id as PRUNED, unlike _rotate_conversation_for_single_turn_target which does. Those ids are therefore not reset. It is pre-existing and currently unreachable, since it needs a target that is both single-turn and holds per-conversation state, and no shipped target is both. Recording them as PRUNED there would close it, but that is a change to TAP's bookkeeping rather than to this hook, so I left it out. Happy to add it if you would rather have it in one go.

Not a breaking change

Nothing is removed or renamed from the public API. The new method has a no-op base implementation, cleanup_conversation_async still works, and no target is required to implement anything. The behavior change is that stateful targets now get told when a conversation is finished, which is the point of the issue.

Tests and Documentation

12 new tests. No existing test was removed. Six existing tests changed: three test_teardown_async_is_noop tests now assert the reset happens, and the three RealtimeTarget conversation-cleanup tests were renamed to the new method.

  • tests/unit/executor/attack/core/test_attack_strategy.py: single-turn and multi-turn id lookup against real contexts, PRUNED conversations included, non-PRUNED references ignored, no id on the base context, and a raising target being swallowed (6 tests).
  • tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py: node conversations collected, best and pruned included, the unused session.conversation_id excluded, no duplicates (4 tests).
  • tests/unit/prompt_target/test_text_target.py: a stateless target inherits the no-op (1 test).
  • tests/unit/prompt_target/target/test_realtime_target.py: the hook closes and removes the connection, swallows close errors, is a no-op for an unknown id, and the deprecated alias warns and delegates (3 renamed, 1 new).

Documentation: added a "Releasing per-conversation state" section to doc/code/targets/0_prompt_targets.md, which is where the PromptTarget contract is explained to target authors, and a matching one to .github/instructions/targets.instructions.md. Both are .md, so there is nothing to run through JupyText.

Checks run locally

uv run pre-commit run --all-files
  18 hooks, all Passed (ruff format, ruff check, ty, async-suffix, no-rest-roles, ...)

uv run pytest tests/unit -q
  14807 passed, 6 skipped

Integration tests were not run because they need credentials.

Targets that hold external state keyed by conversation had no standard way
to release it when an attack finished. RealtimeTarget grew its own
cleanup_conversation_async, and the issue reporter had to monkey-patch
_teardown_async to call it.

PromptTarget gains a no-op reset_conversation_async(*, conversation_id),
and AttackStrategy._teardown_async now hands it every objective-target
conversation the run used. That is the live conversation plus the ones
recorded as PRUNED, since a PromptSendingAttack retry, a Crescendo
backtrack, and the single-turn rotation in multi-turn attacks all leave
earlier conversations behind. TAP keys conversations per tree node instead,
so it overrides the lookup.

The reset runs in the lifecycle finally block, so a target that raises is
logged rather than replacing whatever error the attack was reporting.

RealtimeTarget now implements the hook, and cleanup_conversation_async
delegates to it with a deprecation warning. cleanup_target_async is left
alone since closing the whole target is a different concern.

Towards microsoft#1247
@mahdi-al-hakim

Copy link
Copy Markdown
Author

Hi Richard Lundeen (@rlundeen2) Roman Lutz (@romanlutz), this has been open for a few weeks now. It is a self-contained change (a reset_conversation_async hook on PromptTarget). I recently landed microsoft/RAMPART#141, so I am glad to match whatever conventions you prefer here. Is there anything I can do to help move it along, or someone I should tag for review?

@rlundeen2

Copy link
Copy Markdown
Contributor

(GHCP Generated): I think the underlying problem is real, especially for scenarios. A scenario commonly shares one RealtimeTarget across many concurrent attacks. Conversation IDs prevent state from mixing, but RealtimeTarget caches one websocket per ID. Without conversation-level cleanup, a large or long-running scenario can accumulate open connections until cleanup_target_async() runs.

I do not think AttackStrategy._teardown_async should discover conversation IDs from attack contexts, though. That introduces several concerns:

  • AttackResult already defines the canonical relationship through conversation_id, related_conversations, and get_active_conversation_ids(). _get_objective_conversation_ids() duplicates this contract with getattr checks and attack-specific overrides.
  • TAP demonstrates the fragility: it needs a custom lookup, and its single-turn path replaces a node's conversation ID each turn without recording the old ID as PRUNED, so those earlier resources still leak.
  • The current implementation handles only the objective target. Attacks can also use adversarial targets, scorer targets, and converter targets. Some of their IDs are already stored in related_conversations, but the teardown filters them out. This makes a generic attack-level cleanup guarantee incomplete.
  • Custom attacks can return a valid AttackResult while storing transient IDs in a different context shape. They will silently skip cleanup unless they know to override the new protected method.

I suggest this design:

  1. Keep a no-op conversation-level cleanup hook on PromptTarget, with RealtimeTarget implementing it. This allows a shared target to release one conversation without destroying the target or its other active conversations.
  2. Invoke objective-target cleanup at the AttackExecutor boundary after each run. The executor owns per-run execution, has the objective target, and receives the completed AttackResult.
  3. Use result.get_active_conversation_ids() instead of reconstructing IDs from context shape. Ensure error results also contain the real multi-turn session.conversation_id and all related objective conversations before cleanup runs.
  4. Make the scope explicit: this pass cleans objective-target conversations only. Define adversarial, scorer, and converter target lifecycles separately rather than implying that attack teardown handles every target.
  5. Fix TAP and other rotation paths so every replaced objective conversation is added to related_conversations; otherwise neither result-based nor context-based cleanup can release it.

This keeps the useful part of the proposal for shared realtime targets, while preserving AttackResult as the source of truth and avoiding target-lifecycle policy in each attack implementation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants