fix(client-runtime): recover subscriptions that stop on a healthy connection

packages/client-runtime/src/rpc/client.test.ts

@@ -13,8 +13,10 @@ import * as Fiber from "effect/Fiber";
 import * as Option from "effect/Option";
 import * as Queue from "effect/Queue";
 import * as Ref from "effect/Ref";
+import * as Schema from "effect/Schema";
 import * as Stream from "effect/Stream";
 import * as SubscriptionRef from "effect/SubscriptionRef";
+import * as TestClock from "effect/testing/TestClock";
 import { RpcClientError } from "effect/rpc";
 
 import {
@@ -44,6 +46,8 @@ const INSTALL_DOWNLOADING: RelayClientInstallProgressEvent = {
   stage: "downloading",
 };
 
+const decodeNumber = Schema.decodeUnknownEffect(Schema.Number);
+
 function session(client: WsRpcProtocolClient): RpcSession.RpcSession {
   return {
     client,
@@ -258,13 +262,93 @@ describe("environment RPC", () => {
       for (let attempt = 0; attempt < 100 && subscriptions.length < 2; attempt += 1) {
         yield* Effect.yieldNow;
       }
+      // The pending same-session retry must not resubscribe the replaced session.
+      yield* TestClock.adjust("1 minute");
       yield* Fiber.interrupt(subscriptionFiber);
 
       expect(subscriptions).toEqual(["first", "second"]);
       expect(yield* Ref.get(retryCount)).toBe(0);
     }),
   );
 
+  // The session stays healthy (`closed` never resolves), as when one stream chunk fails to decode,
+  // so the supervisor will not replace it and only resubscribing restores delivery.
+  const resumesOnTheSameSession = Effect.fn("TestEnvironmentRpc.resumesOnTheSameSession")(
+    function* (streamFailure: Stream.Stream<never, unknown>) {
+      const resumedValue = { source: "resumed" } as unknown as ServerLifecycleStreamEvent;
+      const firstSubscribed = yield* Deferred.make<void>();
+      let subscriptions = 0;
+      const observedDefects: unknown[] = [];
+      const client = {
+        [WS_METHODS.subscribeServerLifecycle]: () => {
+          subscriptions += 1;
+          return subscriptions === 1
+            ? Stream.fromEffect(Deferred.succeed(firstSubscribed, undefined)).pipe(
+                Stream.drain,
+                Stream.concat(streamFailure),
+              )
+            : Stream.make(resumedValue).pipe(Stream.concat(Stream.never));
+        },
+      } as unknown as WsRpcProtocolClient;
+      const { activeSession, retryCount, supervisor } = yield* makeHarness();
+
+      yield* SubscriptionRef.set(activeSession, Option.some(session(client)));
+      const resultFiber = yield* subscribe(
+        WS_METHODS.subscribeServerLifecycle,
+        {},
+        {
+          onDefect: (cause) =>
+            Effect.sync(() => {
+              observedDefects.push(Cause.squash(cause));
+            }),
+        },
+      ).pipe(
+        Stream.runHead,
+        Effect.timeoutOption("2 seconds"),
+        Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
+        Effect.forkChild,
+      );
+      yield* Deferred.await(firstSubscribed);
+      yield* TestClock.adjust("2 seconds");
+
+      expect(
+        yield* Fiber.join(resultFiber),
+        "durable subscription stopped delivering after its stream ended on a healthy session",
+      ).toEqual(Option.some(Option.some(resumedValue)));
+      expect(subscriptions).toBe(2);
+      expect(yield* Ref.get(retryCount)).toBe(0);
+      return observedDefects;
+    },
+  );
+
+  it.effect("resubscribes on the same session after a transport failure", () =>
+    Effect.gen(function* () {
+      yield* resumesOnTheSameSession(
+        Stream.fail(
+          new RpcClientError.RpcClientError({
+            reason: new RpcClientError.RpcClientDefect({
+              message: "stream failed",
+              cause: new Error("stream failed"),
+            }),
+          }),
+        ),
+      );
+    }),
+  );
+
+  it.effect("resubscribes on the same session after a stream chunk fails to decode", () =>
+    Effect.gen(function* () {
+      // Mirrors RpcClient, which decodes each chunk with `Effect.orDie`.
+      const decodeDefect = decodeNumber("not a number").pipe(Effect.orDie);
+      const observedDefects = yield* resumesOnTheSameSession(
+        Stream.fromEffect(decodeDefect).pipe(Stream.drain),
+      );
+
+      expect(observedDefects).toHaveLength(1);
+      expect(Schema.isSchemaError(observedDefects[0])).toBe(true);
+    }),
+  );
+
   it.effect("surfaces domain subscription failures without reconnecting", () =>
     Effect.gen(function* () {
       const domainError = new Error("terminal subscription rejected");

packages/client-runtime/src/rpc/client.ts

@@ -107,6 +107,7 @@ export const isRpcClientError = Schema.is(RpcClientError.RpcClientError);
 const isEnvironmentAuthorizationError = Schema.is(EnvironmentAuthorizationError);
 
 const MAX_EXPECTED_FAILURE_RETRY_DELAY_MS = 30_000;
+const STREAM_FAILURE_RETRY_DELAY = Duration.seconds(1);
 
 export type EnvironmentRpcInput<TTag extends EnvironmentRpcTag> = Parameters<RpcMethod<TTag>>[0];
 
@@ -324,6 +325,19 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
                         : Effect.void,
                     ),
                     Stream.catchCause((cause) => {
+                      const retryAfter = (base: Duration.Input) => {
+                        const retryDelay = Duration.millis(
+                          Math.min(
+                            Duration.toMillis(base) * 2 ** expectedFailureRetries,
+                            MAX_EXPECTED_FAILURE_RETRY_DELAY_MS,
+                          ),
+                        );
+                        expectedFailureRetries += 1;
+                        return Stream.fromEffect(Effect.sleep(retryDelay)).pipe(
+                          Stream.drain,
+                          Stream.concat(subscribeToSession()),
+                        );
+                      };
                       const hasOnlyExpectedFailures =
                         cause.reasons.length > 0 &&
                         cause.reasons.every((reason) => reason._tag === "Fail");
@@ -332,17 +346,34 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
                         cause.reasons.every(
                           (reason) => reason._tag === "Fail" && isRpcClientError(reason.error),
                         );
-                      if (isTransportFailure) {
+                      // The RPC client dies with a SchemaError when one stream chunk fails to
+                      // decode, ending only that request while the socket stays open.
+                      const isChunkDecodeDefect =
+                        cause.reasons.length > 0 &&
+                        cause.reasons.every(
+                          (reason) => reason._tag === "Die" && Schema.isSchemaError(reason.defect),
+                        );
+                      // Both end one stream while the session can stay healthy, so nothing else
+                      // would replace it. A dead session is replaced by the supervisor, which
+                      // interrupts this retry through switchMap.
+                      if (isTransportFailure || isChunkDecodeDefect) {
                         return Stream.fromEffect(
                           Effect.logWarning(
-                            "Durable RPC subscription lost its transport; waiting for the next session.",
+                            "Durable RPC subscription stream ended; resubscribing.",
                             {
                               cause: Cause.pretty(cause),
                               method: tag,
                               environmentId: supervisor.target.environmentId,
                             },
                           ),
-                        ).pipe(Stream.drain);
+                        ).pipe(
+                          Stream.drain,
+                          Stream.concat(
+                            retryAfter(
+                              options?.retryExpectedFailureAfter ?? STREAM_FAILURE_RETRY_DELAY,
+                            ),
+                          ),
+                        );
                       }
                       if (hasOnlyExpectedFailures && options?.onExpectedFailure !== undefined) {
                         const handled = Stream.fromEffect(options.onExpectedFailure(cause)).pipe(
@@ -358,19 +389,8 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
                         ) {
                           return handled;
                         }
-                        const retryDelay = Duration.millis(
-                          Math.min(
-                            Duration.toMillis(options.retryExpectedFailureAfter) *
-                              2 ** expectedFailureRetries,
-                            MAX_EXPECTED_FAILURE_RETRY_DELAY_MS,
-                          ),
-                        );
-                        expectedFailureRetries += 1;
                         return handled.pipe(
-                          Stream.concat(
-                            Stream.fromEffect(Effect.sleep(retryDelay)).pipe(Stream.drain),
-                          ),
-                          Stream.concat(subscribeToSession()),
+                          Stream.concat(retryAfter(options.retryExpectedFailureAfter)),
                         );
                       }
                       return Stream.failCause(cause);