feat: State persistence recovery fix

[saharannaveen]

Summary

Descrition
When an agent pod is killed mid-conversation (rolling update, OOM, node eviction), the in-flight run is lost — the user sees
a frozen chat with no response, and the conversation cannot be resumed. There is no mechanism to track which runs were
active, persist their state, or recover them on a replacement pod.

Changes

template-agent (14 commits, 2680 lines added)

New module: deep_agent/aegra/lifecycle.py — Lifecycle state persistence engine that tracks inflight runs in Redis and
recovers them after pod restarts:

• register_inflight — records run metadata (thread_id, checkpoint_id, user tokens) in Redis when a run starts, with a
  TTL-based lease
• update_inflight_checkpoint — updates checkpoint progress on each graph node completion
• deregister_inflight — removes the Redis key when a run finishes or fails
• persist_inflight_runs — called during graceful shutdown (SIGTERM), marks all local inflight runs as interrupted in both
  Redis and Postgres
• resume_interrupted_runs — called during startup, picks up interrupted runs using FOR UPDATE SKIP LOCKED (prevents
  duplicate processing across replicas), claims them via lease, and re-enqueues through the Aegra worker executor
• Token encryption via Fernet (derived from SSO_CLIENT_SECRET) for sensitive SSO tokens stored in Redis inflight records

Modified: deep_agent/aegra/startup.py — Calls resume_interrupted_runs on pod startup to recover any runs left behind by a
killed predecessor

Modified: deep_agent/aegra/shutdown.py — Calls persist_inflight_runs during graceful shutdown to mark active runs as
interrupted before the pod terminates

Modified: deep_agent/aegra/graph.py — Integrates lifecycle tracking into the agent graph: registers inflight at run start,
updates checkpoint on node completion, deregisters on completion/failure

Testing: 641-line unit test suite (test_lifecycle.py) + 318-line integration tests + 3 shell scripts for Kind cluster
pod-kill testing (kind-pod-kill-test.sh, pod-kill-and-recover.sh, inspect-db.sh)

User Flow

1. User sends a message in the chat UI — the agent starts processing.
2. The agent pod is killed mid-run (rolling update, OOM, kubectl delete pod).
3. Shutdown hook fires: persist_inflight_runs() marks the active run as interrupted in Redis + Postgres with the latest
   checkpoint_id.
4. The SSE stream drops — the UI detects the error and starts recovery polling (every 5s).
5. A replacement pod starts up and calls resume_interrupted_runs().
6. The new pod claims the interrupted run (lease-based, FOR UPDATE SKIP LOCKED), loads the checkpoint from Postgres, and
   resumes the graph from the last completed node.
7. The run completes on the new pod — the final response is written to the thread state.
8. The UI's recovery poller detects the completed thread state and renders the response — the user sees the answer as if
   nothing happened.

[nirmchan]

Hi @saharannaveen There are some minor issues have given inline review comments kindly check /address also the CI is failing with some issues, may need to be resolved. LGTM otherwise

[coderabbitai]

Walkthrough

Adds optional lifecycle persistence for compiled graphs. Redis tracks inflight runs and checkpoints. PostgreSQL stores leases, interruption metadata, and recovery status. Shutdown persists active runs, while startup initializes Aegra and resumes eligible runs. New settings control persistence and recovery. Tests cover encryption, Redis tracking, lease recovery, startup behavior, and pod coordination. The mock MCP server now uses FastMCP Streamable HTTP.

Suggested reviewers: np-compete, vishnusrichand

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the primary change: state persistence and recovery for interrupted runs.
Description check	✅ Passed	The description directly explains lifecycle persistence, interrupted-run recovery, startup and shutdown integration, and related tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

🚀 Post-Merge Actions

• Update changelog

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 58.13008% with 103 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
deep_agent/aegra/startup.py	34.61%	51 Missing ⚠️
deep_agent/aegra/lifecycle.py	73.72%	36 Missing ⚠️
deep_agent/aegra/graph.py	20.00%	12 Missing ⚠️
deep_agent/aegra/shutdown.py	63.63%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!

[NP-compete]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 9

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deep_agent/aegra/lifecycle.py`:
- Around line 445-530: Refactor persist_inflight_runs to expire all collected
run_ids with one database UPDATE using a run_id = ANY(...) filter, rather than
calling _expire_run_lease once per run. Preserve the existing running-status
condition, lease metadata updates, logging, deregistration, and persisted-count
behavior; update _expire_run_lease or replace its usage so only a single
synchronous Postgres connection is opened during shutdown.
- Around line 343-369: Update get_inflight_run_ids to batch the per-key hash
reads through a Redis pipeline rather than calling hgetall inside the loop.
Preserve filtering for the current POD_ID, collection of non-empty run_id
values, and the existing unavailable/error fallback behavior.

In `@deep_agent/aegra/startup.py`:
- Around line 160-172: Cap the interrupted-run recovery query in the startup
recovery flow by adding the same 20-row limit used by the stale-run diagnostic
query. Keep the recovered_run_ids collection and subsequent UPDATE behavior
unchanged, while ensuring the SELECT does not load more than 20 rows.

In `@scripts/inspect-db.sh`:
- Around line 55-68: Replace direct $THREAD_ID interpolation in the thread and
forensics SQL commands with psql variable binding, passing the identifier
through the psql invocation and referencing the bound variable in each query.
Apply the same protection to the additional affected query range while
preserving the existing query behavior and output.

In `@scripts/kind-pod-kill-test.sh`:
- Line 33: Replace the hardcoded username and password in the DB connection
string with values sourced from environment variables or the project’s existing
secret configuration, while preserving the current host, port, and database
name.

In `@scripts/pod-kill-and-recover.sh`:
- Line 104: Update the background Aegra command in the pod recovery script to
write logs to a securely created temporary file using mktemp, rather than the
predictable /tmp/agent.log path. Reuse the generated temporary filename for
output redirection and preserve the existing background startup behavior.

In `@tests/integration/test_lifecycle_integration.py`:
- Around line 30-75: Update the _FakeRedis test double to implement scan_iter
with the same prefix-pattern behavior as keys, yielding matching keys after
cleaning expired entries. Ensure get_inflight_run_ids can scan this mock
successfully so lifecycle assertions validate actual returned run IDs rather
than the module’s exception fallback.

In `@tests/unit/test_lifecycle.py`:
- Around line 465-476: Add pytest asyncio markers to test_no_interrupted_runs
and every other async test in TestResumeInterruptedRuns in
tests/unit/test_lifecycle.py (anchor site, lines 465-476), and to
test_resume_after_persist and test_stale_lease_reclaimed in
tests/integration/test_lifecycle_integration.py (sibling site, lines 177-206).
Alternatively, confirm the project config enables asyncio_mode = auto; do not
leave these async tests unmarked when that setting is unavailable.
- Around line 87-97: Replace the mocked `_get_encryption_secret` call in
`test_encryption_secret_uses_sso_client_secret` with an invocation of the real
`lifecycle._get_encryption_secret` function while retaining the patched settings
containing `SSO_CLIENT_SECRET`; assert that it returns the configured secret so
the test exercises the implementation.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: f1d52010-19bc-42a7-99ba-13a4b01ae37a

📥 Commits

Reviewing files that changed from the base of the PR and between 64c6065 and 00ea95e.

📒 Files selected for processing (14)

• deep_agent/aegra/graph.py
• deep_agent/aegra/lifecycle.py
• deep_agent/aegra/shutdown.py
• deep_agent/aegra/startup.py
• deep_agent/src/infrastructure/subagents.py
• deep_agent/src/settings.py
• scripts/inspect-db.sh
• scripts/kind-pod-kill-test.sh
• scripts/pod-kill-and-recover.sh
• tests/integration/conftest.py
• tests/integration/test_lifecycle_integration.py
• tests/mocks/mock_mcp_server.py
• tests/unit/aegra/test_graph.py
• tests/unit/test_lifecycle.py

🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

• redhat-data-and-ai/template-mcp (manual)
• redhat-data-and-ai/template-ui (manual)

[saharannaveen]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

scripts/inspect-db.sh (1)

102-108: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Compare the claimant with the original worker before labeling recovery, because a non-null, non-dash claimed_by only proves that a run is claimed and mislabels ordinary single-worker runs as recovered by a different pod.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/inspect-db.sh` around lines 102 - 108, Update the “WORKER POD
HISTORY” query in the runs case to compare claimed_by with the run’s original
worker identity, and only include runs where those values differ; retain the
existing non-null/non-dash filters and recovered-by-pod output.

scripts/kind-pod-kill-test.sh (1)

172-192: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Fail when claimed_by does not match a running pod instead of killing the first pod, because the fallback can terminate an unrelated replica and leave the target worker alive.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/kind-pod-kill-test.sh` around lines 172 - 192, Update the KILL
selection logic around HANDLING_POD and POD_TO_KILL so that no fallback pod is
selected when no running pod matches the claimed handler. Fail the script with
an error and stop before kubectl delete, ensuring only the explicitly matched
pod can be terminated.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deep_agent/aegra/lifecycle.py`:
- Around line 361-364: Update the pipeline flow around client.pipeline() and
pipeline.execute() to process keys in bounded batches instead of queuing every
scanned key at once. Execute each batch, collect or handle its results
incrementally, and ensure no unbounded list of keys or hashes is materialized
while preserving the existing hgetall behavior.

In `@scripts/inspect-db.sh`:
- Around line 177-182: Update the checkpoint aggregation query in the $PG
invocation to cast metadata::jsonb->>'step' to a numeric type before applying
MIN and MAX, so from_step and to_step use numeric ordering while preserving the
existing run and count aggregation.

---

Outside diff comments:
In `@scripts/inspect-db.sh`:
- Around line 102-108: Update the “WORKER POD HISTORY” query in the runs case to
compare claimed_by with the run’s original worker identity, and only include
runs where those values differ; retain the existing non-null/non-dash filters
and recovered-by-pod output.

In `@scripts/kind-pod-kill-test.sh`:
- Around line 172-192: Update the KILL selection logic around HANDLING_POD and
POD_TO_KILL so that no fallback pod is selected when no running pod matches the
claimed handler. Fail the script with an error and stop before kubectl delete,
ensuring only the explicitly matched pod can be terminated.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: d1722fd1-eec6-4c63-9bbe-4d12e1bb8cb0

📥 Commits

Reviewing files that changed from the base of the PR and between 00ea95e and 46bde3d.

📒 Files selected for processing (7)

• deep_agent/aegra/lifecycle.py
• deep_agent/aegra/startup.py
• scripts/inspect-db.sh
• scripts/kind-pod-kill-test.sh
• scripts/pod-kill-and-recover.sh
• tests/integration/test_lifecycle_integration.py
• tests/unit/test_lifecycle.py

🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

• redhat-data-and-ai/template-mcp (manual)
• redhat-data-and-ai/template-ui (manual)

[saharannaveen]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

♻️ Duplicate comments (1)

deep_agent/aegra/lifecycle.py (1)

361-364: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

Avoid unbounded key materialization

As per path instructions, the 50-key pipelines still rely on keys = list(client.scan_iter(...)), which retains every matching key in memory; consume the iterator incrementally into bounded batches.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@deep_agent/aegra/lifecycle.py` around lines 361 - 364, Update the
key-processing flow around the batch_size loop to consume client.scan_iter(...)
incrementally instead of materializing all keys with list(...). Accumulate at
most 50 keys per batch, process each batch before collecting the next, and
preserve the existing run_ids behavior.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Duplicate comments:
In `@deep_agent/aegra/lifecycle.py`:
- Around line 361-364: Update the key-processing flow around the batch_size loop
to consume client.scan_iter(...) incrementally instead of materializing all keys
with list(...). Accumulate at most 50 keys per batch, process each batch before
collecting the next, and preserve the existing run_ids behavior.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: b740cdaf-7b2e-4f0e-a058-07f7d51e0605

📥 Commits

Reviewing files that changed from the base of the PR and between 46bde3d and 4ccccec.

📒 Files selected for processing (3)

• deep_agent/aegra/lifecycle.py
• scripts/inspect-db.sh
• scripts/kind-pod-kill-test.sh

🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

• redhat-data-and-ai/template-mcp (manual)
• redhat-data-and-ai/template-ui (manual)

[saharannaveen]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deep_agent/aegra/graph.py`:
- Around line 385-407: The cached graph must not retain per-request lifecycle
state such as user identity or the encrypted refresh token. Update the flow
around build_execution_context and the _lifecycle_* setattr calls so this
context is supplied per request rather than stored on the shared compiled
object, or extend cache_key to include the user-specific identity before
_graph_cache lookup/storage; preserve shared graph reuse without cross-user
context leakage.

In `@deep_agent/aegra/lifecycle.py`:
- Around line 577-622: Change the resume batch flow around the SELECT/UPDATE
logic to claim all eligible rows and commit the transaction before invoking
resume_fn. Store the claimed run_id and thread_id values, then process them
afterward so resume_fn and its success/error updates do not run while the FOR
UPDATE locks are held; preserve per-run result recording, logging, and
deregister_inflight behavior.
- Around line 221-243: The register_inflight flow in lifecycle.py performs
synchronous psycopg.connect and UPDATE work on the async agent request path;
make this database write non-blocking by using an async PostgreSQL connection or
offloading the entire existing connection, cursor, commit, and update sequence
via asyncio.to_thread. Preserve the current started_by_pod update and debug
logging behavior while ensuring register_inflight does not block the event loop.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: 16ad62de-44ff-447c-a103-baf6277fa831

📥 Commits

Reviewing files that changed from the base of the PR and between 4ccccec and 5f64817.

📒 Files selected for processing (11)

• deep_agent/aegra/graph.py
• deep_agent/aegra/lifecycle.py
• deep_agent/aegra/shutdown.py
• deep_agent/aegra/startup.py
• deep_agent/src/infrastructure/subagents.py
• deep_agent/src/settings.py
• tests/integration/conftest.py
• tests/integration/test_lifecycle_integration.py
• tests/mocks/mock_mcp_server.py
• tests/unit/aegra/test_graph.py
• tests/unit/test_lifecycle.py

🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

• redhat-data-and-ai/template-mcp (manual)
• redhat-data-and-ai/template-ui (manual)