@@ -26,7 +26,8 @@ import * as ConnectionWakeups from "../connection/wakeups.ts";
2626import { safeErrorLogAttributes } from "../errors/safeLog.ts" ;
2727import * as Persistence from "../platform/persistence.ts" ;
2828import { 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" ;
3031import * as ShellSnapshotLoader from "./shellSnapshotHttp.ts" ;
3132import type { DeferredShellSnapshot } from "./shellPullRequests.ts" ;
3233import {
@@ -81,6 +82,16 @@ function withoutDeferredPullRequests(value: DeferredShellSnapshot): Orchestratio
8182
8283const 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+
8495export 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