feat: expose caught-up-to-tail signal on read sessions#88
Open
infiniteregrets wants to merge 1 commit into
Open
feat: expose caught-up-to-tail signal on read sessions#88infiniteregrets wants to merge 1 commit into
infiniteregrets wants to merge 1 commit into
Conversation
Greptile SummaryThis PR wraps the existing
Confidence Score: 4/5Safe to merge; the core logic is correct and matches the Go SDK caught-up semantics closely. The ReadSession implementation is carefully written — cancellation propagation, waiter lifecycle across reconnects, error normalization, and the delivery-consumed handshake all look correct. Both findings are non-blocking: _mark_behind() is called on every update including heartbeats, creating a transient reset invisible in practice due to single-threaded asyncio; and the close()-before-start path is correct but worth a clarifying comment for future maintainers. src/s2_sdk/_read_session.py — the _run() loop and close() early-exit path are the most complex sections and deserve the most attention on future changes. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant User
participant ReadSession
participant _run task
participant run_read_session
User->>ReadSession: "stream.read_session(start=...)"
Note over ReadSession: Stores generator ref, no task yet
User->>ReadSession: aenter / aiter / caught_up()
ReadSession->ReadSession: "_ensure_started() -> create_task(_run)"
_run task->>run_read_session: aiter(updates)
run_read_session-->>_run task: _ReadSessionUpdate(batch=None) [reconnect]
_run task->ReadSession: _mark_behind() -> is_caught_up=False
run_read_session-->>_run task: _ReadSessionUpdate(batch, caught_up_tail)
_run task->ReadSession: _mark_behind()
_run task->>ReadSession._deliveries: put(delivery)
_run task->>_run task: await delivery.consumed
User->>ReadSession: await anext()
ReadSession->>ReadSession._deliveries: "get() -> delivery"
ReadSession->ReadSession: "_mark_caught_up(tail) -> resolve waiters"
ReadSession->>_run task: delivery.consumed.set()
ReadSession-->>User: batch
User->>ReadSession: await caught_up()
ReadSession-->>User: StreamPosition (immediately if caught up)
User->>ReadSession: aclose() / aexit
ReadSession->>_run task: task.cancel()
_run task->>run_read_session: aclose()
_run task->ReadSession: _finish(error=None)
_run task->ReadSession: _signal_end() -> drain queue, put _END
ReadSession-->>User: (closed)
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant User
participant ReadSession
participant _run task
participant run_read_session
User->>ReadSession: "stream.read_session(start=...)"
Note over ReadSession: Stores generator ref, no task yet
User->>ReadSession: aenter / aiter / caught_up()
ReadSession->ReadSession: "_ensure_started() -> create_task(_run)"
_run task->>run_read_session: aiter(updates)
run_read_session-->>_run task: _ReadSessionUpdate(batch=None) [reconnect]
_run task->ReadSession: _mark_behind() -> is_caught_up=False
run_read_session-->>_run task: _ReadSessionUpdate(batch, caught_up_tail)
_run task->ReadSession: _mark_behind()
_run task->>ReadSession._deliveries: put(delivery)
_run task->>_run task: await delivery.consumed
User->>ReadSession: await anext()
ReadSession->>ReadSession._deliveries: "get() -> delivery"
ReadSession->ReadSession: "_mark_caught_up(tail) -> resolve waiters"
ReadSession->>_run task: delivery.consumed.set()
ReadSession-->>User: batch
User->>ReadSession: await caught_up()
ReadSession-->>User: StreamPosition (immediately if caught up)
User->>ReadSession: aclose() / aexit
ReadSession->>_run task: task.cancel()
_run task->>run_read_session: aclose()
_run task->ReadSession: _finish(error=None)
_run task->ReadSession: _signal_end() -> drain queue, put _END
ReadSession-->>User: (closed)
Prompt To Fix All With AIFix the following 2 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 2
src/s2_sdk/_read_session.py:155-166
**`_mark_behind()` called unconditionally, including for caught-up heartbeats**
Every iteration of the inner loop calls `_mark_behind()` before inspecting the update. For heartbeat updates — an empty batch with a non-`None` `caught_up_tail` — this causes the sequence `_mark_behind()` → `_mark_caught_up()` within the same event-loop tick. The Go SDK's equivalent `setBehind()` is [skipped for pure heartbeats](https://github.com/s2-streamstore/s2-sdk-go/blob/HEAD/s2/read.go#L566-L568) (`!caughtUp || hasRecords`). Because asyncio is single-threaded, the transient `False` state is never observable through `is_caught_up()`, but the docstring's stated semantics ("A later batch that does not reach a reported tail or a reconnect resets it") do not precisely match the code: every batch, including heartbeats with a tail, momentarily resets the flag. This could silently break if the invariant is ever relied upon in a future refactor.
### Issue 2 of 2
src/s2_sdk/_read_session.py:93-113
**`close()` with no task skips `_is_caught_up` reset on error path**
When `close()` is called before iteration starts (`_task is None`), `_finish()` is called with no argument, which does NOT clear `_is_caught_up`. Since the task never ran this is always `False` today, so no observable issue exists. More noteworthy: `_finish` sets `rejection = error or _session_closed_error()` where `error` is `None`, so subsequent calls to `caught_up()` correctly raise `ReadSessionClosedError`. The path is correct but a brief comment explaining why `_finish()` is called without an error argument here would help future maintainers.
Reviews (1): Last reviewed commit: "feat: expose caught-up-to-tail signal on..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+155
to
+166
| async def _run(self) -> None: | ||
| error: BaseException | None = None | ||
| cancelled = False | ||
| updates = aiter(self._updates) | ||
| try: | ||
| async for update in updates: | ||
| self._mark_behind() | ||
| batch = update.batch | ||
| if batch is None or not batch.records: | ||
| if update.caught_up_tail is not None: | ||
| self._mark_caught_up(update.caught_up_tail) | ||
| continue |
There was a problem hiding this comment.
_mark_behind() called unconditionally, including for caught-up heartbeats
Every iteration of the inner loop calls _mark_behind() before inspecting the update. For heartbeat updates — an empty batch with a non-None caught_up_tail — this causes the sequence _mark_behind() → _mark_caught_up() within the same event-loop tick. The Go SDK's equivalent setBehind() is skipped for pure heartbeats (!caughtUp || hasRecords). Because asyncio is single-threaded, the transient False state is never observable through is_caught_up(), but the docstring's stated semantics ("A later batch that does not reach a reported tail or a reconnect resets it") do not precisely match the code: every batch, including heartbeats with a tail, momentarily resets the flag. This could silently break if the invariant is ever relied upon in a future refactor.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/s2_sdk/_read_session.py
Line: 155-166
Comment:
**`_mark_behind()` called unconditionally, including for caught-up heartbeats**
Every iteration of the inner loop calls `_mark_behind()` before inspecting the update. For heartbeat updates — an empty batch with a non-`None` `caught_up_tail` — this causes the sequence `_mark_behind()` → `_mark_caught_up()` within the same event-loop tick. The Go SDK's equivalent `setBehind()` is [skipped for pure heartbeats](https://github.com/s2-streamstore/s2-sdk-go/blob/HEAD/s2/read.go#L566-L568) (`!caughtUp || hasRecords`). Because asyncio is single-threaded, the transient `False` state is never observable through `is_caught_up()`, but the docstring's stated semantics ("A later batch that does not reach a reported tail or a reconnect resets it") do not precisely match the code: every batch, including heartbeats with a tail, momentarily resets the flag. This could silently break if the invariant is ever relied upon in a future refactor.
How can I resolve this? If you propose a fix, please make it concise.
Comment on lines
+93
to
+113
| async def close(self) -> None: | ||
| """Close the session.""" | ||
| if self._closed: | ||
| return | ||
| if self._task is None: | ||
| self._finish() | ||
| self._signal_end() | ||
| return | ||
| current = asyncio.current_task() | ||
| cancellation_count = current.cancelling() if current is not None else 0 | ||
| self._task.cancel() | ||
| try: | ||
| await self._task | ||
| except asyncio.CancelledError: | ||
| pass | ||
| finally: | ||
| if not self._closed: | ||
| self._finish() | ||
| self._signal_end() | ||
| if current is not None and current.cancelling() > cancellation_count: | ||
| raise asyncio.CancelledError |
There was a problem hiding this comment.
close() with no task skips _is_caught_up reset on error path
When close() is called before iteration starts (_task is None), _finish() is called with no argument, which does NOT clear _is_caught_up. Since the task never ran this is always False today, so no observable issue exists. More noteworthy: _finish sets rejection = error or _session_closed_error() where error is None, so subsequent calls to caught_up() correctly raise ReadSessionClosedError. The path is correct but a brief comment explaining why _finish() is called without an error argument here would help future maintainers.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/s2_sdk/_read_session.py
Line: 93-113
Comment:
**`close()` with no task skips `_is_caught_up` reset on error path**
When `close()` is called before iteration starts (`_task is None`), `_finish()` is called with no argument, which does NOT clear `_is_caught_up`. Since the task never ran this is always `False` today, so no observable issue exists. More noteworthy: `_finish` sets `rejection = error or _session_closed_error()` where `error` is `None`, so subsequent calls to `caught_up()` correctly raise `ReadSessionClosedError`. The path is correct but a brief comment explaining why `_finish()` is called without an error argument here would help future maintainers.
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.