Skip to content

fix(memory): restore session history when compaction replacement is cancelled - #4298

Open
nothariharan wants to merge 2 commits into
openai:mainfrom
nothariharan:fix/compaction-restore-on-cancel
Open

fix(memory): restore session history when compaction replacement is cancelled#4298
nothariharan wants to merge 2 commits into
openai:mainfrom
nothariharan:fix/compaction-restore-on-cancel

Conversation

@nothariharan

Copy link
Copy Markdown

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

…ancelled

CancelledError is a BaseException, so the Exception-only restore path left sessions empty after a successful clear during compaction replacement.

@seratch seratch left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Author

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.

…writes

Add a session-owned mutation lock around snapshot/replace/restore and wrapper mutators so cancel-restore cannot rewrite past a newer concurrent write.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants