fix(sessions): restore session history when compaction replacement is cancelled

[nothariharan]

Summary

This pull request fixes a data-loss path in OpenAIResponsesCompactionSession where cancelling a run during compaction replacement could permanently wipe the underlying session history.

PR #3117 already restored previous history when replacement failed with a normal Exception (for example, a failed add_items() after clear_session()). That restore path does not run for asyncio.CancelledError, because on supported Python versions CancelledError is a BaseException, not an Exception.

What goes wrong today

Compaction replacement is intentionally destructive:

1. Snapshot the full underlying history.
2. clear_session() — history is now empty.
3. add_items(compacted_output) — write the compacted replacement.

If cancellation lands after step 2 and before step 3 finishes, the except Exception handlers from #3117 do not run. The compacted write never lands, restore never runs, and the session remains empty.

previous history  -->  clear_session()  -->  [CancelledError]  -->  empty session
                              ^
                              restore skipped (CancelledError is not Exception)

This is especially sharp for immediate cancel (task.cancel()), which is delivered at await points. The post-clear add_items() await is exactly such a point.

What this PR changes

• Treat clear → add as one replacement transaction for both Exception and CancelledError.
• On CancelledError, restore previous history with the same clear-failed vs post-clear routing used for Exception.
• Await restore through a shielded drain loop so a second cancel during restore cannot leave the rewrite unfinished before CancelledError reaches the caller.
• Keep the existing Exception restore behavior unchanged (no shield on that path).

Restore helpers now accept BaseException so cancellation can be logged through the existing warning helpers.

Relationship to nearby work

Change	Relationship
#3117 / #3116	Same restore contract; this PR closes the CancelledError gap left by Exception-only handlers
#4293	Different bug (limited input window permanently dropping older history). This PR does not change limit / candidate loading behavior

Test plan

Added regression coverage next to the #3117 tests in tests/memory/test_openai_responses_compaction_session.py:

• Restore when replacement add_items() raises CancelledError after a successful clear
• Restore when clear_session() mutates then raises CancelledError
• Adversarial second cancel while restore is mid-flight (task.cancel() + yield while restore is gated), asserting history is already restored when CancelledError surfaces

Local verification:

• uv run ruff format / uv run ruff check
• uv run pyright / uv run mypy on the touched compaction module
• uv run pytest -q tests/memory/test_openai_responses_compaction_session.py (51 passed)
• Focused cancel/restore subset (8 passed)
• Full make tests via the verification script could not be run here because make is unavailable in this Windows environment; broader pytest run showed only unrelated local tracing environment failures

Issue number

N/A — discovered by code-path review against the #3117 restore contract; no open issue tracks the cancel-specific gap.

Checks

• I've added new tests, if relevant
• I've run formatting and linting on the changed files
• I've confirmed the compaction session test suite passes
• If using Codex, I've run /review before submitting this PR

[seratch]

Thanks for the focused cancellation fix. The main path works, but the rollback currently has no ownership boundary over the session state it rewrites. I reproduced this with SQLiteSession: pause cancellation restore before its clear, append a newer item through the same OpenAIResponsesCompactionSession, then resume restore. The rollback deletes the newer item, while _session_items still retains it, so persistent history and the compaction cache diverge.

Before merge, please add one session-owned mutation lock covering the snapshot, clear/add replacement, and restore, and use the same lock for wrapper add_items(), pop_item(), and clear_session(). Add a controlled interleaving regression test proving the newer write waits and survives in chronological order after cancellation recovery.

[nothariharan]

Yeah cool will get on it, should have done that earlier mb.

I can reproduce the failure mode you described: if cancellation restore is paused before its clear, a newer add_items() through the same OpenAIResponsesCompactionSession can land, and then restore clears/rewrites from the pre-cancel snapshot, deleting the newer item while _session_items still reflects it.

I'll add a session-owned mutation lock that covers:

1. snapshot + clear/add replacement + restore, and
2. the wrapper mutators add_items(), pop_item(), and clear_session()

I'll also add a controlled interleaving regression test with SQLiteSession that pauses restore before clear, appends a newer item, resumes restore, and asserts the newer write waits on the lock and survives in chronological order after recovery.

[seratch]

Thanks for addressing the original rollback ownership issue. A broader cancellation probe found one remaining data-loss path. If replacement add_items() first raises a normal Exception, then the caller cancels while restore is awaiting, the Exception path uses shield_restore=False; restore is cancelled, SQLite history remains empty, and the wrapper caches retain the previous history.

Please make restore settlement cancellation-safe regardless of whether replacement first failed with Exception or CancelledError. Remove the conditional shield behavior and drain every restore before propagating cancellation. Add a controlled SQLite regression test that gates restore, cancels the outer compaction task, proves it remains pending until restore is released, and then asserts CancelledError, restored history, cache agreement, a released lock, and no orphan restore task. Existing behavior should still re-raise the original exception when no cancellation occurs.

[nothariharan]

Thanks — good catch. Confirmed on the current head: the Exception path still uses shield_restore=False, so if replacement add_items() fails with a normal Exception and the caller cancels while restore is awaiting, restore is aborted, the SQLite history can stay empty, and the wrapper caches keep the old view.

I'll make restore settlement cancellation-safe for both Exception and CancelledError (always drain restore before cancel propagates), drop the conditional shield, and add a controlled SQLite regression that gates restore, cancels the outer compaction task, and asserts CancelledError + restored history + cache agreement + released lock with no orphan restore. Existing “no cancel → original Exception re-raised” behavior will stay covered.