Skip to content

feat(checkpoint): add indirect prompt injection state recovery - #3548

Open
zyzhou5 wants to merge 4 commits into
NVIDIA-NeMo:amahishi/gym-prefix-cuts-lineagefrom
zyzhou5:zezhou/workplace-checkpoint-state
Open

feat(checkpoint): add indirect prompt injection state recovery#3548
zyzhou5 wants to merge 4 commits into
NVIDIA-NeMo:amahishi/gym-prefix-cuts-lineagefrom
zyzhou5:zezhou/workplace-checkpoint-state

Conversation

@zyzhou5

@zyzhou5 zyzhou5 commented Sep 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What does this PR do?

Indirect Prompt Injection keeps tool-mutated environment state in memory. Add a resource checkpoint adapter so a replacement rollout attempt can restore that state and continue without reseeding the environment.

  • Bind sessions to rollout/attempt identity while preserving cookie-based lookup for legacy requests.
  • Export versioned snapshots of the complete nested environment dictionary, with JSON validation and independent copies.
  • Validate the entire restore batch before activation; reject duplicate identities and malformed state.
  • Remove environment state and its execution binding during verification or explicit retirement.

Request receipts, revisions, and stale-attempt fencing use the shared checkpoint stack. Based on #3412 (amahishi/gym-prefix-cuts-lineage, which includes #3349), at 90ff732e4cf4ca4e41c840d2ed14867baa2b3524, matching the base used by #3634. This PR adds only IPI support and tests. Related to #3240.

Validation

Add 17 focused cases in resources_servers/indirect_prompt_injection/tests/test_checkpoint.py:

  • Exercise the real resource prepare/commit/restore/resume endpoints and checkpoint files at three boundaries: after seeding, after a chart update, and after sending a message.
  • Restore into a fresh server under the replacement attempt, then compare continuation tool outputs, revisions, final environment state, and rewards with the original server. Check receipt replay, stale-attempt rejection, and continuation without the original session cookie.
  • Check nested snapshot isolation, successful batch replacement, atomic rejection of invalid/duplicate batches, JSON validation, execution lookup, verifier-error cleanup, and idempotent retirement.

Validation after syncing to the current base: 356 tests passed across the tracked IPI suite and shared resource/agent checkpoint tests. The unchanged IPI application module reached 99.50% coverage in the earlier 309-test run. Repository-wide pre-commit run --all-files and git diff --check passed. The final diff contains only the IPI adapter and its checkpoint tests.

Rollouts

The unchanged production adapter was also exercised in a manual mid-training recovery run on 2026-09-20 with Qwen3-4B-Instruct-2507 (jobs 7306846 and 7307095). A separate job restored the checkpoint saved after training step 1; both jobs completed step 3 with matching selected-rollout state, continuation tool outputs, and rewards. Tested Gym 05953c7 and RL df109a0 with local experiment instrumentation and an RL metric-transport fix. This covered one completed-tool boundary; RFC-wrapper integration and token-level interruptions were not exercised.

@copy-pr-bot

copy-pr-bot Bot commented Sep 18, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@zyzhou5
zyzhou5 changed the base branch from main to amahishi/gym-turn-level-recovery September 18, 2026 22:59
@yaoyu-33 yaoyu-33 added area:environment Individual environments, benchmarks, verifiers, and environment-specific resources servers complexity:medium Single-domain change with interacting parts or a moderate review surface feature New capabilities, enhancements, or enablement work needs-review PR is ready for code review and waiting on a reviewer labels Sep 18, 2026
@zyzhou5
zyzhou5 force-pushed the zezhou/workplace-checkpoint-state branch from 3342943 to 05953c7 Compare September 19, 2026 22:44
@zyzhou5 zyzhou5 changed the title feat: workplace checkpoint state feat(checkpoint): add indirect prompt injection state recovery Sep 19, 2026
@github-actions github-actions Bot added the sla:triage-overdue Review assignment is over the one-business-day SLA label Sep 21, 2026
@github-actions github-actions Bot mentioned this pull request Sep 21, 2026
Open
@macandro96
macandro96 force-pushed the amahishi/gym-turn-level-recovery branch from 5b0b444 to f4fcf8c Compare September 23, 2026 00:36
@zyzhou5
zyzhou5 force-pushed the zezhou/workplace-checkpoint-state branch from e44fb49 to c259973 Compare September 23, 2026 23:35
@zyzhou5
zyzhou5 changed the base branch from amahishi/gym-turn-level-recovery to amahishi/gym-prefix-cuts-lineage September 23, 2026 23:36
@macandro96
macandro96 force-pushed the amahishi/gym-prefix-cuts-lineage branch from 8a51b31 to 90ff732 Compare September 28, 2026 20:58
@zyzhou5
zyzhou5 force-pushed the zezhou/workplace-checkpoint-state branch from c259973 to 7de74a8 Compare September 28, 2026 21:31
@macandro96
macandro96 force-pushed the amahishi/gym-prefix-cuts-lineage branch from 90ff732 to 76747e0 Compare September 29, 2026 04:40
@zyzhou5
feat(checkpoint): restore indirect prompt injection state
1b8314d 
Add versioned environment snapshots and rollout-attempt session bindings for the indirect prompt injection resources server. Validate restore batches before activation and retire restored state through the existing checkpoint lifecycle.

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5
test(checkpoint): cover indirect prompt injection recovery
319067b 
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5
zyzhou5 force-pushed the zezhou/workplace-checkpoint-state branch from 7de74a8 to 319067b Compare September 29, 2026 16:50
@zyzhou5
fix(checkpoint): preserve IPI cookie session access
2089379 
Use the fresh session cookie as the resource state key and bind rollout identity to it. Preserve execution isolation when cookies are reused, keep reseeding on the existing binding, and clean up cookie-only verification when checkpointing is disabled. Require execution identity before completing a checkpoint-managed session.

Add HTTP regressions for mixed identity and cookie access, reused sessions, reseeding, cleanup, and verification fencing. Validation: 368 tests passed, IPI application coverage 99.52%, and all-files pre-commit passed.

Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@zyzhou5
refactor(checkpoint): simplify IPI state restore
0e6ee8d 
Signed-off-by: Zeyu Zhou <zezhou@nvidia.com>
@ananthsub ananthsub mentioned this pull request Oct 1, 2026
Draft
@zyzhou5 zyzhou5 mentioned this pull request Oct 1, 2026
Draft
7 tasks done
ananthsub added a commit that referenced this pull request Oct 1, 2026
@ananthsub
feat(checkpoint): export indirect prompt injection session state
d88d2c2 
The indirect prompt injection resources server keeps a mutable environment per
cookie session, so it now declares checkpoint_mode "exported" and implements
the three session hooks: export each live session's environment, restore a
batch only after every state validates, and drop a retired session.
Verification keeps the default "wait" mode because it deletes the session.

Export leaves out a session the server already dropped, such as one whose
verification raised, so a later commit does not fail on it. The environment is
copied through JSON, which keeps its key order: tool outputs serialize the
environment as stored, so a restored session answers byte for byte as the
original would.

Ported from #3548 onto the v2 checkpoint hooks. The execution
identity map, header identity lookups, mutation receipts, revisions, and the
verify identity check are not carried: the v2 participant keys sessions by
cookie session ID and relies on ordering instead of receipts.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub added a commit that referenced this pull request Oct 1, 2026
@ananthsub
feat(checkpoint): export indirect prompt injection session state
e297cd8 
The indirect prompt injection resources server keeps a mutable environment per
cookie session, so it now declares checkpoint_mode "exported" and implements
the three session hooks: export each live session's environment, restore a
batch only after every state validates, and drop a retired session.
Verification keeps the default "wait" mode because it deletes the session.

Export leaves out a session the server already dropped, such as one whose
verification raised, so a later commit does not fail on it. The environment is
copied through JSON, which keeps its key order: tool outputs serialize the
environment as stored, so a restored session answers byte for byte as the
original would.

Ported from #3548 onto the v2 checkpoint hooks. The execution
identity map, header identity lookups, mutation receipts, revisions, and the
verify identity check are not carried: the v2 participant keys sessions by
cookie session ID and relies on ordering instead of receipts.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub added a commit that referenced this pull request Oct 2, 2026
@ananthsub
feat(checkpoint): export indirect prompt injection session state
3bd2b9c 
The indirect prompt injection resources server keeps a mutable environment per
cookie session, so it now declares checkpoint_mode "exported" and implements
the three session hooks: export each live session's environment, restore a
batch only after every state validates, and drop a retired session.
Verification keeps the default "wait" mode because it deletes the session.

Export leaves out a session the server already dropped, such as one whose
verification raised, so a later commit does not fail on it. The environment is
copied through JSON, which keeps its key order: tool outputs serialize the
environment as stored, so a restored session answers byte for byte as the
original would.

Ported from #3548 onto the v2 checkpoint hooks. The execution
identity map, header identity lookups, mutation receipts, revisions, and the
verify identity check are not carried: the v2 participant keys sessions by
cookie session ID and relies on ordering instead of receipts.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub added a commit that referenced this pull request Oct 2, 2026
@ananthsub
feat(checkpoint): export indirect prompt injection session state
2304606 
The indirect prompt injection resources server keeps a mutable environment per
cookie session, so it now declares checkpoint_mode "exported" and implements
the three session hooks: export each live session's environment, restore a
batch only after every state validates, and drop a retired session.
Verification keeps the default "wait" mode because it deletes the session.

Export leaves out a session the server already dropped, such as one whose
verification raised, so a later commit does not fail on it. The environment is
copied through JSON, which keeps its key order: tool outputs serialize the
environment as stored, so a restored session answers byte for byte as the
original would.

Ported from #3548 onto the v2 checkpoint hooks. The execution
identity map, header identity lookups, mutation receipts, revisions, and the
verify identity check are not carried: the v2 participant keys sessions by
cookie session ID and relies on ordering instead of receipts.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub added a commit that referenced this pull request Oct 2, 2026
@ananthsub
feat(checkpoint): export indirect prompt injection session state
41588ad 
The indirect prompt injection resources server keeps a mutable environment per
cookie session, so it now declares checkpoint_mode "exported" and implements
the three session hooks: export each live session's environment, restore a
batch only after every state validates, and drop a retired session.
Verification keeps the default "wait" mode because it deletes the session.

Export leaves out a session the server already dropped, such as one whose
verification raised, so a later commit does not fail on it. The environment is
copied through JSON, which keeps its key order: tool outputs serialize the
environment as stored, so a restored session answers byte for byte as the
original would.

Ported from #3548 onto the v2 checkpoint hooks. The execution
identity map, header identity lookups, mutation receipts, revisions, and the
verify identity check are not carried: the v2 participant keys sessions by
cookie session ID and relies on ordering instead of receipts.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

This branch has not been deployed

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

Reviewers

No reviews

Assignees

No one assigned

Labels

area:environment Individual environments, benchmarks, verifiers, and environment-specific resources servers complexity:medium Single-domain change with interacting parts or a moderate review surface feature New capabilities, enhancements, or enablement work needs-review PR is ready for code review and waiting on a reviewer sla:triage-overdue Review assignment is over the one-business-day SLA

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@zyzhou5 @yaoyu-33