-
Notifications
You must be signed in to change notification settings - Fork 9
fix(streaming): StreamTaskMessageFull closes the coalescing buffer #426
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
a75aa68
2896295
b462ec0
87a6ad0
6772a40
61a507e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -420,15 +420,17 @@ async def close(self) -> TaskMessage: | |
| if not self.task_message: | ||
| raise ValueError("Context not properly initialized - no task message") | ||
|
|
||
| if self._is_closed: | ||
| return self.task_message # Already done | ||
|
|
||
| # Drain any buffered deltas before announcing DONE so consumers see the | ||
| # full sequence in order. | ||
| # Reap the buffer (stopping its ticker) before the _is_closed | ||
| # short-circuit, so a context already marked done by a Full update can't | ||
| # leave the ticker orphaned. Draining here also lets consumers see the | ||
| # full delta sequence in order before DONE. | ||
| if self._buffer is not None: | ||
| await self._buffer.close() | ||
| self._buffer = None | ||
|
|
||
| if self._is_closed: | ||
| return self.task_message # Already done (buffer reaped above) | ||
|
|
||
| # Send the DONE event | ||
| done_event = StreamTaskMessageDone( | ||
|
greptile-apps[bot] marked this conversation as resolved.
|
||
| parent_task_message=self.task_message, | ||
|
|
@@ -486,6 +488,15 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate | | |
| await self._buffer.add(update) | ||
| return update | ||
|
|
||
| # A Full ends the stream and supersedes buffered deltas. Drain and stop | ||
| # the buffer BEFORE publishing the Full, so leftover deltas land in order | ||
| # (deltas -> Full) instead of trailing the terminal Full as a stale | ||
| # duplicate tail. This also stops the ticker, which would otherwise be | ||
| # orphaned when __aexit__'s close() short-circuits on _is_closed. | ||
| if isinstance(update, StreamTaskMessageFull) and self._buffer is not None: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I saw that Greptile also commented about this and looks like it resolved its own comment but it looks to me like it's still possible that we could have a race here? e.g.
In this case, consumers will still see a delta after full since it just falls through to publishing directly. If this is something we care about (and my logic checks out and you also believe there's a race), we could have a terminal-in-progress flag like _is_closing set before awaiting buffer close to reject deltas?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good point, started introducing |
||
| await self._buffer.close() | ||
| self._buffer = None | ||
|
greptile-apps[bot] marked this conversation as resolved.
Outdated
|
||
|
|
||
| result = await self._streaming_service.stream_update(update) | ||
|
eberki-scale marked this conversation as resolved.
|
||
|
|
||
| if isinstance(update, StreamTaskMessageDone): | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.