PR #7: fix: drain LiteLLM callback tasks before agentic_session exit

• Repository: Kelet-ai/python-sdk
• Author: AlmogBaku
• State: closed | merged: True
• Link: https://github.com/Kelet-ai/python-sdk/pull/7

Description (problem / solution / changelog)

User description

Summary

• Add _drain_background_logging_tasks in _context.py that polls asyncio.all_tasks() minus the current task and LiteLLM's long-lived _worker_task, waiting for a ~50ms idle window or a 5s hard timeout.
• Call it from _AgenticSessionContext.__aexit__ before self._exit() so LiteLLM's async callbacks get a turn on the live event loop.
• Debug-log (not silently swallow) any drain failure or missing LiteLLM internals.

Why

Single-completion scenarios (extended-thinking, simple-reasoning, etc.) were losing their litellm_request and raw_gen_ai_request spans. LiteLLM fires its OTEL callbacks via asyncio.create_task(...) through a background queue worker. When the user's async main() returned, asyncio.run() cancelled those pending tasks before they could create the spans, so the BatchSpanProcessor had nothing to export — the user's telemetry silently dropped.

The drain happens in aexit (not atexit) specifically because we need a live loop to pump the pending task set.

Upstream follow-up

Filed https://github.com/BerriAI/litellm/issues/25978 requesting a public await GLOBAL_LOGGING_WORKER.drain(...) API. Once that lands we can drop the private-attribute coupling (the one remaining probe of _worker_task).

Test plan

• uv run pytest — 112 SDK tests pass, including 8 new drain tests:
  • idle-return path
  • pending-task wait
  • 5s timeout bound (no-hang proof)
  • LiteLLM _worker_task exemption via injected sys.modules stub
  • missing-LiteLLM-module path
  • graceful degradation when LiteLLM refactors away _worker_task (debug-log assertion)
  • aexit call-order wiring
  • aexit swallows drain exceptions + logs
• e2e simple-reasoning — litellm_request + raw_gen_ai_request spans present (Bedrock reasoning text captured)
• e2e flat (3 sibling agents) — all 4 request spans present
• e2e nested — parent tool-call → sub-agent → tool-response flow intact
• e2e deep-nesting (3 levels) — all 5 request spans present
• e2e agent-return — two sequential completions under same agent
• e2e instruction-override — system-instruction change between calls

Review feedback addressed

• Moved drain into _context.py (lifecycle logic colocated; no circular-import workaround).
• asyncio.get_event_loop() → get_running_loop().
• LiteLLM internal surface reduced from 5 private attrs → 1 (_worker_task), with a debug log on each fallback path so refactor regressions are observable.
• except Exception: pass in aexit now _logger.debug(..., exc_info=True).
• Drain tuning constants factored out so tests can shrink them without patching asyncio.sleep.

🤖 Generated with Claude Code

---

Generated description

Below is a concise technical summary of the changes proposed in this PR: Implement <code>_drain_background_logging_tasks</code> so LiteLLM instrumentation callbacks can finish on the live event loop and log any failures before <code>_AgenticSessionContext.aexit</code> calls <code>_exit</code>. Bump the SDK to 1.4.0 to ship the telemetry fix plus its regression tests.

<table><tr><th>Topic</th><th>Details</th><tr><td><a href=https://baz.co/changes/Kelet-ai/python-sdk/7?tool=ast&topic=Other>Other</a> </td><td>Other files<details><summary>Modified files (1)</summary><ul><li>src/kelet/_configure.py</li></ul></details><details><summary>Latest Contributors(2)</summary><table><tr><th>User</th><th>Commit</th><th>Date</th></tr><tr><td>[email protected]</td><td>feat: add LiteLLM and ...</td><td>April 16, 2026</td></tr> <tr><td>[email protected]</td><td>fix(KEL-374): remove d...</td><td>April 07, 2026</td></tr></table></details></td></tr> <tr><td><a href=https://baz.co/changes/Kelet-ai/python-sdk/7?tool=ast&topic=Session+exit+drain>Session exit drain</a> </td><td>Implement <code>_drain_background_logging_tasks</code> and call it from <code>_AgenticSessionContext.__aexit__</code> so pending LiteLLM OTEL callbacks on the live loop complete before <code>_exit</code> runs, while still exempting the long-lived worker task and surfacing drain failures via <code>_logger</code>.<details><summary>Modified files (1)</summary><ul><li>src/kelet/_context.py</li></ul></details><details><summary>Latest Contributors(1)</summary><table><tr><th>User</th><th>Commit</th><th>Date</th></tr><tr><td>[email protected]</td><td>fix(KEL-342): propagat...</td><td>March 23, 2026</td></tr></table></details></td></tr> <tr><td><a href=https://baz.co/changes/Kelet-ai/python-sdk/7?tool=ast&topic=Drain+tests>Drain tests</a> </td><td>Exercise <code>_drain_background_logging_tasks</code> under idle, pending, timeout, LiteLLM exemption, missing module, attribute-absence, baseline snapshot, exception, and cancellation scenarios so the agentic session lifetime reliably flushes telemetry spans.<details><summary>Modified files (1)</summary><ul><li>tests/test_drain.py</li></ul></details><details><summary>Latest Contributors(0)</summary><table><tr><th>User</th><th>Commit</th><th>Date</th></tr></table></details></td></tr> <tr><td><a href=https://baz.co/changes/Kelet-ai/python-sdk/7?tool=ast&topic=Release+bump>Release bump</a> </td><td>Publish the SDK as version 1.4.0 so users receive the new drain behavior and test suite updates.<details><summary>Modified files (1)</summary><ul><li>uv.lock</li></ul></details><details><summary>Latest Contributors(1)</summary><table><tr><th>User</th><th>Commit</th><th>Date</th></tr><tr><td>[email protected]</td><td>feat: add LiteLLM and ...</td><td>April 16, 2026</td></tr></table></details></td></tr></table> This pull request is reviewed by Baz. Review like a pro on <a href=https://baz.co/changes/Kelet-ai/python-sdk/7?tool=ast>(Baz)</a>.

Changed files

• src/kelet/_configure.py (modified, +1/-3)
• src/kelet/_context.py (modified, +134/-1)
• tests/test_drain.py (added, +248/-0)
• uv.lock (modified, +1/-1)