Skip to content

Commit bb012ea

Browse files
fix(client-runtime): a shell snapshot buffered across a session change keeps its own resume key
Items from the previous session can arrive after the next session's input was built, which tagged an old snapshot with the new session's key. Each item now carries its session's key, and only full snapshots record it. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 769bda2 commit bb012ea

1 file changed

Lines changed: 31 additions & 13 deletions

File tree

  • packages/client-runtime/src/state

‎packages/client-runtime/src/state/shell.ts‎

Lines changed: 31 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ import * as ConnectionWakeups from "../connection/wakeups.ts";
2626
import { safeErrorLogAttributes } from "../errors/safeLog.ts";
2727
import * as Persistence from "../platform/persistence.ts";
2828
import { runCachePersistence } from "./cachePersistence.ts";
29-
import { subscribeDynamic } from "../rpc/client.ts";
29+
import { subscribeDynamicWithSession } from "../rpc/client.ts";
30+
import type { RpcSession } from "../rpc/session.ts";
3031
import * as ShellSnapshotLoader from "./shellSnapshotHttp.ts";
3132
import type { DeferredShellSnapshot } from "./shellPullRequests.ts";
3233
import {
@@ -81,6 +82,16 @@ function withoutDeferredPullRequests(value: DeferredShellSnapshot): Orchestratio
8182

8283
const SHELL_SYNCHRONIZATION_ERROR_MESSAGE = "Could not synchronize environment data.";
8384

85+
function groupBySession<A>(entries: ReadonlyArray<readonly [RpcSession, A]>) {
86+
const groups: Array<[RpcSession, Array<A>]> = [];
87+
for (const [session, value] of entries) {
88+
const last = groups.at(-1);
89+
if (last?.[0] === session) last[1].push(value);
90+
else groups.push([session, [value]]);
91+
}
92+
return groups;
93+
}
94+
8495
export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")(function* () {
8596
const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
8697
const cache = yield* Persistence.EnvironmentCacheStore;
@@ -108,7 +119,8 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
108119
// The omissions and server build of the subscription whose snapshot the
109120
// in-memory shell last came from; null while it holds only cached data.
110121
const authoritativeOmission = yield* Ref.make<string | null>(null);
111-
const activeSubscriptionOmission = yield* Ref.make<string | null>(null);
122+
// Items buffered from a previous session keep that session's key.
123+
const sessionOmissions = new WeakMap<RpcSession, string | null>();
112124
const latestLiveSnapshot = yield* Ref.make<Option.Option<OrchestrationV2ShellSnapshot>>(
113125
Option.none(),
114126
);
@@ -209,12 +221,13 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
209221

210222
const applyItems = Effect.fn("EnvironmentShellState.applyItems")(function* (
211223
items: ReadonlyArray<OrchestrationV2ShellStreamItem>,
224+
omission: string | null,
212225
pendingLinkIds?: ReadonlySet<ThreadId>,
213226
) {
214227
const initial = yield* SubscriptionRef.get(state);
215228
let waiting = yield* Ref.get(awaitingCompletion);
216229
let next = initial;
217-
let receivedSnapshot = false;
230+
let receivedAuthoritativeSnapshot = false;
218231
let deltas: Array<ShellStreamEvent> = [];
219232
const flushDeltas = () => {
220233
if (deltas.length > 0 && Option.isSome(next.snapshot)) {
@@ -268,7 +281,7 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
268281
? undefined
269282
: { resolvedRepositoryIdentityRoots: item.resolvedRepositoryIdentityRoots },
270283
);
271-
receivedSnapshot = true;
284+
if (authoritative) receivedAuthoritativeSnapshot = true;
272285
next = {
273286
snapshot: Option.some(nextSnapshot),
274287
status: waiting ? "synchronizing" : "live",
@@ -289,11 +302,8 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
289302
yield* Ref.set(latestLiveSnapshot, next.snapshot);
290303
}
291304
yield* SubscriptionRef.set(state, next);
292-
if (receivedSnapshot) {
293-
const omission = yield* Ref.get(activeSubscriptionOmission);
294-
if (omission !== null) {
295-
yield* Ref.set(authoritativeOmission, omission);
296-
}
305+
if (receivedAuthoritativeSnapshot && omission !== null) {
306+
yield* Ref.set(authoritativeOmission, omission);
297307
}
298308
if (next.snapshot !== initial.snapshot && Option.isSome(next.snapshot)) {
299309
yield* Queue.offer(persistence, next.snapshot.value);
@@ -375,7 +385,7 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
375385

376386
yield* setSynchronizing;
377387
yield* Effect.forkScoped(
378-
subscribeDynamic(
388+
subscribeDynamicWithSession(
379389
ORCHESTRATION_V2_WS_METHODS.subscribeShell,
380390
Effect.fn("EnvironmentShellState.makeSubscribeInput")(function* (session) {
381391
const { omitIdleSubagents, omitOlderSettled, omissionKey } =
@@ -402,7 +412,7 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
402412
omissionKey: null,
403413
})),
404414
);
405-
yield* Ref.set(activeSubscriptionOmission, omissionKey);
415+
sessionOmissions.set(session, omissionKey);
406416
yield* Ref.set(awaitingCompletion, true);
407417
yield* setSynchronizing;
408418

@@ -445,7 +455,7 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
445455
? []
446456
: httpSnapshot.value.threads.map((thread) => thread.id),
447457
);
448-
yield* applyItems([{ kind: "snapshot", snapshot }], pendingLinkIds);
458+
yield* applyItems([{ kind: "snapshot", snapshot }], omissionKey, pendingLinkIds);
449459
canResume = true;
450460
current = yield* SubscriptionRef.get(state);
451461
if (Option.isSome(current.snapshot)) {
@@ -472,7 +482,15 @@ export const makeEnvironmentShellState = Effect.fn("EnvironmentShellState.make")
472482
retryExpectedFailureAfter: "250 millis",
473483
resubscribe: foregroundResubscriptions,
474484
},
475-
).pipe(Stream.runForEachArray((items) => applyItems(items))),
485+
).pipe(
486+
Stream.runForEachArray((entries) =>
487+
Effect.forEach(
488+
groupBySession(entries),
489+
([session, items]) => applyItems(items, sessionOmissions.get(session) ?? null),
490+
{ discard: true },
491+
),
492+
),
493+
),
476494
);
477495
yield* SubscriptionRef.changes(supervisor.state).pipe(
478496
Stream.runForEach((connectionState) => {

0 commit comments

Comments
 (0)