Fix ui.timer cancelling @await_on_shutdown callbacks on shutdown

[Kshot3000]

Motivation

#6331 reports that @background_tasks.await_on_shutdown protection is lost for handlers dispatched through a client context. The issue lists three affected entry points: Client.safe_invoke (app.on_connect), handle_event (ui.button(on_click=...)) and Timer (ui.timer).

Re-checking all three against current main with the User test fixture (protected callback in flight, then background_tasks.teardown()):

• on_connect shorthand — now completes: since Report exceptions from fire-and-forget tasks to ui.on_exception #6363 these paths pass the original awaitable to background_tasks.create(..., context=...), and create() checks the original awaitable for the _AwaitOnShutdown marker.
• on_click shorthand — now completes, same reason (create_or_defer(..., context=...)).
• ui.timer(..., once=True) — still cancelled: the timer's callback never reaches background_tasks at all. It runs as a child asyncio task inside _invoke_callback, awaited by the timer's loop task, so teardown's cancellation of the timer task propagates straight into the protected invocation. A @await_on_shutdown save/backup started by a timer is silently dropped on shutdown.

This PR fixes the remaining Timer path.

Implementation

In Timer._invoke_callback, when the callback's result is an _AwaitOnShutdown awaitable, it is handed to background_tasks.create(...) — which registers it in running_tasks and _await_tasks_on_shutdown, so teardown() awaits it instead of cancelling it — and awaited through asyncio.shield(...), so cancelling the timer (or its invocation task) does not propagate into the protected work. Repeating timers keep their pacing: the invocation is still awaited before the next tick; only the cancellation propagation changes.

Trade-offs / notes:

• handle_exceptions=False on the inner create(): exceptions from the protected callback are handled by _invoke_callback's existing in-context handler exactly like unprotected callbacks. Verified that a raising protected timer callback is reported to ui.on_exception exactly once, identical to the unprotected baseline.
• Unprotected callbacks are untouched: verified a plain timer callback is still cancelled by teardown(), and that an in-flight protected invocation of a repeating timer also completes through teardown().
• The new regression test test_timer_invocation_is_awaited_on_shutdown asserts the user-visible contract (protected timer work completes through teardown(), mirroring test_queued_lazy_task_is_awaited_on_shutdown); it fails on unpatched main (events == ['cancelled']) and passes with this change.
• Related observation, not addressed here: Client.handle_exception still wraps results in helpers.await_with_context(...) before create(), which hides the marker the same way — happy to extend the fix if you want it in scope.

Test report: tests/test_background_tasks.py 8/8 user-fixture tests pass (incl. the new one, red→green verified); tests/test_timer.py user-fixture tests pass. Selenium Screen tests could not run in my sandbox (no Chrome), so those are left to CI.

Analysis and fix drafted with AI assistance (per AGENTS.md pair-programming), verified locally as described above.

Progress

• The PR title is a short phrase starting with a verb like "Add ...", "Fix ...", "Update ...", "Remove ...", etc.
• The implementation is complete.
• This PR does not address a security issue. (Security fixes must be coordinated via the security advisory process before opening a PR.)
• Pytests have been added/updated.
• Documentation is not necessary. (This restores the documented @await_on_shutdown behavior for timers; no API or doc changes.)
• No breaking changes to the public API.

---

Submitted by Kshot3000 (@Kshot9000 on X) via the BUG FIXER MAN public bug-fix log (github.com/Kshot3000/BUG-FIXER-MAN). No bounty was posted on #6331; the fix is offered freely — tips welcome (PayPal kyleblake0659@gmail.com · BTC 3GnR7TWBXAB3pPztBWpNF4LMNEX5yX8vZK · ETH 0xA1d3CEB7bD707847c3c6aB59d30D385FE8DD85Fc · SOL DLg1ua1ufewQ81J7dQomwWhXzhZreq4jEwTaQJR31Exb).

[evnchn]

Welcome to NiceGUI, and thanks for tracking this one down so carefully.

The fix works: a protected ui.timer or app.timer callback now survives shutdown, but it also loses its UI context, so one argument needs adding before merge.

Passing context=self._get_context() to create() (inline suggestion) fixes that; everything else I probed matches main, apart from one cancel behaviour I left as a question below.

What the fix does on a real shutdown (MRE + table)

Real server and real headless browser, with app.shutdown() fired with a 1 s protected save in flight for each timer:

@background_tasks.await_on_shutdown
async def save(tag: str) -> None:
    log(f'{tag}:start')
    if tag == 'ui.timer once':
        asyncio.get_running_loop().call_later(0.5, app.shutdown)
    try:
        await asyncio.sleep(1.0)
        log(f'{tag}:done')
    except asyncio.CancelledError:
        log(f'{tag}:cancelled')
        raise

app.timer(0.1, lambda: save('app.timer repeat'))

def build() -> None:
    ui.timer(0.1, lambda: save('ui.timer once'), once=True)
    ui.timer(0.1, lambda: save('ui.timer repeat'))
# run either as script mode (build() at top level) or inside @ui.page('/')

callback in flight at shutdown	main, script mode	main, @ui.page	PR, script mode	PR, @ui.page	Implications
ui.timer(..., once=True)	cancelled	cancelled	done	done	shutdown now awaits it
repeating ui.timer	cancelled	cancelled	done	done	shutdown now awaits it
repeating app.timer	cancelled	cancelled	done	done	app.timer fixed too

The new test fails with timer.py reverted to main (assert ['cancelled'] == ['done']) and passes on the PR, so it guards the fix.

The context regression, and the probes that are unchanged (User fixture, main vs PR vs PR + suggestion)

probe	main	PR	PR + suggestion	Implications
protected callback calls ui.label() / ui.context.client	works	RuntimeError (slot stack empty)	works	regression; suggestion fixes it
same with an unprotected async callback	works	works	works	unprotected path untouched
sync callback	works	works	works	unprotected path untouched
protected callback raises ValueError	app.on_exception 1×, page ui.on_exception 1×	same	same	exception routing kept
plain async callback at teardown()	cancelled	cancelled	cancelled	no protection leaks out
repeating protected timer: protected tasks alive at once	0	1	1	no pile-up across ticks
deactivate() on a repeating protected timer	stops ticking	stops ticking	stops ticking	deactivate still works

Why: Slot.stacks is keyed by asyncio task id, so the create()d task starts with an empty stack.
create(context=...) from #6363 re-enters the slot and the client inside the new task.

Suites: tests/test_timer.py and tests/test_background_tasks.py pass on the PR and on PR + suggestion, Screen tests included, no skips.
test_cleanup and test_no_leak_when_client_deleted in test_timer.py fail now and then locally, on unmodified main too, so a red there in CI is probably not this PR.

One behaviour change for maintainers to confirm (not a request for this PR)

Because the protected work is shielded, cancel(with_current_invocation=True) no longer stops it, and neither does deleting the timer, which goes through that same call:

action while a protected callback runs	main	PR	Implications
timer.cancel(with_current_invocation=True)	cancelled	runs to completion	explicit cancel now ignored
timer's parent .clear()-ed	cancelled	runs to completion	work outlives its element

@await_on_shutdown is documented as shutdown-only, so this may or may not be wanted; it reads to me like a maintainer call rather than something to fold into this PR.

※

[evnchn]

Required: the new task starts with an empty slot stack, because NiceGUI keys slot stacks by asyncio task.
So a protected callback that touches the UI (ui.label(...), ui.context.client) now raises RuntimeError: The current slot cannot be determined...; it worked on main.
Passing the timer's context, which #6363 added to create(), keeps it (indentation as autopep8 wants it):

Suggested change

	task = background_tasks.create(result, name=str(self.callback), handle_exceptions=False)
	task = background_tasks.create(result, name=str(self.callback), handle_exceptions=False,
	context=self._get_context())

With this applied, the protected callback sees its client and slot again, and a raising callback still reaches app.on_exception and the page's ui.on_exception exactly once.
app.timer (where _get_context() is a nullcontext()) still completes through teardown.

[Kshot3000]

Applied in 7534417 — background_tasks.create(...) now passes context=self._get_context(), exactly as suggested.

Verified locally: tests/test_timer.py and tests/test_background_tasks.py pass (User-fixture tests; Screen tests not run here), and I re-ran your context probe — a protected ui.timer callback calling ui.label() sees its client and slot stack again and completes through teardown() instead of raising RuntimeError.

Left the cancel(with_current_invocation=True) behaviour unchanged, since you flagged it as a maintainer call rather than part of this PR.

[evnchn]

Approving 7534417: with the context fix, a protected timer callback can use the UI again, and test_timer plus test_background_tasks pass (26 passed).

Optional: the test below would keep the context from slipping again; it fails without your last commit and passes with it.

Test for tests/test_background_tasks.py

async def test_protected_timer_callback_keeps_its_context(user: User):
    @background_tasks.await_on_shutdown
    async def work() -> None:
        ui.label('added by protected timer')

    @ui.page('/')
    def page():
        ui.timer(0.01, work, once=True)

    await user.open('/')
    await user.should_see('added by protected timer')

※

[Kshot3000]

Added the suggested regression test in 2a8fe03 — test_protected_timer_callback_keeps_its_context in tests/test_background_tasks.py, exactly as you proposed.

Verified locally: the test fails without the context fix in 7534417 (the callback raises when it tries to create the label) and passes with it; tests/test_background_tasks.py passes (9 passed, Screen-fixture test not run here).

[evnchn]

Re-approving 2a8fe03: the new test fails without the context fix (RuntimeError: ... the slot stack for this task is empty) and passes with it; test_background_tasks plus test_timer give 27 passed.

※