feat(checkpoint): export indirect prompt injection session state

[ananthsub]

What changed and why

The indirect prompt injection resources server keeps a mutable environment for each session: tool calls update patient charts, send messages, and so on. Before this change it was restart-only by default, so every live session blocked a checkpoint until its rollout was retired and restarted from its input. Now a rollout interrupted by a crash can continue from its last checkpoint.

• The server declares checkpoint_mode = "exported" and implements the three session hooks, keyed by the cookie session ID:
  • export_session_states returns each live session's environment as an IPICheckpointState;
  • restore_session_states validates every state first, then installs them all, so an invalid state installs nothing;
  • retire_session_state drops the session.
• checkpoint_verify stays at the default wait, because /verify deletes the session's environment and cannot be replayed.
• Export leaves out a session the server no longer holds. /verify drops the session in its finally block even when verification raises, and the participant only stops tracking a session on a successful /verify. Without this, the next commit after a failed verification would fail.
• The environment is copied through JSON, which keeps its key order and rejects values JSON cannot carry. Tool outputs serialize the environment as stored, so a restored session answers byte for byte as the original would; a test compares every tool output and the verify reply against an uncheckpointed run.

How it works

Where this PR sits in the overall flow

The highlighted part is what this PR adds.

flowchart LR
  C["Controller<br/>NeMo RL, or rollout collection"]
  CO["Coordination<br/>prepare, commit, restore, resume, retire"]
  K["Control plane on every server<br/>phases, fencing, lease, storage"]
  subgraph G["One participant per Gym server"]
    E["Environment server<br/>episode steps"]
    M["Policy model<br/>held responses, generation cuts"]
    A["Agent<br/>sessions parked at boundaries"]
    R["Resources server<br/>session state"]
  end
  W["Inference worker<br/>stages cut prefixes"]
  D[("Checkpoint directory<br/>records, then manifest")]
  C --> CO --> K
  K --> E & M & A & R
  M --> W
  G --> D
  classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937
  class R this

Loading

One checkpoint, a crash, and the restore, end to end:

sequenceDiagram
  participant C as Controller
  participant G as Gym participants
  participant D as Checkpoint directory
  C->>G: prepare, in order environment, model, agent, resources
  Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held
  G-->>C: prepared, or blockers at the deadline
  C->>G: commit with the episodes the controller continues
  G->>D: each participant writes its records, then its manifest
  C->>C: publish the checkpoint with the controller's own state
  C->>G: resume, in order resources, agent, model, environment
  Note over C,G: crash - every Gym process dies
  C->>G: restore the checkpoint in fresh processes, all or nothing
  D-->>G: records installed under attempt + 1
  C->>G: resume
  C->>G: /run as attempt + 1 continues each episode from its boundary
  G-->>C: a late call from attempt 0 gets 409 stale_attempt

Loading

This PR

Each session's environment, saved at a checkpoint. It is copied through JSON, which keeps its key order, so a restored session answers byte for byte as the original would.

sequenceDiagram
  participant A as Agent
  participant I as Indirect prompt injection server
  A->>I: /seed_session with the environment
  A->>I: tool calls read and change the environment
  Note over I: commit - export_session_states saves each environment
  Note over A,I: crash, then restore into a fresh server
  Note over I: restore_session_states validates every state, then installs them all
  A->>I: tool calls continue on the same cookie session
  A->>I: /verify scores and drops the session, so verify mode stays wait

Loading

Where this sits in the stack

This PR builds on the partial-rollout checkpointing stack: #3882 (core) through #3889 (end-to-end suite), plus #3893 (rollout collection). It is one of five PRs that port the environments the old stack checkpointed onto the new hooks. Each one is based on #3889 and can be reviewed on its own, except Blackjack, which builds on the Gymnasium fix:

• feat(checkpoint): checkpoint Workplace Assistant sessions #3894 (ananthsub/partial-ckpt-workplace-assistant): feat(checkpoint): checkpoint Workplace Assistant sessions
• fix(gymnasium): build Gymnasium servers on the shared resources server setup #3895 (ananthsub/partial-ckpt-gymnasium): fix(gymnasium): build Gymnasium servers on the shared resources server setup
• feat(checkpoint): export Blackjack game state for partial-rollout checkpoints #3896 (ananthsub/partial-ckpt-blackjack): feat(checkpoint): export Blackjack game state for partial-rollout checkpoints (on the Gymnasium PR)
• feat(checkpoint): export indirect prompt injection session state #3897 (ananthsub/partial-ckpt-indirect-prompt-injection): feat(checkpoint): export indirect prompt injection session state (this PR)
• feat(checkpoint): continue proof refinement agent turns from their boundaries #3898 (ananthsub/partial-ckpt-proof-refinement): feat(checkpoint): continue proof refinement agent turns from their boundaries

Relationship to the old stack

This supersedes #3548, which added the same recovery on the old (v1) checkpoint stack. Ported from #3548 onto the v2 checkpoint hooks.

Carried over:

• the strict, versioned checkpoint state model;
• validate-then-install restore;
• the restore-and-continue test parametrized over the checkpoint point (after seed, after a chart update, after a message), plus the aliasing, invalid-state, and non-JSON export tests.

Not carried, because the v2 stack removes the mechanisms they served:

• the execution_to_session map and the identity lookups from rollout headers, since v2 keys resources state by the cookie session ID;
• mutation receipts, revisions, and the receipt-replay test, since v2 relies on ordering (a tool call runs between agent boundaries, so it blocks prepare);
• the verify identity check and the explicit terminal marking, since the participant learns a session ended from a successful /verify;
• the duplicate-identity and stale-attempt-identity tests. A restore batch is keyed by session ID, so it cannot hold duplicates, and the participant fences retired sessions itself.

Issue

No separate issue. This is part of the partial-rollout checkpointing work, as the port of draft #3548.

Validation

• pytest -q resources_servers/indirect_prompt_injection/tests: 282 passed, including 16 new tests in test_checkpoint.py. They drive the server through its real /ng-control/v1/checkpoint/* routes over ASGI:
  • a seeded session is checkpointed after 0, 1, or 2 tool calls, restored into a fresh server instance, and continued through the remaining tool calls and /verify. Tool outputs and the verify reply match an uncheckpointed run exactly. After verification, the restored server and its participant hold no session;
  • an invalid checkpoint fails restore through the control route and installs nothing;
  • restore with any invalid state in the batch installs nothing;
  • a retired session is gone from the server, and its cookie is refused with stale_attempt while another session keeps working;
  • a session dropped by a failed verification is left out of the next commit, and that commit succeeds;
  • export and restore do not alias live state, and export rejects values JSON cannot carry.
• pytest -q tests/unit_tests/test_checkpoint_*.py (with RAY_TMPDIR=/tmp): 115 passed.
• ruff check and ruff format on the server: clean.
• pre-commit run --files on the changed files: passed.

Rollout evidence

Pending. No model rollout has been run yet. A smoke run should start the indirect prompt injection environment with checkpointing on, take a checkpoint mid-episode, restart the resources server from that checkpoint, and confirm the rollouts continue and their rewards match uncheckpointed runs. End-to-end use also needs the agent lane.

Compatibility

• Nothing changes when checkpointing is off. The hooks only run when the checkpoint participant is installed.
• With checkpointing on, live sessions of this server no longer block prepare. They are exported and restored instead of restarting their rollouts.
• The stored state format (schema_version: 1, the environment as a JSON object) is new and does not read checkpoints written by feat(checkpoint): add indirect prompt injection state recovery #3548.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.