fix: hold references to the persistence tasks

[noron12234]

The problem

Every write to the data layer is fired and forgotten:

if data_layer := get_data_layer():
    try:
        asyncio.create_task(data_layer.create_step(step_dict))
        self.persisted = True
    except Exception as e:
        ...

The task is discarded, and the event loop only keeps a weak reference:

Important: Save a reference to the result of this function, to avoid a task disappearing mid-execution. The event loop only keeps weak references to tasks. A task that isn't referenced elsewhere may get garbage collected at any time, even before it's done.

For logging or telemetry that would cost a datapoint. Here it costs user data: a collected task means a message, step or element is never written. There is no exception and nothing in the logs — the write simply never happens. self.persisted = True is set regardless, so the object believes it was saved. The failure only surfaces later as a thread with holes in it.

Note the try/except does not help: create_task schedules and returns immediately, so an error inside the coroutine never reaches that handler either.

Sites

Eight, across four modules — every place a data_layer.* write is spawned:

file	call
context.py:98	update_thread during http context init
element.py:215	create_element
message.py:116	update_step
message.py:135	delete_step
message.py:150	create_step
step.py:345	update_step
step.py:370	delete_step
step.py:396	create_step

The change

Each module keeps a set and discards the task in a done callback:

_persistence_tasks: set[asyncio.Task] = set()
...
_task = asyncio.create_task(data_layer.create_step(step_dict))
_persistence_tasks.add(_task)
_task.add_done_callback(_persistence_tasks.discard)

Module-level rather than instance-level because Message, Step and Element are short-lived — an instance-scoped set would be collected along with the object it was meant to outlive. discard rather than remove so a double callback cannot raise. The set stays bounded by the number of in-flight writes.

No control flow, ordering or awaited behaviour changes. These writes were fire-and-forget before and remain so; they simply can no longer be collected mid-write.

Verification

$ ruff check chainlit/context.py chainlit/element.py chainlit/message.py chainlit/step.py
Found 2 errors.          # both pre-existing on main, unchanged by this PR
$ ruff format --check ...
4 files already formatted
$ python -m py_compile ...

I want to be straight about what I did not do: I did not run the test suite, so this is verified by lint and compile only. The change is mechanical — bind the task, register it, deregister on completion — but if you would like the suite run before merging, say so and I will.

No test is added either. The failure is a garbage-collection race, so a test would have to force a GC at a chosen moment and assert a task did not vanish, which is flaky by construction. Happy to add one if you have a shape in mind.

Found with an AST scan for create_task / ensure_future results discarded as bare expression statements (excluding TaskGroup.create_task, which does hold strong references). The scan reports other hits under llama_index/, openai/, mistralai/ and socket.py; those spawn UI-emit and stream work rather than data-layer writes, so I left them out to keep this diff to one concern. Happy to follow up.

---

Summary by cubic

Prevents lost data writes by holding strong references to all persistence tasks. Previously, writes used bare asyncio.create_task(...) and tasks could be garbage collected before reaching the data layer; now each task is tracked until completion while keeping fire-and-forget semantics.

• Track in-flight tasks via a module-level _persistence_tasks: set[asyncio.Task] with add_done_callback(_persistence_tasks.discard) in context.py, element.py, message.py, and step.py (including Step.__enter__/__exit__ dispatches of send()/update()).
• Add backend/tests/test_persistence_tasks.py to pin the invariant (tasks held during writes and released on completion) and verify the set drains and stays bounded.
• No API, flow, or ordering changes; memory use is bounded by in-flight writes; no migration required.

^{Written for commit 551e38e. Summary will update on new commits.}

[cubic-dev-ai]

No issues found across 4 files

_{Re-trigger cubic}

[github-actions]

This PR is stale because it has been open for 14 days with no activity.

[dokterbob]

@noron12234 Thanks for the contrib. Could you add some unit tests for this functionality?

[noron12234]

@dokterbob Added in 551e38e — backend/tests/test_persistence_tasks.py, 10 tests.

One per write site (3 message, 3 step, 1 element, 1 context), plus the two Step.__enter__/__exit__ dispatches and a check that the set is a live registry rather than an accumulator.

Shape. The data layer is an AsyncMock whose method blocks on an asyncio.Event, so the write is deterministically still in flight when the assertion runs. Each test asserts the module holds exactly one strong reference to the in-flight task, then opens the gate and asserts the reference is discarded once the task is done.

Assertions are on the delta of each module's task set, not its absolute contents. test_message.py::test_update_with_data_layer patches asyncio.create_task with a MagicMock, so a mock lands in _persistence_tasks and is never discarded (its add_done_callback is also a mock). Absolute assertions passed alone and failed in the full suite; the delta approach is immune to it and doesn't require touching that test.

What I did not write, and why. I said in the description that a GC test would be flaky by construction. That was the wrong reason — I checked, and the real problem is that it would be vacuous:

mode=event  alive_after_gc=1 coroutine_completed=True
mode=sleep  alive_after_gc=1 coroutine_completed=True

That's fire-and-forget create_task with the result discarded, on unpatched CPython, gc.collect() called while suspended. The task survives either way — a suspended task is reachable from whatever it awaits (the Event's _waiters deque, or a timer handle on the loop), and that thing holds a callback pointing back at the task. So "force a collection, assert it survived" passes just as happily against the unfixed code. It would look like a regression test and test nothing.

What is left is the invariant: a strong reference exists for the duration of the write. That does mean the tests reference the private _persistence_tasks. I think that's the right trade here, but say the word if you'd rather they only assert through public surface.

Verification (this time I ran it):

• pytest backend/tests → 790 passed. Two files error at collection on google.auth / asyncpg; both are missing optional deps, present on main before this branch.
• New file alone, 20 consecutive runs → 10 passed every time, no flakes.
• Reverted the fix with git apply -R, kept the tests → all 10 fail. Restored → all 10 pass.
• ruff check / ruff format --check / mypy backend/ all clean; the repo's pre-commit hook ran and passed.

Being straight about the failure mode when reverted: the tests fail with AttributeError: module 'chainlit.message' has no attribute '_persistence_tasks', not with a caught garbage collection. They pin the fix's invariant, they don't independently reproduce the bug. Given the above, I don't think anything stronger is actually available here.