FEAT: Add reset_conversation_async hook to PromptTarget - #2320
FEAT: Add reset_conversation_async hook to PromptTarget#2320Mahdi Alhakim (mahdi-al-hakim) wants to merge 1 commit into
Conversation
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
|
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? |
|
(GHCP Generated): I think the underlying problem is real, especially for scenarios. A scenario commonly shares one I do not think
I suggest this design:
This keeps the useful part of the proposal for shared realtime targets, while preserving |
Description
Implements item 3 from #1247, which Roman Lutz (@romanlutz) scoped on 2026-07-03 as:
Today a target that holds state per conversation has no standard place to release it.
RealtimeTargetgrew its owncleanup_conversation_async, nothing calls it, and the issue reporter had to monkey-patch_teardown_asyncto get it invoked.What this adds
PromptTarget.reset_conversation_async(*, conversation_id), a no-op on the base so stateless targets are unaffected.AttackStrategy._teardown_asyncthat hands the target every objective-target conversation the run used.RealtimeTargetimplements the hook.cleanup_conversation_asyncnow 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:
PromptSendingAttackmints a fresh id per retry (prompt_sending.py:154), somax_attempts_on_failure=3leaves three behind.CrescendoAttackmints one per backtrack (crescendo.py:817).MultiTurnAttackStrategy._rotate_conversation_for_single_turn_targetmints 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_conversationsasConversationType.PRUNED, so_get_objective_conversation_idsreturns the live conversation plus those.TAP is the exception.
TAPAttackContextinheritssessionbut never uses it, and keeps one objective conversation per tree node plusbest_conversation_id. The base lookup would have reset asession.conversation_idthat is never sent anywhere while leaking every node, soTreeOfAttacksWithPruningAttackoverrides the lookup.Teardown must not swallow the real error
_teardown_asyncruns in thefinallyof_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-typesPromptTargethits 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_asyncthat did nothing, five saying so in a comment and three in the docstring. Now that the base does something, each would have neededawait 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.TreeOfAttacksWithPruningAttackkeeps a real override of the id lookup rather than of teardown.BargeInAttack's override noted that its session closes its own connection inrun_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_asyncandcleanup_target_asynctogether, with naming to be reconciled.cleanup_conversation_asyncis the same concept as the new hook, so it is deprecated withremoved_in="1.3.0"(current version plus two minors, per the style guide). It has no callers anywhere inpyrit/,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_asyncis left alone. Closing the whole target is a different lifetime from releasing one conversation, and it is not declared onPromptTargetat all today, only onRealtimeTargetand as a stub onTextTarget. 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_idper turn against a single-turn target (tree_of_attacks.py:562and:623) without recording the abandoned id asPRUNED, unlike_rotate_conversation_for_single_turn_targetwhich 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 asPRUNEDthere 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_asyncstill 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_nooptests now assert the reset happens, and the threeRealtimeTargetconversation-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 unusedsession.conversation_idexcluded, 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 thePromptTargetcontract 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
Integration tests were not run because they need credentials.