fix(livekit): flush traceAI's spans when a LiveKit job shuts down (TH-8730) - #282

What does this PR do?

traceai-livekit now sends the spans it is still holding when a LiveKit job shuts down, so voice traces arrive with their root agent_session span. enable_http_attribute_mapping() / enable_grpc_attribute_mapping() add a shutdown callback to the job they run in that flushes the tracer provider. Agents need no code change: the examples and the README already call enable_* in the entrypoint. They also register a second flush that runs when the job process exits, which sends what LiveKit ends after the shutdown callbacks (its own job_shutdown and job spans) and covers tracing set up outside a job. A new flush_on_job_shutdown() lets set-up done outside a job (for example at process start) also flush as each call ends.

Why?

Linear: TH-8730. Every voice trace from the phase 1 e2e run on the dev box (livekit-agents 1.8.3) was missing exactly its agent_session span, so the Traces tab showed "Root span not stored yet".

The cause, reproduced on my laptop:

• traceAI exports through a batch processor that sends every 5 s.
• On Linux, LiveKit runs each job in a forkserver process, which exits with os._exit() and skips atexit.
• At job end LiveKit flushes only its own processors.

So the spans that end last, the root among them, stayed queued and were dropped. Running a fake job through LiveKit's own job-process executor:

How it works

• flush_on_job_shutdown(provider=None) finds the job with get_job_context(), called with no arguments because required= only exists from livekit-agents 1.5.3. It adds one async shutdown callback per job and provider.
• It also registers, once per process, a multiprocessing finalizer at exit priority 10. LiveKit's forkserver children run these finalizers before os._exit(), so it flushes every registered provider just before the process stops. Outside a job this is all that's registered, with a debug line.
  • A forkserver child clears its finalizers just before its target runs, after any module it imported while unpickling (an agent module doing set-up at its top) has registered one. An after-fork hook makes ours again right after the clear, and it's also made again whenever it's gone or belongs to another process.
  • Registration is under a lock (LiveKit's thread executor can enable from several jobs at once). Each provider flushes in its own thread under one shared deadline, so one that hangs doesn't hold up the others.
• Everything is under a broad except, so on any livekit-agents 1.x it registers or is skipped and never raises into the entrypoint.
• Both flushes run provider.force_flush() in a daemon thread and wait for it at most 3 s, well inside LiveKit's 10 s exit budget:
  • OpenTelemetry's batch force_flush ignores its timeout, and LiveKit joins the loop's default executor before the process exits, so asyncio.to_thread would let a hung collector hold the job open.
  • A failure or timeout logs one warning; the callback never raises.
• It flushes and never shuts down: LiveKit can run several jobs in one process.

How was it tested?

• Unit tests added / updated (pytest): new tests/test_shutdown_flush.py (24 cases), plus tests/_job_process_entrypoint.py and tests/_job_process_setup_at_import.py. The shared fixtures moved from test_integration.py to tests/conftest.py.
  • One test runs a real LiveKit job process started with forkserver and checks both agent_session and LiveKit's job_shutdown reach traceAI's exporter before the process exits.
  • A second one sets tracing up at import of a module only the job process loads, and checks agent_session, job_shutdown and job_entrypoint arrive. On the previous commit that set-up sent nothing at all.
  • Others cover the flush running off the event loop in a daemon thread, a blocked or failing flush not holding up or breaking shutdown, registering once per job and provider, the exit flush (registered once per process, flushing every provider, bounded), the hook that re-arms it in a child, a child keeping its inherited providers, one registration under concurrent set-up, a hung provider not blocking the others, a failing exit flush logging a warning, both waits fitting LiveKit's shutdown_process_timeout (read from ServerOptions), no warning outside a job, and older LiveKit job APIs, including the pre-1.5.3 one during a job.
• Integration tests pass (or N/A — no gateway behavior changed): tests/test_integration.py tests/test_mapping.py tests/test_shutdown_flush.py in one run, in CI's order: 82 passed (livekit-agents 1.8.6 from the lock, Python 3.11).
• ruff check passes on the touched files.
• Public types, env vars, or SDK behavior changes are documented in the relevant README: a "Flushing at the end of a job" section, and a CHANGELOG line under 0.1.2. 0.1.2 is not on PyPI yet, so this joins it.
• Bugs planted one at a time, all caught by a test:
  • the flush removed;
  • enable_* no longer calling it;
  • the flush run on the event loop, or on the loop's default executor;
  • a non-daemon thread;
  • the bound removed, or put back to 10 s;
  • get_job_context(required=False);
  • the exit flush not registered, unbounded, registered on every call, or flushing only the first provider;
  • the after-fork hook not registered, not re-arming, or re-arming on the pid alone;
  • the registration lock removed;
  • the exit flush run one provider after another;
  • a child clearing its inherited providers;
  • the exit-flush warning downgraded, the exit cap raised to 5 s, and the per-job list appended only for the first provider;
  • the flush's error handling removed;
  • the once-per-job guard removed;
  • the broad except narrowed.

Checklist

• Branch is off dev (the PR goes into dev; main is for releases)
• Commit messages follow Conventional Commits (feat:, fix:, docs:, chore: …)
• No TODOs or commented-out code left in
• No real API keys or secrets in the diff
• If prose was added or changed: checked against docs/VOCABULARY.md (N/A: not in this repo)

Notes for reviewers

• Known limits (also in the README):

  • A process LiveKit kills (for example when a worker drains and the job overruns its exit time) sends nothing more.
  • Before livekit-agents 1.3, LiveKit closes the session at the same time as the job-shutdown flush, so that flush can miss agent_session; the exit flush still sends it.
  • A second AgentSession still closing during the shutdown callbacks is sent by the exit flush, not the job-shutdown one.
  • In LiveKit Cloud observability mode the session's parent span is recorded on LiveKit's own provider, so Future AGI has no stored root either way.

• Changed after the last review (commit 7316d78, for review 5600437004 on cf0d4e4): no reviewer has seen these yet.

  • The Reviewer's first Medium: set-up at import in a Linux job process now keeps its exit flush. The fix is the after-fork hook, plus making the finalizer again when it's gone or belongs to another process. On cf0d4e4 that set-up exported nothing. With this change a probe run of it sent all 12 spans; the test asserts the three that matter: agent_session, job_shutdown and job_entrypoint.
  • The Reviewer's second Medium: the README says job_entrypoint is sent only when tracing is set up before the job starts. With set-up inside the entrypoint, the job trace is rooted at job_shutdown.
  • The Lows:
    • the registration lock;
    • a flush thread per provider;
    • the budget test reads LiveKit's own timeout;
    • tests for the three plants that survived the review.
  • Other changes:
    • A child now keeps the providers it inherited instead of clearing them. OpenTelemetry empties their queues at fork, so nothing is sent twice.
    • The flush runs inline if no thread can start.

  One cold review (Opus) found nothing above Low. I fixed the two Lows it raised: an order-dependent test, and a silent skip when a thread can't start. All 24 planted bugs above are caught.

• Follow-ups (Lows from the round 3 review, left for a later PR):

  • If a flush thread can't start at exit, the fallback flushes inline with no deadline. cf0d4e4 skipped that provider instead.
  • No test has two hung providers, so a bug that gave each thread the full 3 s would pass.
  • No test catches the exit-flush warning's any changed to all.
  • The after-fork hook swaps the lock itself. os.register_at_fork would be cleaner, and would also cover plain os.fork children.
  • Your own fork-start children now flush the providers they inherited at exit, over connections opened before the fork.
  • On Python 3.13, fork and forkserver children run atexit, so Future AGI's provider shutdown there can wait past 3 s. This was already the case before this PR.
  • In a job process, the exit flush's warnings don't reach LiveKit's logs.
  • traceAI never reports a failed flush beyond a warning in the log.
  • If another shutdown callback raises, LiveKit stops waiting for ours; the exit flush still runs.

• The last check is a rerun of the voice agent on the dev box with this branch installed, to confirm the root span is stored and the Traces tab opens the trace.

🤖 Generated with Claude Code