Skip to content

feat: State persistence and recovery for agent runs across pod restarts #150

Description

@saharannaveen

Problem

When an agent pod is killed mid-conversation (due to rolling updates, OOM, or node eviction), the in-flight run is lost. The user sees a frozen chat with no way to resume. There is no mechanism to track active runs, persist their state, or recover them on a replacement pod.

Expected Behavior

When a pod is killed mid-run:

  1. The shutdown hook should mark the active run as interrupted in Redis + Postgres with the latest checkpoint_id
  2. A replacement pod should claim the interrupted run on startup and resume from the last completed graph node
  3. The UI should detect the resumed run completing and render the response seamlessly

Proposed Solution

Add a lifecycle state persistence engine (lifecycle.py) that:

  • Registers in-flight runs in Redis with TTL-based leases
  • Updates checkpoint progress on each graph node completion
  • Persists state on shutdown (SIGTERM) — marks all local in-flight runs as interrupted
  • Resumes on startup — claims interrupted runs using FOR UPDATE SKIP LOCKED to prevent duplicate processing across replicas, then re-enqueues through the Aegra worker executor
  • Encrypts tokens via Fernet (derived from SSO_CLIENT_SECRET) for sensitive data stored in Redis

Files to modify

  • startup.py — call resume_interrupted_runs on pod startup
  • shutdown.py — call persist_inflight_runs during graceful shutdown
  • graph.py — integrate lifecycle tracking (register, update, deregister)

Acceptance Criteria

  • In-flight runs are tracked in Redis with TTL-based leases
  • Graceful shutdown persists all active run state to Redis + Postgres
  • New pod resumes interrupted runs from the last checkpoint on startup
  • No duplicate processing across replicas (lease-based claiming with FOR UPDATE SKIP LOCKED)
  • Sensitive tokens are encrypted at rest in Redis
  • Unit and integration tests cover the lifecycle module

References

  • Internal tracker: RHITAIF-206

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions