Repository: thijsmie/pocketbase

Issue:
Bug: Realtime connection timeout after inactivity in PocketBase
# Description
When I use the `subscribe` method of the package and run it in an infinite loop, if no modifications are made to the specified Collection/record_id within 10 minutes, I encounter the following error: `httpx.RemoteProtocolError: peer closed connection without sending complete message body (incomplete chunked read)`.

# Steps to reproduce
Using this test code:
```py
from pocketbase import PocketBase
from pocketbase.services.realtime import RealtimeEvent
from datetime import datetime
import asyncio


CONNECTION_URL = "CONNECTION_URL"
ADMIN_EMAIL = "ADMIN_EMAIL"
ADMIN_PASSWORD = "ADMIN_PASSWORD"
COLLECTION_NAME = "COLLECTION_NAME"


async def callback(event: RealtimeEvent) -> None:
    """Callback function for handling Realtime events.

    Args:
        event (RealtimeEvent): The event object containing information about the record change.
    """

    at = datetime.now().isoformat()
    print(f"[{at}] {event['action'].upper()}: {event['record']}")


async def realtime_updates():
    """Establishes a PocketBase connection, authenticates, and subscribes to Realtime events."""

    try:
        # Instantiate the PocketBase connector
        pb = PocketBase(CONNECTION_URL)

        # Authenticate as an admin
        await pb.admins.auth.with_password(ADMIN_EMAIL, ADMIN_PASSWORD)

        # Get the collection object
        col = pb.collection(COLLECTION_NAME)

        # Subscribe to Realtime events for the specific record ID
        unsub = await col.subscribe(
            callback=callback,
            record_id="*"
            )

        # Infinite loop to wait for events (adjusted from the second snippet)
        while True:
            await asyncio.sleep(10)  # Sleep for an hour to avoid hitting PocketBase's rate limits

    except Exception as e:
        print(f"Error: {e}")

    finally:
        # Unsubscribe if still active
        if unsub:
            try:
                await unsub()
            except Exception as e:
                print(f"Error unsubscribing: {e}")


if __name__ == "__main__":
    asyncio.run(realtime_updates())
```

I run the code and i wait for 5-10 mins.

# Expected behavior
It should not Close the connection

# [Optional] Stacktrace
```
Task exception was never retrieved
future: <Task finished name='Task-3' coro=<RealtimeService._make_connection() done, defined at /mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/pocketbase/services/realtime.py:48> exception=RemoteProtocolError('peer closed connection without sending complete message body (incomplete chunked read)')>
Traceback (most recent call last):
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_transports/default.py", line 69, in map_httpcore_exceptions
    yield
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_transports/default.py", line 254, in __aiter__
    async for part in self._httpcore_stream:
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_async/connection_pool.py", line 367, in __aiter__
    raise exc from None
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_async/connection_pool.py", line 363, in __aiter__
    async for part in self._stream:
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_async/http11.py", line 349, in __aiter__
    raise exc
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_async/http11.py", line 341, in __aiter__
    async for chunk in self._connection._receive_response_body(**kwargs):
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_async/http11.py", line 210, in _receive_response_body
    event = await self._receive_event(timeout=timeout)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_async/http11.py", line 220, in _receive_event
    with map_exceptions({h11.RemoteProtocolError: RemoteProtocolError}):
  File "/usr/lib/python3.11/contextlib.py", line 155, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpcore/_exceptions.py", line 14, in map_exceptions
    raise to_exc(exc) from exc
httpcore.RemoteProtocolError: peer closed connection without sending complete message body (incomplete chunked read)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/lib/python3.11/contextlib.py", line 222, in __aexit__
    await self.gen.athrow(typ, value, traceback)
  File "/usr/lib/python3.11/contextlib.py", line 222, in __aexit__
    await self.gen.athrow(typ, value, traceback)
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_client.py", line 1624, in stream
    yield response
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx_sse/_api.py", line 70, in aconnect_sse
    yield EventSource(response)
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/pocketbase/services/realtime.py", line 51, in _make_connection
    async for message in sse.aiter_sse():
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx_sse/_api.py", line 39, in aiter_sse
    async for line in self._response.aiter_lines():
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_models.py", line 963, in aiter_lines
    async for text in self.aiter_text():
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_models.py", line 950, in aiter_text
    async for byte_content in self.aiter_bytes():
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_models.py", line 929, in aiter_bytes
    async for raw_bytes in self.aiter_raw():
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_models.py", line 987, in aiter_raw
    async for raw_stream_bytes in self.stream:
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_client.py", line 149, in __aiter__
    async for chunk in self._stream:
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_transports/default.py", line 253, in __aiter__
    with map_httpcore_exceptions():
  File "/usr/lib/python3.11/contextlib.py", line 155, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/mnt/d/Documents/Projets/myProject/venv/lib/python3.11/site-packages/httpx/_transports/default.py", line 86, in map_httpcore_exceptions
    raise mapped_exc(message) from exc
httpx.RemoteProtocolError: peer closed connection without sending complete message body (incomplete chunked read)
```

# Environment
 - Operating system: `WSL Version 2`
 - Python version: `Python 3.11.0rc1`
 - Upstream PocketBase version: `PocketBase v0.21.3`
 - pocketbase-async version: `0.8.3`


Discussion:

Patch:
diff --git a/examples/realtime_updates.py b/examples/realtime_updates.py
index e27022c..3d919d0 100644
--- a/examples/realtime_updates.py
+++ b/examples/realtime_updates.py
@@ -2,7 +2,7 @@ import asyncio
 from datetime import datetime
 
 from pocketbase import PocketBase
-from pocketbase.services.realtime import RealtimeEvent
+from pocketbase.models.dtos import RealtimeEvent
 
 CONNECTION_URL = "http://localhost:8123"
 ADMIN_EMAIL = "test@example.com"
@@ -24,6 +24,7 @@ async def callback(event: RealtimeEvent) -> None:
 
 async def realtime_updates():
     """Establishes a PocketBase connection, authenticates, and subscribes to Realtime events."""
+    unsubscribe = None
 
     try:
         # Instantiate the PocketBase connector
diff --git a/src/pocketbase/models/dtos.py b/src/pocketbase/models/dtos.py
index 23b1bf7..bdd440a 100644
--- a/src/pocketbase/models/dtos.py
+++ b/src/pocketbase/models/dtos.py
@@ -1,4 +1,4 @@
-from typing import Generic, NotRequired, TypedDict, TypeVar
+from typing import Generic, Literal, NotRequired, TypedDict, TypeVar
 
 from pocketbase.utils.types import JsonType
 
@@ -127,3 +127,8 @@ class AuthResult(TypedDict):
 class AdminAuthResult(TypedDict):
     token: str
     admin: AdminModel
+
+
+class RealtimeEvent(TypedDict):
+    action: Literal["create", "update", "delete"]
+    record: Record
diff --git a/src/pocketbase/services/realtime.py b/src/pocketbase/services/realtime.py
index 6223431..09586a1 100644
--- a/src/pocketbase/services/realtime.py
+++ b/src/pocketbase/services/realtime.py
@@ -5,12 +5,13 @@ from asyncio.tasks import Task
 from collections import defaultdict
 from collections.abc import Awaitable, Callable
 from contextlib import suppress
-from typing import TYPE_CHECKING, Literal, TypedDict
+from typing import TYPE_CHECKING, Any
 from urllib.parse import quote
 
+from httpx import ReadError, RemoteProtocolError
 from httpx_sse import aconnect_sse
 
-from pocketbase.models.dtos import Record
+from pocketbase.models.dtos import RealtimeEvent
 from pocketbase.models.options import CommonOptions
 from pocketbase.services.base import Service
 
@@ -18,11 +19,6 @@ if TYPE_CHECKING:
     from pocketbase.client import PocketBase, PocketBaseInners
 
 
-class RealtimeEvent(TypedDict):
-    action: Literal["create", "update", "delete"]
-    record: Record
-
-
 Callback = Callable[[RealtimeEvent], Awaitable[None]]
 
 
@@ -46,30 +42,46 @@ class RealtimeService(Service):
         await sentinel.wait()
 
     async def _make_connection(self, sentinel: asyncio.Event) -> None:
+        headers: dict[str, Any] = {}
+        last_event_id: Any | None = None
         try:
-            async with aconnect_sse(self._in.client, "GET", self.__base_sub_path__, timeout=900) as sse:
-                async for message in sse.aiter_sse():
-                    if message.event == "PB_CONNECT":
-                        self._client_id = message.id
-                        await self._transmit_subscriptions()
-                        sentinel.set()
-                    elif message.event in self._subscriptions:
-                        for callback in self._subscriptions[message.event]:
-                            try:
-                                await callback(message.json())
-                            except:  # noqa: E722
-                                # We never want any exception to break the realtime handler.
-                                logging.exception("Unhandled exception in realtime event handler")
+            while True:
+                try:
+                    async with aconnect_sse(
+                        self._in.client, "GET", self.__base_sub_path__, headers=headers, timeout=900
+                    ) as sse:
+                        async for message in sse.aiter_sse():
+                            if message.event == "PB_CONNECT":
+                                self._client_id = message.id
+                                await self._transmit_subscriptions(force=True)
+                                sentinel.set()
+                                continue
+
+                            last_event_id = message.id
+
+                            if message.event in self._subscriptions:
+                                for callback in self._subscriptions[message.event]:
+                                    try:
+                                        await callback(message.json())
+                                    except:  # noqa: E722
+                                        # We never want any exception to break the realtime handler.
+                                        logging.exception("Unhandled exception in realtime event handler")
+                except (ReadError, asyncio.TimeoutError, RemoteProtocolError):
+                    logging.debug("Connection lost, reconnecting automatically")
+                    if last_event_id:
+                        headers["Last-Event-ID"] = last_event_id
+
         finally:
             # Disconnected, reset to plain state.
+            logging.exception("Connection to realtime endpoint lost")
             self._connection = None
             self._client_id = None
             self._last_transmit = set()
             sentinel.set()
 
-    async def _transmit_subscriptions(self) -> None:
+    async def _transmit_subscriptions(self, force: bool = False) -> None:
         to_transmit = set(self._subscriptions.keys())
-        if to_transmit == self._last_transmit or not self._client_id:
+        if not force and (to_transmit == self._last_transmit or not self._client_id):
             return
 
         self._last_transmit = to_transmit


Test patch:
diff --git a/tests/test_realtime.py b/tests/test_realtime.py
index a19f6e0..a343442 100644
--- a/tests/test_realtime.py
+++ b/tests/test_realtime.py
@@ -1,11 +1,17 @@
 import asyncio
+import json
+from contextlib import asynccontextmanager
+from unittest.mock import MagicMock, patch
+from urllib.parse import quote
+from uuid import uuid4
 
+import httpx
+from httpx_sse import ServerSentEvent
 from pocketbase import PocketBase
-from pocketbase.models.dtos import CollectionModel
 from pocketbase.services.realtime import RealtimeEvent
 
 
-async def test_realtime(admin_client: PocketBase) -> CollectionModel:
+async def test_realtime(admin_client: PocketBase) -> None:
     await admin_client.collections.create(
         {
             "name": "test",
@@ -51,3 +57,68 @@ async def test_realtime(admin_client: PocketBase) -> CollectionModel:
     assert event_payload["action"] == "update"
     assert event_payload["record"] == record
     await unsub()
+
+
+# The test below takes five minutes at least (until we get disconnected)
+# so we don't run it normally, you'll have to turn it on manually
+# Feature request, configurable timeout for realtime connections?
+
+
+class MockAsyncIteratorCallable:
+    def __init__(self, generator_function):
+        self.generator_function = generator_function
+        self.case = -1
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            self.case += 1
+            return await self.generator_function(self.case)
+        except StopIteration:
+            raise StopAsyncIteration
+
+
+async def test_realtime_all_records(admin_client: PocketBase) -> None:
+    def handler(_r):
+        return httpx.Response(200, json={})
+
+    transport = httpx.MockTransport(handler)
+    pb = PocketBase("http://bla.com")
+    pb._inners.client = httpx.AsyncClient(base_url="http://bla.com", transport=transport)
+
+    col = pb.collection("test_realtime_disconnect_all_records")
+
+    event_counter = 0
+
+    async def callback(event: RealtimeEvent) -> None:
+        nonlocal event_counter
+        event_counter += 1
+
+    async def behaviour(case: int) -> ServerSentEvent:
+        match case:
+            case 0:
+                return ServerSentEvent(event="PB_CONNECT", data=str(uuid4()), id="something")
+            case 1:
+                return ServerSentEvent(
+                    event=quote("test_realtime_disconnect_all_records"),
+                    data=json.dumps({"action": "create", "record": {"title": "Hi"}}),
+                    id="something",
+                )
+            case 2:
+                await asyncio.sleep(0.02)
+                raise TimeoutError()
+
+    @asynccontextmanager
+    async def repl_aconnect_sse(*args, **kwargs):
+        mock = MagicMock()
+        mock.aiter_sse.return_value = MockAsyncIteratorCallable(behaviour)
+        yield mock
+
+    with patch("pocketbase.services.realtime.aconnect_sse", repl_aconnect_sse):
+        unsub = await col.subscribe_all(callback=callback)
+        await asyncio.sleep(0.1)
+        await unsub()
+
+    assert event_counter > 1