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, including Step.__enter__/__exit__ dispatches. Every data-layer write is now tracked until completion so tasks can’t be garbage collected mid-execution.

• Bug Fixes
  • Track in-flight persistence tasks in a module-level _persistence_tasks: set[asyncio.Task] in context.py, element.py, message.py, and step.py (including context-manager send()/update() calls).
  • Replace bare asyncio.create_task(...) with _task = asyncio.create_task(...); add to the set, then _task.add_done_callback(_persistence_tasks.discard).
  • No behavior change to flow or ordering; writes remain fire-and-forget but are no longer dropped.

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

[cubic-dev-ai]

No issues found across 4 files

_{Re-trigger cubic}