fix(client-runtime): a resumed thread renders its catch-up once instead of stepping through stale states

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

@@ -281,6 +281,7 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
   makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
   mapStream: (
     session: RpcSession,
+    input: EnvironmentRpcInput<TTag>,
     stream: Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>,
   ) => Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>>,
   options?: SubscriptionOptions<TTag>,
@@ -333,6 +334,7 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
                       });
                       const stream = mapStream(
                         session,
+                        input,
                         Stream.suspend(() => method(input)),
                       ).pipe(
                         Stream.onFirst(() =>
@@ -467,7 +469,7 @@ export function subscribeDynamic<TTag extends EnvironmentSubscriptionRpcTag>(
   EnvironmentRpcStreamFailure<TTag>,
   EnvironmentSupervisor.EnvironmentSupervisor
 > {
-  return subscribeDynamicMapped(tag, makeInput, (_session, stream) => stream, options);
+  return subscribeDynamicMapped(tag, makeInput, (_session, _input, stream) => stream, options);
 }
 
 /** Tags each value before `switchMap` can buffer it across a session change. */
@@ -483,7 +485,29 @@ export function subscribeDynamicWithSession<TTag extends EnvironmentSubscription
   return subscribeDynamicMapped(
     tag,
     makeInput,
-    (session, stream) => stream.pipe(Stream.map((value) => [session, value] as const)),
+    (session, _input, stream) => stream.pipe(Stream.map((value) => [session, value] as const)),
+    options,
+  );
+}
+
+/**
+ * Tags each value with the input of the attempt that produced it. Every
+ * `makeInput` call starts an attempt, including a resubscribe on the same
+ * session; `makeInput` must return a fresh object so the input identifies it.
+ */
+export function subscribeDynamicWithInput<TTag extends EnvironmentSubscriptionRpcTag>(
+  tag: TTag,
+  makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
+  options?: SubscriptionOptions<TTag>,
+): Stream.Stream<
+  readonly [input: EnvironmentRpcInput<TTag>, value: EnvironmentRpcStreamValue<TTag>],
+  EnvironmentRpcStreamFailure<TTag>,
+  EnvironmentSupervisor.EnvironmentSupervisor
+> {
+  return subscribeDynamicMapped(
+    tag,
+    makeInput,
+    (_session, input, stream) => stream.pipe(Stream.map((value) => [input, value] as const)),
     options,
   );
 }
@@ -506,7 +530,7 @@ export const subscribeVcsBranch = (
   subscribeDynamicMapped(
     WS_METHODS.subscribeVcsBranch,
     () => Effect.succeed(input),
-    (session, stream) =>
+    (session, _input, stream) =>
       Stream.unwrap(
         session.initialConfig.pipe(
           Effect.mapError(

packages/client-runtime/src/state/threads-atoms.test.ts

@@ -304,6 +304,8 @@ describe("createEnvironmentThreadStateAtoms", () => {
       });
       const unmount = h.registry.mount(h.stateAtom);
       const first = yield* Queue.take(h.subscriptions);
+      yield* Queue.offer(first.events, { kind: "synchronized" });
+      yield* observeState(h.registry, h.stateAtom, (state) => state.status === "live");
       const loading = yield* h.loadEarlier().pipe(Effect.forkScoped);
       const older = yield* Queue.take(h.olderLoads);
       expect(h.registry.get(h.stateAtom).history.loading).toBe(true);

packages/client-runtime/src/state/threads-sync.test.ts

@@ -23,8 +23,8 @@ 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 Scope from "effect/Scope";
 import * as Scheduler from "effect/Scheduler";
+import * as Scope from "effect/Scope";
 import * as Schema from "effect/Schema";
 import * as Stream from "effect/Stream";
 import * as SubscriptionRef from "effect/SubscriptionRef";
@@ -133,6 +133,9 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
   readonly getTurnItem?: (input: {
     readonly itemId: TurnItemId;
   }) => Effect.Effect<{ readonly item: OrchestrationV2TurnItem | null }>;
+  readonly streamFor?: (
+    subscription: number,
+  ) => Stream.Stream<OrchestrationV2ThreadStreamItem> | undefined;
 }) {
   const inputs = yield* Queue.unbounded<TestThreadInput>();
   const observed = yield* Queue.unbounded<EnvironmentThreadState>();
@@ -168,17 +171,19 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
       readonly acceptCompactTurnItems?: true;
     }) =>
       Stream.unwrap(
-        Ref.set(lastSubscribeAfterSequence, input.afterSequence).pipe(
-          Effect.andThen(
-            Ref.set(lastRequestCompletionMarker, input.requestCompletionMarker === true),
-          ),
-          Effect.andThen(Ref.set(lastAcceptBoundedSnapshot, input.acceptBoundedSnapshot)),
-          Effect.andThen(Ref.set(lastAcceptCompactTurnItems, input.acceptCompactTurnItems)),
-          Effect.andThen(SubscriptionRef.update(subscriptionCount, (count) => count + 1)),
-          Effect.as(
-            options?.completionMarker === true
-              ? streamFrom(inputs)
-              : Stream.concat(Stream.make(synchronized()), streamFrom(inputs)),
+        Effect.all([
+          Ref.set(lastSubscribeAfterSequence, input.afterSequence),
+          Ref.set(lastRequestCompletionMarker, input.requestCompletionMarker === true),
+          Ref.set(lastAcceptBoundedSnapshot, input.acceptBoundedSnapshot),
+          Ref.set(lastAcceptCompactTurnItems, input.acceptCompactTurnItems),
+        ]).pipe(
+          Effect.andThen(SubscriptionRef.updateAndGet(subscriptionCount, (count) => count + 1)),
+          Effect.map(
+            (count) =>
+              options?.streamFor?.(count) ??
+              (options?.completionMarker === true
+                ? streamFrom(inputs)
+                : Stream.concat(Stream.make(synchronized()), streamFrom(inputs))),
           ),
         ),
       ),
@@ -316,6 +321,15 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
   };
 });
 
+const awaitSubscriptionCount = (
+  harness: { readonly subscriptionCount: SubscriptionRef.SubscriptionRef<number> },
+  count: number,
+) =>
+  SubscriptionRef.changes(harness.subscriptionCount).pipe(
+    Stream.filter((current) => current >= count),
+    Stream.runHead,
+  );
+
 const snapshot = (
   projection: OrchestrationV2ThreadProjection,
   snapshotSequence = 1,
@@ -2489,34 +2503,30 @@ describe("EnvironmentThreads", () => {
       }),
   );
 
-  it.effect("keeps replayed updates synchronizing until the completion marker arrives", () =>
+  it.effect("publishes a 20-event catch-up as one state change at its marker", () =>
     Effect.gen(function* () {
       const harness = yield* makeHarness({ cached: BASE_PROJECTION, completionMarker: true });
-      yield* awaitThreadState(
-        harness.observed,
-        (value) => value.status === "synchronizing" && Option.isSome(value.data),
-      );
+      yield* awaitSubscriptionCount(harness, 1);
       expect(yield* Ref.get(harness.lastRequestCompletionMarker)).toBe(true);
-
-      yield* Queue.offer(
-        harness.inputs,
-        titleUpdated("Caught-up title", CACHED_SNAPSHOT_SEQUENCE + 1),
-      );
-      const catchingUp = yield* awaitThreadState(
-        harness.observed,
-        (value) =>
-          value.status === "synchronizing" &&
-          Option.isSome(value.data) &&
-          value.data.value.thread.title === "Caught-up title",
-      );
-      expect(catchingUp.status).toBe("synchronizing");
-
+      // The observer replays the current state when it starts; take it so clearing drops it.
+      yield* Queue.take(harness.observed);
+      yield* Queue.clear(harness.observed);
+      // Separate frames, as the server sends a replay.
+      for (let index = 1; index <= 20; index += 1) {
+        yield* Queue.offer(
+          harness.inputs,
+          titleUpdated(`Title ${index}`, CACHED_SNAPSHOT_SEQUENCE + index),
+        );
+        yield* Effect.yieldNow;
+      }
       yield* Queue.offer(harness.inputs, synchronized());
-      const live = yield* awaitThreadState(
-        harness.observed,
-        (value) => value.status === "live" && Option.isSome(value.data),
-      );
-      expect(Option.getOrThrow(live.data).thread.title).toBe("Caught-up title");
+      const live = yield* Queue.take(harness.observed);
+      expect(live.status).toBe("live");
+      expect(Option.getOrThrow(live.data).thread.title).toBe("Title 20");
+
+      yield* Queue.offer(harness.inputs, titleUpdated("Live title", CACHED_SNAPSHOT_SEQUENCE + 21));
+      const next = yield* Queue.take(harness.observed);
+      expect(Option.getOrThrow(next.data).thread.title).toBe("Live title");
     }),
   );
 
@@ -2552,6 +2562,107 @@ describe("EnvironmentThreads", () => {
     }),
   );
 
+  // A long batch yields to the scheduler while it applies, so the next attempt
+  // can start while the replaced attempt's batch and items are still landing.
+  const LONG_BATCH = 1000;
+
+  it.effect(
+    "does not apply a replaced attempt's snapshot after the next attempt reads its cursor",
+    () =>
+      Effect.gen(function* () {
+        const sent = yield* Deferred.make<void>();
+        const harness = yield* makeHarness({
+          cached: BASE_PROJECTION,
+          streamFor: (subscription) =>
+            subscription === 1
+              ? Stream.fromArray([
+                  synchronized(),
+                  snapshot(BASE_PROJECTION, 8),
+                  ...Array.from({ length: LONG_BATCH }, (_, index) =>
+                    titleUpdated(`Title ${index + 9}`, index + 9),
+                  ),
+                  snapshot({
+                    ...BASE_PROJECTION,
+                    thread: { ...BASE_PROJECTION.thread, title: "Stale" },
+                  }),
+                ]).pipe(
+                  Stream.rechunk(1),
+                  Stream.concat(
+                    Stream.fromEffect(Deferred.succeed(sent, undefined)).pipe(Stream.drain),
+                  ),
+                  Stream.concat(Stream.never),
+                )
+              : undefined,
+        });
+        yield* Deferred.await(sent);
+        yield* Queue.offer(harness.wakeups, "application-active");
+        yield* awaitSubscriptionCount(harness, 2);
+        yield* Queue.offer(harness.inputs, titleUpdated("Live", LONG_BATCH + 9));
+        const titles: Array<string | undefined> = [];
+        yield* awaitThreadState(harness.observed, (value) => {
+          titles.push(Option.getOrNull(value.data)?.thread.title);
+          return Option.getOrNull(value.data)?.thread.title === "Live";
+        });
+        // The next attempt started with the snapshot still buffered: its cursor
+        // is past the replaced snapshot's sequence but short of the batch end.
+        const cursor = yield* Ref.get(harness.lastSubscribeAfterSequence);
+        expect(cursor).toBeGreaterThan(8);
+        expect(cursor).toBeLessThan(LONG_BATCH + 8);
+        expect(titles).not.toContain("Stale");
+      }).pipe(
+        // A small budget lets the next attempt start between the consumer's
+        // batches, before it takes the queued snapshot.
+        Effect.provideService(Scheduler.MaxOpsBeforeYield, 50),
+      ),
+  );
+
+  it.effect("does not apply a replaced attempt's late marker without its dropped prefix", () =>
+    Effect.gen(function* () {
+      const sent = yield* Deferred.make<void>();
+      const replay = [
+        questionResolved(CACHED_SNAPSHOT_SEQUENCE + 1),
+        ...Array.from({ length: 40 }, (_, index) =>
+          titleUpdated(`Title ${index}`, CACHED_SNAPSHOT_SEQUENCE + 2 + index),
+        ),
+        synchronized(),
+      ];
+      const harness = yield* makeHarness({
+        cached: PROJECTION_WITH_QUESTION,
+        completionMarker: true,
+        streamFor: (subscription) =>
+          subscription === 1
+            ? Stream.fromArray(replay).pipe(
+                Stream.rechunk(1),
+                Stream.concat(
+                  Stream.fromEffect(Deferred.succeed(sent, undefined)).pipe(Stream.drain),
+                ),
+                Stream.concat(Stream.never),
+              )
+            : undefined,
+      });
+      yield* Deferred.await(sent);
+      yield* Queue.offer(harness.wakeups, "application-active");
+      yield* awaitSubscriptionCount(harness, 2);
+      // The next attempt started before the replaced attempt's buffered marker landed.
+      expect(yield* Ref.get(harness.lastSubscribeAfterSequence)).toBe(CACHED_SNAPSHOT_SEQUENCE);
+      yield* Queue.offerAll(harness.inputs, [
+        ...replay,
+        titleUpdated("Live", CACHED_SNAPSHOT_SEQUENCE + 42),
+      ]);
+      const live = yield* awaitThreadState(
+        harness.observed,
+        (value) => Option.getOrNull(value.data)?.thread.title === "Live",
+      );
+      expect(Option.getOrThrow(live.data).runtimeRequests.map((request) => request.status)).toEqual(
+        ["resolved"],
+      );
+    }).pipe(
+      // The replaced attempt's items fill the switch buffer, and a small budget
+      // lets the next attempt start while the consumer still drains them.
+      Effect.provideService(Scheduler.MaxOpsBeforeYield, 5),
+    ),
+  );
+
   it.effect("skips an unknown event type and resumes after it", () =>
     Effect.gen(function* () {
       const harness = yield* makeHarness({ cached: BASE_PROJECTION, completionMarker: true });