feat/river schema v2 compatibility

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "nanoid>=2.0.0",
     "protobuf>=5.28.3",
     "pydantic-core>=2.20.1",
-    "websockets>=12.0",
+    "websockets>=13.0,<14",
     "opentelemetry-sdk>=1.28.2",
     "opentelemetry-api>=1.28.2",
 ]

src/replit_river/__init__.py

@@ -1,3 +1,4 @@
+from . import v2
 from .client import Client
 from .error_schema import RiverError
 from .rpc import (
@@ -20,4 +21,5 @@
     "subscription_method_handler",
     "upload_method_handler",
     "stream_method_handler",
+    "v2",
 ]

src/replit_river/client.py

@@ -226,9 +226,11 @@ def _trace_procedure(
     span_handle = _SpanHandle(span)
     try:
         yield span_handle
+        span_handle.set_status(StatusCode.OK)
     except GeneratorExit:
         # This error indicates the caller is done with the async generator
         # but messages are still left. This is okay, we do not consider it an error.
+        span_handle.set_status(StatusCode.OK)
         raise
     except RiverException as e:
         span.record_exception(e, escaped=True)
@@ -239,7 +241,6 @@ def _trace_procedure(
         span_handle.set_status(StatusCode.ERROR, f"{type(e).__name__}: {e}")
         raise e
     finally:
-        span_handle.set_status(StatusCode.OK)
         span.end()