fix(short-term-memory): protect shared summarization, clear state on reset, close reliably - #1720

…sted summary

Review on #1720 found gaps in the close() and reset fixes from its
first commit.

close() set _closed in a finally inside the write lock. A close
cancelled while waiting for the lock never reached it, and a close
cancelled while waiting for summarization set it without clearing
anything, after which every later close() returned early from
_do_reset(), so the cleanup could never finish.

close() now sets _closed before any await, and runs its cleanup in a
task it shields: the cleanup runs to completion even if close() is
cancelled, and a repeated close() awaits the same task. The cleanup
still waits for in-flight summarization rather than cancelling it, so
the result is persisted for the next instance. That is deliberate:
with instance_cache_size at 0 the instance cache closes an instance as
soon as its last reference is released, which is at the end of nearly
every request, so a close that cancelled summarization would cancel it
in nearly every request that starts one.

clear_memory() dropped only the in-memory summary. The copy saved
through save_short_term_memory() survived, so a new instance for the
session restored a summary of what had been cleared, including after
delete_session_episodes(). clear_memory() now cancels in-flight
summarization, whose result would summarize the cleared content (no
query can be awaiting the worker while the write lock is held), and
overwrites the persisted copy with an empty one. close() leaves the
persisted copy alone: releasing a cached instance must keep it for the
next one.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_013zM4AyQyFwwUv3RNmMKwgY
Signed-off-by: Zubiao Xiong <[email protected]>

Review on #1720 found a window in clear_memory(). It cancels the
summarization worker and waits for it to stop, but if clear_memory()
was itself cancelled during that wait, it released the write lock while
the worker was still running.

An add_episodes() that evicted in that window queued its batch behind
the stopping worker: summarize() starts no worker while one is running,
and the cancelled worker exits without draining the queue. The batch
then waited for a later eviction to start a worker, and was discarded
if close() ran first. With instance_cache_size at 0, close() runs at
the end of nearly every request. A query in that window failed with
CancelledError although nothing had cancelled it, because the shield in
wait_until_done() passes the worker's cancellation on to its waiters.
The cancelled clear also left behind the episodes and the summary it
was meant to clear. The cancellation came with the previous commit;
before it, clear_memory() waited for summarization instead.

clear_memory() now runs in a task it shields, as close() does, so it
holds the write lock until the worker has stopped, and finishes
clearing even when the call is cancelled.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_013zM4AyQyFwwUv3RNmMKwgY
Signed-off-by: Zubiao Xiong <[email protected]>

Review on #1720 found that a clear_memory() already waiting for the
lock did nothing if close() was called meanwhile: close() sets _closed
before it takes the lock, and clear_memory() checked _closed only once
it had the lock, so the queued clear returned normally without
clearing anything, the persisted summary included. add_episodes(),
get_summary() and get_short_term_memory_context() made the same check
and failed with ShortTermMemoryClosedError in that position, where on
main they ran before the close. Both came with an earlier commit on
this branch, the one that made close() set _closed first.

Each of these methods now checks _closed when it is called, before it
takes the lock, so a call made before close() queues ahead of the
cleanup and completes first. A call made after close() is still
rejected, and clear_memory() still does nothing then.

test_close_cancelled_while_waiting_for_the_lock_still_closes now waits
for the cleanup before checking that it cleared: the add it tries is
rejected at once, before the cleanup has had a turn, where it used to
wait for the lock behind it.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_013zM4AyQyFwwUv3RNmMKwgY
Signed-off-by: Zubiao Xiong <[email protected]>

Review on #1720 found that _clear_memory() overwrote the persisted
summary before clearing the episodes, so a failing
save_short_term_memory() raised with the summary already cleared but
the episodes still readable. The episodes are now cleared first, and
the error still propagates, so the caller knows the stored copy may
remain.

Review also found that clear_memory() cannot cancel summarization that
a query is already waiting on: the query holds the lock, so the clear
waits for that summary and then discards it. On main, clear_memory()
always waited. This behavior is unchanged; the docstring now says so,
and a new test pins that the summary does not survive the clear.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_013zM4AyQyFwwUv3RNmMKwgY
Signed-off-by: Zubiao Xiong <[email protected]>

Review on #1720 pointed out that close()'s docstring said the instance
rejects new work as soon as close() is called, while delete_episode()
and get_summary_length() still run after it. The docstring now names
what happens: add_episodes(), get_summary() and
get_short_term_memory_context() raise ShortTermMemoryClosedError, and
clear_memory() does nothing. No behavior changes.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_013zM4AyQyFwwUv3RNmMKwgY
Signed-off-by: Zubiao Xiong <[email protected]>