Skip to content

Restore complete planned-restart chat recovery #540

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hamzamerzic merged 1 commit into from
Aug 2, 2026
+15 −96
Merged

Restore complete planned-restart chat recovery #540

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 6 additions & 42 deletions backend/app/chat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1643,12 +1643,6 @@ async def prepare_restart_intents(
# not delay an opted-in chat after its reset, but a bad/early reset timestamp
# also must not launch a whole parked batch in one burst.
LIMIT_AUTO_RESUME_STAGGER_SECS = 30.0
# Planned restarts may recover many exact turns at once. Restoring all of them
# in one lifespan sweep can consume the host before ordinary shell requests get
# a chance to run. Automatic recovery therefore uses a small process-wide
# admission ceiling; owner sends remain uncapped and take the available slots
# first, while durable restart parks stay due for the next sweep.
RESTART_AUTO_RESUME_MAX_ACTIVE = 2
_next_limit_auto_resume_at = 0.0
_RESTART_AUTHORIZATION_UNSET = object()

Expand Down Expand Up @@ -1677,14 +1671,6 @@ def _claim_limit_auto_resume_slot(now: float | None = None) -> bool:
return True


def _restart_auto_resume_capacity_available() -> bool:
"""Whether automatic restart recovery may add another live turn."""
return (
len(registry.all_alive_chat_ids())
< RESTART_AUTO_RESUME_MAX_ACTIVE
)


def _has_unanswered_question(chat: models.Chat | None) -> bool:
"""Whether the durable tail is waiting for an owner answer.

Expand Down Expand Up @@ -1773,8 +1759,9 @@ async def _auto_resume_chat(
# attribution, the exact latest park, and provider ownership at the
# actual claim point. Provider-limit retries are staggered by the
# reset sweep, rather than blocked on unrelated live chats. Planned-
# restart continuations are admitted against the process-wide recovery
# ceiling so a restart cannot crowd ordinary shell work off the host.
# restart continuations are different: they are the exact, owner-opted
# set that was already live together before the restart, so each chat
# may reclaim its own slot independently.
async with chat_queue.get_lock(chat_id):
with SessionLocal() as check_db:
chat = check_db.query(models.Chat).filter(
Expand Down Expand Up @@ -1849,11 +1836,6 @@ async def _auto_resume_chat(
resume_app_id = (
park.initiated_by_app_id if restart_park else None
)
if (
restart_park
and not _restart_auto_resume_capacity_available()
):
return False
if not mark_starting(chat_id):
return False
claimed = True
Expand Down Expand Up @@ -1959,11 +1941,9 @@ async def sweep_reset_parks(
at most one starts per sweep and launches are spaced even when unrelated
chats are live. App-attributed provider-limit runs never auto-resume.
Planned-restart continuations reclaim the exact set that was already live
before the restart and preserve each run's attribution, but only up to a
small process-wide recovery ceiling. Additional exact parks remain due
for a later tick, leaving capacity for owner-triggered work. A deferred
enabled chat stays pending while notify-only chats in the same due batch
still resolve normally.
before the restart, preserve each run's attribution, and may resume
independently. A staggered enabled chat stays pending for a later tick,
while notify-only chats in the same due batch still resolve normally.
App-attributed messages newly queued behind either kind of park still
require an ordinary app-owned handoff rather than being swept into the
synthetic continuation.
Expand Down Expand Up @@ -2078,13 +2058,6 @@ def wants_auto_resume(chat, run) -> bool:
# but keep walking so a later notify-only chat is not held hostage by
# another chat's auto-resume preference.
continue
if (
restart_auto_resume
and not _restart_auto_resume_capacity_available()
):
# Keep the exact park untouched. A later sweep admits it after another
# turn settles; an owner send can still start immediately in the meantime.
continue
if auto_resume:
try:
prepared = await _await_ack(get_writer().submit(
Expand Down Expand Up @@ -2134,15 +2107,6 @@ def wants_auto_resume(chat, run) -> bool:
queue_due_notification(chat_id, run)
continue

if (
restart_auto_resume
and not _restart_auto_resume_capacity_available()
):
# Capacity can disappear while the actor prepares the durable park.
# Leaving resume_pending intact is intentional: the next sweep retries
# without duplicating the continuation or its notification.
continue

if prepared.get("notify"):
queue_due_notification(chat_id, run)
if (
Expand Down
63 changes: 9 additions & 54 deletions backend/tests/test_limit_park.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@
while draining, and resolves deleted chats silently.
(e) Auto-resume is policy-controlled (off = notify only). Provider-limit
retries ignore unrelated live work and launch with a short stagger,
while an accepted planned restart restores the exact previously-live set
through a bounded recovery ceiling. Each resumed turn combines its
preserved queue + a "continue" into one continuation.
while an accepted planned restart resumes the exact previously-live set
together. Each resumed turn combines its preserved queue + a "continue"
into one continuation.
(f) The parks are observable: /api/debug/status lists parked runs.
(g) A planned restart reuses the same exact-run state with a due-now time;
crashes, unanswered questions, and app-owned work stay manual.
Expand Down Expand Up @@ -1433,10 +1433,10 @@ def test_sweep_starts_only_one_of_two_opted_chats(owner_token, monkeypatch):
chat_mod.discard_starting(cid)


def test_sweep_bounds_restart_recovery_and_retries_the_exact_remainder(
def test_sweep_restarts_every_opted_chat_in_the_accepted_batch(
owner_token, monkeypatch,
):
"""Restart recovery preserves every exact park without a launch stampede."""
"""A restart restores the exact set that was already concurrent."""
del owner_token
nonce = "restart-nonce-batch"
monkeypatch.setattr(
Expand All @@ -1463,20 +1463,11 @@ def _schedule(**kw):
)

try:
first = _run_sweep()
assert len(first) == chat_mod.RESTART_AUTO_RESUME_MAX_ACTIVE
assert {item["chat_id"] for item in scheduled} == set(first)
deferred = (set(chat_ids) - set(first)).pop()
assert _run_row(f"rt-{deferred}")["status"] == "parked"
assert {chat_id for kind, chat_id in events if kind == "notify"} == set(
first
)

# A settled recovery frees one slot. The next sweep starts the exact
# remaining park rather than consuming or replacing it.
chat_mod.discard_starting(first[0])
assert _run_sweep() == [deferred]
assert set(_run_sweep()) == set(chat_ids)
assert {item["chat_id"] for item in scheduled} == set(chat_ids)
assert [kind for kind, _ in events[:len(chat_ids)]] == [
"schedule", "schedule", "schedule",
]
assert {chat_id for kind, chat_id in events if kind == "notify"} == set(
chat_ids
)
Expand Down Expand Up @@ -1634,42 +1625,6 @@ async def scenario():
chat_mod.discard_starting(cid)


def test_restart_auto_resume_locked_claim_preserves_park_at_capacity():
"""Owner work already using the recovery budget keeps the exact park due."""
cid = "restart-capacity-locked-claim"
park_token = f"park-{cid}"
nonce = "restart-capacity-nonce"
_seed_chat(cid, auto_restart=True)
_seed_run(
cid,
park_token,
status="resume_pending",
park_reason="restart",
restart_nonce=nonce,
started_offset=-30,
)
blockers = [
_Handle(f"restart-capacity-live-{index}-{uuid.uuid4()}")
for index in range(chat_mod.RESTART_AUTO_RESUME_MAX_ACTIVE)
]
for blocker in blockers:
registry.register(blocker)

try:
assert asyncio.run(chat_mod._auto_resume_chat(
cid,
park_token=park_token,
restart_authorization=nonce,
)) is False
finally:
for blocker in blockers:
registry.unregister(blocker.chat_id, blocker.kind)

assert _run_row(park_token)["status"] == "resume_pending"
assert _chat_row(cid)["pending"] == []
assert not chat_mod.is_chat_running(cid)


def test_auto_resume_locked_claim_rejects_superseded_park():
"""The selected park can become stale while the sweep waits on the lock."""
cid = "auto-superseded-locked-claim"
Expand Down
Loading