Skip to content

fix(live): stop an early-closed merged live stream from hanging - #7358

Open
harshal-96 wants to merge 1 commit into
google:mainfrom
harshal-96:fix/live-merge-close-hang
Open

harshal-96 wants to merge 1 commit into
google:mainfrom
harshal-96:fix/live-merge-close-hang

Conversation

@harshal-96

Copy link
Copy Markdown

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

Problem:

_merge_live_event_streams merges two pumps into a one-slot queue, merged. When the caller closes the stream while merged is full, the cleanup cancels _pump_queued_events. Its finally then waits forever in await merged.put(done_sentinel), because nothing reads merged anymore, so aclose() never returns. Details and a reproduction script are in the issue.

Solution:

The merge now sets a local consumer_done flag in its finally, before the cleanup starts. _pump_queued_events only puts the sentinel if the consumer is still reading. Once the consumer is gone, nobody is waiting for the sentinel.

  • While the consumer is still reading, nothing changes: a normal end or a failing pump still delivers the sentinel.
  • A pump that was already blocked on that put before the flag was set is still released by the cleanup's cancel, as before.
  • merged keeps its one-slot size, so the backpressure is unchanged.

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added test_merge_live_event_streams_closes_while_merged_queue_is_full to tests/unittests/live/test__runner_utils.py. It fails 20 of 20 runs without the fix and passes 20 of 20 with it.

All results below are from environments built like CI (uv sync --extra test --no-install-package lancedb):

  • tests/unittests/live, tests/unittests/integrations/livekit, tests/unittests/runners and tests/unittests/test_runners.py: 415 passed, 1 skipped, 4 xfailed on Python 3.10, 3.11, 3.12, 3.13 and 3.14. Without the fix, one or both LiveKit run_live tests time out on each version.
  • Full tests/unittests on Python 3.12: 16389 passed, 1 failed. The one failure is test_import_loading.py::test_entry_point_loads_only_allowlisted_packages[agent], which already fails on main for an unrelated reason. main itself gives 16386 passed, 3 failed: that same test plus the two LiveKit tests fixed here.
  • mypy, compared the way CI does it (errors on the PR vs errors on main): no new errors.

Manual End-to-End (E2E) Tests:

  • The reproduction script from the issue now prints closed and exits. On main it hangs at closing....
  • A stress script ran 500 trials, each with random queued and agent events and closing at a random point. With the fix, no trial hung, no event was lost or duplicated, and no task was left behind. On main, 131 of 500 hung on Linux and 157 of 500 on Windows (Python 3.12).
  • I have not tested against a real LiveKit room with Gemini Live. The LiveKit tests that drive a real Runner through run_live are the closest check, and they pass.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

The change is small: 6 added lines and 1 changed line in src/google/adk/live/_runner_utils.py, plus one new test.

_merge_live_event_streams feeds a one-slot queue from two pumps. When the
caller closes the stream while that queue is full, cleanup cancels the
queued-events pump, whose finally block then waits forever to put the
done sentinel, since nothing reads the queue anymore. This hung both
LiveKit run_live tests on some machines.

Skip the sentinel once the consumer has stopped reading.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Closing a live event stream early can hang forever in _merge_live_event_streams

2 participants