diff --git a/packages/client-runtime/src/rpc/client.ts b/packages/client-runtime/src/rpc/client.ts index ee70400ca47a..829fc609f466 100644 --- a/packages/client-runtime/src/rpc/client.ts +++ b/packages/client-runtime/src/rpc/client.ts @@ -281,6 +281,7 @@ function subscribeDynamicMapped( makeInput: (session: RpcSession) => Effect.Effect>, mapStream: ( session: RpcSession, + input: EnvironmentRpcInput, stream: Stream.Stream, EnvironmentRpcStreamFailure>, ) => Stream.Stream>, options?: SubscriptionOptions, @@ -333,6 +334,7 @@ function subscribeDynamicMapped( }); const stream = mapStream( session, + input, Stream.suspend(() => method(input)), ).pipe( Stream.onFirst(() => @@ -467,7 +469,7 @@ export function subscribeDynamic( EnvironmentRpcStreamFailure, 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 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( + tag: TTag, + makeInput: (session: RpcSession) => Effect.Effect>, + options?: SubscriptionOptions, +): Stream.Stream< + readonly [input: EnvironmentRpcInput, value: EnvironmentRpcStreamValue], + EnvironmentRpcStreamFailure, + 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( diff --git a/packages/client-runtime/src/state/threads-atoms.test.ts b/packages/client-runtime/src/state/threads-atoms.test.ts index 460f7e593990..9aaa605998a5 100644 --- a/packages/client-runtime/src/state/threads-atoms.test.ts +++ b/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); diff --git a/packages/client-runtime/src/state/threads-sync.test.ts b/packages/client-runtime/src/state/threads-sync.test.ts index 5a771646f03f..352b6919d946 100644 --- a/packages/client-runtime/src/state/threads-sync.test.ts +++ b/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 | undefined; }) { const inputs = yield* Queue.unbounded(); const observed = yield* Queue.unbounded(); @@ -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 }, + 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(); + 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 = []; + 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(); + 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 }); diff --git a/packages/client-runtime/src/state/threads.ts b/packages/client-runtime/src/state/threads.ts index b21e688162cd..666e1eae7b76 100644 --- a/packages/client-runtime/src/state/threads.ts +++ b/packages/client-runtime/src/state/threads.ts @@ -35,7 +35,7 @@ import * as ManagedRelay from "../relay/managedRelay.ts"; import { isAuthorizationFailure, request, - subscribeDynamic, + subscribeDynamicWithInput, subscriptionRetryDelay, } from "../rpc/client.ts"; import { parseThreadKey, threadKey } from "./entities.ts"; @@ -644,10 +644,13 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make ) { if (item.kind === "synchronized") { yield* Ref.set(awaitingCompletion, false); - yield* SubscriptionRef.update(state, (current) => - Option.isSome(current.data) && current.status !== "deleted" && Option.isNone(current.error) - ? { ...current, status: "live" as const, error: Option.none() } - : current, + yield* SubscriptionRef.updateSome(state, (current) => + Option.isSome(current.data) && + current.status !== "deleted" && + current.status !== "live" && + Option.isNone(current.error) + ? Option.some({ ...current, status: "live" as const }) + : Option.none(), ); return; } @@ -679,12 +682,41 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make yield* applyEventsLocked([item]); }); + // Items held until the completion marker apply together, so a replay-only resume renders once. + // A new attempt drops the previous attempt's held and late items (unapplied, so its cursor + // replays them). Both run under applyLock: each batch lands wholly before or after an attempt. + let attempt: unknown; + let catchUp: OrchestrationV2ThreadStreamItem[] = []; + const startAttempt = applyLock.withPermits(1)( + Effect.sync(() => { + attempt = undefined; + catchUp = []; + }), + ); // Keep this per-batch streaming path untraced for the same reason as applyEventsLocked. - const applyItems = Effect.fnUntraced(function* ( - items: ReadonlyArray, + const receiveItems = Effect.fnUntraced(function* ( + received: ReadonlyArray, ) { yield* applyLock.withPermits(1)( Effect.gen(function* () { + let completes = false; + for (const [input, item] of received) { + if (input !== attempt) continue; + catchUp.push(item); + if (item.kind === "synchronized") completes = true; + } + if (catchUp.length === 0) return; + if (!completes && (yield* Ref.get(awaitingCompletion))) { + // A warm resume keeps its "live" status until it learns it is behind. + yield* SubscriptionRef.updateSome(state, (current) => + current.status === "live" + ? Option.some({ ...current, status: "synchronizing" as const }) + : Option.none(), + ); + return; + } + const items = catchUp; + catchUp = []; let events: SequencedItem[] = []; for (const item of items) { if ( @@ -694,6 +726,9 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make events.push(item); continue; } + // Clearing the flag first lets the events before the marker publish + // "live", so a catch-up ending in the marker is one state change. + if (item.kind === "synchronized") yield* Ref.set(awaitingCompletion, false); yield* applyEventsLocked(events); events = []; yield* applyItemLocked(item); @@ -1001,9 +1036,10 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make never, EnvironmentSupervisor.EnvironmentSupervisor > => - subscribeDynamic( + subscribeDynamicWithInput( ORCHESTRATION_V2_WS_METHODS.subscribeThread, Effect.fn("EnvironmentThreadState.makeSubscribeInput")(function* (session) { + yield* startAttempt; let current = yield* SubscriptionRef.get(state); // A delete event already cleared this thread. Park without opening the socket. if (current.status === "deleted" && httpMisses === 0) { @@ -1091,7 +1127,7 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make const canResume = Option.isSome(current.data); const acceptBoundedSnapshot = yield* Ref.get(acceptsBoundedSocketSnapshots); - return { + const input = { threadId, ...(canResume ? { afterSequence: sequence } : {}), requestCompletionMarker: true as const, @@ -1105,6 +1141,8 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make : {}), acceptDedupedBodies: true as const, }; + attempt = input; + return input; }), { // A socket-only open that fails retries over HTTP, which tells a @@ -1129,7 +1167,7 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make }, ).pipe( Stream.runForEachArray((items) => - applyItems(items).pipe( + receiveItems(items).pipe( Effect.andThen( Effect.sync(() => { defectRetries = 0;