@@ -294,19 +294,36 @@ function makeProviderAdapter(
294294 readonly hasPendingBackgroundWork ?: Effect . Effect < boolean > ;
295295 readonly hangSessionScopeClose ?: boolean ;
296296 readonly beforeUnload ?: Effect . Effect < void > ;
297+ /** Registers the process's closeCount finalizer before `beforeOpen` runs. */
298+ readonly spawnBeforeOpen ?: boolean ;
297299 } = { } ,
298300) : ProviderAdapterV2Shape {
301+ const countClose = Effect . addFinalizer ( ( ) =>
302+ Ref . update ( state , ( current ) => ( {
303+ ...current ,
304+ closeCount : current . closeCount + 1 ,
305+ } ) ) ,
306+ ) ;
299307 return {
300308 instanceId : ProviderInstanceId . make ( "codex" ) ,
301309 driver : CODEX_DRIVER ,
302310 getCapabilities : ( ) => Effect . succeed ( options . capabilities ?? CodexCapabilities ) ,
303311 planSelectionTransition : ( ) => Effect . succeed ( { type : "apply_on_next_turn" } ) ,
304312 openSession : ( input ) =>
305313 Effect . gen ( function * ( ) {
314+ if ( options . spawnBeforeOpen === true ) {
315+ yield * countClose ;
316+ }
317+ if ( options . mcpConfigs !== undefined && options . spawnBeforeOpen === true ) {
318+ yield * Ref . update ( options . mcpConfigs , ( configs ) => [
319+ ...configs ,
320+ McpProviderSession . readMcpProviderSession ( input . threadId ) ,
321+ ] ) ;
322+ }
306323 if ( options . beforeOpen !== undefined ) {
307324 yield * options . beforeOpen ( input ) ;
308325 }
309- if ( options . mcpConfigs !== undefined ) {
326+ if ( options . mcpConfigs !== undefined && options . spawnBeforeOpen !== true ) {
310327 yield * Ref . update ( options . mcpConfigs , ( configs ) => [
311328 ...configs ,
312329 McpProviderSession . readMcpProviderSession ( input . threadId ) ,
@@ -328,12 +345,9 @@ function makeProviderAdapter(
328345 eventQueues,
329346 } ;
330347 } ) ;
331- yield * Effect . addFinalizer ( ( ) =>
332- Ref . update ( state , ( current ) => ( {
333- ...current ,
334- closeCount : current . closeCount + 1 ,
335- } ) ) ,
336- ) ;
348+ if ( options . spawnBeforeOpen !== true ) {
349+ yield * countClose ;
350+ }
337351 if ( options . hangSessionScopeClose === true ) {
338352 // Registered last so it runs first on scope close, wedging the
339353 // close before the closeCount finalizer, like a provider process
@@ -410,6 +424,7 @@ function makeTestLayer(input: {
410424 readonly hasPendingBackgroundWork ?: Effect . Effect < boolean > ;
411425 readonly hangSessionScopeClose ?: boolean ;
412426 readonly beforeUnload ?: Effect . Effect < void > ;
427+ readonly spawnBeforeOpen ?: boolean ;
413428 readonly serverSettingsLayer ?: ReturnType < typeof ServerSettings . layerTest > ;
414429 readonly projectServiceLayer ?: Layer . Layer < ProjectService . ProjectService > ;
415430} ) {
@@ -432,6 +447,7 @@ function makeTestLayer(input: {
432447 ? { }
433448 : { hangSessionScopeClose : input . hangSessionScopeClose } ) ,
434449 ...( input . beforeUnload === undefined ? { } : { beforeUnload : input . beforeUnload } ) ,
450+ ...( input . spawnBeforeOpen === undefined ? { } : { spawnBeforeOpen : input . spawnBeforeOpen } ) ,
435451 } ) ,
436452 ) ;
437453 const providerEventIngestorTestLayer = ProviderEventIngestor . layer . pipe (
@@ -892,6 +908,127 @@ it.effect("ProviderSessionManagerV2 opens a duplicate session only once", () =>
892908 } ) ,
893909) ;
894910
911+ it . effect ( "ProviderSessionManagerV2 cleans up an open interrupted mid-handshake" , ( ) =>
912+ Effect . gen ( function * ( ) {
913+ const state = yield * Ref . make ( emptyState ) ;
914+ const mcpConfigs = yield * Ref . make <
915+ ReadonlyArray < McpProviderSession . McpProviderSessionConfig | undefined >
916+ > ( [ ] ) ;
917+ const handshakeStarted = yield * Deferred . make < void > ( ) ;
918+ const holdHandshake = yield * Ref . make ( true ) ;
919+ const effect = Effect . gen ( function * ( ) {
920+ const eventSink = yield * EventSink . EventSinkV2 ;
921+ const idAllocator = yield * IdAllocator . IdAllocatorV2 ;
922+ const manager = yield * ProviderSessionManager . ProviderSessionManagerV2 ;
923+ const registry = yield * McpSessionRegistry . McpSessionRegistry ;
924+ const projectionStore = yield * ProjectionStore . ProjectionStoreV2 ;
925+ const now = yield * DateTime . now ;
926+ const threadId = ThreadId . make ( "thread-provider-session-manager-interrupted-open" ) ;
927+ const providerSessionId = yield * idAllocator . allocate . providerSession ( {
928+ providerInstanceId : modelSelection . instanceId ,
929+ threadId,
930+ } ) ;
931+ yield * eventSink . write ( {
932+ events : [ yield * makeThreadCreatedEvent ( { idAllocator, threadId, now } ) ] ,
933+ } ) ;
934+
935+ const opening = yield * manager
936+ . open ( { threadId, providerSessionId, modelSelection, runtimePolicy } )
937+ . pipe ( Effect . forkScoped ) ;
938+ yield * Deferred . await ( handshakeStarted ) ;
939+ const issued = ( yield * Ref . get ( mcpConfigs ) ) . at ( - 1 ) ;
940+ const token = issued ?. authorizationHeader . replace ( / ^ B e a r e r \s + / , "" ) ;
941+ assert . isDefined ( token ) ;
942+ assert . isDefined ( yield * registry . resolve ( token ! ) ) ;
943+
944+ // A Stop while the provider is still starting.
945+ yield * Fiber . interrupt ( opening ) ;
946+
947+ // The process started for this open is stopped, and the credential minted
948+ // for it revoked.
949+ assert . equal ( ( yield * Ref . get ( state ) ) . closeCount , 1 ) ;
950+ assert . isUndefined ( yield * registry . resolve ( token ! ) ) ;
951+ assert . isUndefined ( McpProviderSession . readMcpProviderSession ( threadId ) ) ;
952+ assert . isTrue ( Option . isNone ( yield * manager . get ( providerSessionId ) ) ) ;
953+
954+ // Nothing of the interrupted open is left behind: the next open starts a
955+ // fresh process with a fresh credential that a later release revokes,
956+ // which a leaked reservation would prevent.
957+ yield * Ref . set ( holdHandshake , false ) ;
958+ yield * manager . open ( { threadId, providerSessionId, modelSelection, runtimePolicy } ) ;
959+ const replacement = ( yield * Ref . get ( mcpConfigs ) ) . at ( - 1 ) ;
960+ const replacementToken = replacement ?. authorizationHeader . replace ( / ^ B e a r e r \s + / , "" ) ;
961+ assert . isDefined ( replacementToken ) ;
962+ assert . notEqual ( replacementToken , token ) ;
963+ const projection = yield * projectionStore . getThreadProjection ( threadId ) ;
964+ assert . equal ( projection . providerSessions . at ( - 1 ) ?. status , "ready" ) ;
965+
966+ yield * manager . close ( providerSessionId ) ;
967+ assert . isUndefined ( yield * registry . resolve ( replacementToken ! ) ) ;
968+ } ) ;
969+
970+ yield * effect . pipe (
971+ Effect . provide (
972+ makeTestLayer ( {
973+ state,
974+ idleTimeoutMs : 60_000 ,
975+ mcpConfigs,
976+ beforeOpen : ( ) =>
977+ Ref . get ( holdHandshake ) . pipe (
978+ Effect . flatMap ( ( hold ) =>
979+ hold
980+ ? Deferred . succeed ( handshakeStarted , undefined ) . pipe ( Effect . andThen ( Effect . never ) )
981+ : Effect . void ,
982+ ) ,
983+ ) ,
984+ // The process is spawned before the handshake that is interrupted.
985+ spawnBeforeOpen : true ,
986+ } ) ,
987+ ) ,
988+ ) ;
989+ } ) ,
990+ ) ;
991+
992+ it . effect ( "ProviderSessionManagerV2 stops a session still opening when its layer shuts down" , ( ) =>
993+ Effect . gen ( function * ( ) {
994+ const state = yield * Ref . make ( emptyState ) ;
995+ const handshakeStarted = yield * Deferred . make < void > ( ) ;
996+ yield * Effect . gen ( function * ( ) {
997+ const eventSink = yield * EventSink . EventSinkV2 ;
998+ const idAllocator = yield * IdAllocator . IdAllocatorV2 ;
999+ const manager = yield * ProviderSessionManager . ProviderSessionManagerV2 ;
1000+ const now = yield * DateTime . now ;
1001+ const threadId = ThreadId . make ( "thread-provider-session-manager-shutdown-during-open" ) ;
1002+ const providerSessionId = yield * idAllocator . allocate . providerSession ( {
1003+ providerInstanceId : modelSelection . instanceId ,
1004+ threadId,
1005+ } ) ;
1006+ yield * eventSink . write ( {
1007+ events : [ yield * makeThreadCreatedEvent ( { idAllocator, threadId, now } ) ] ,
1008+ } ) ;
1009+ // Detached, like an open the layer does not own: only the session
1010+ // scope's parent can stop its process when the layer closes.
1011+ yield * manager
1012+ . open ( { threadId, providerSessionId, modelSelection, runtimePolicy } )
1013+ . pipe ( Effect . forkDetach ) ;
1014+ yield * Deferred . await ( handshakeStarted ) ;
1015+ assert . equal ( ( yield * Ref . get ( state ) ) . closeCount , 0 ) ;
1016+ } ) . pipe (
1017+ Effect . provide (
1018+ makeTestLayer ( {
1019+ state,
1020+ idleTimeoutMs : 60_000 ,
1021+ beforeOpen : ( ) =>
1022+ Deferred . succeed ( handshakeStarted , undefined ) . pipe ( Effect . andThen ( Effect . never ) ) ,
1023+ spawnBeforeOpen : true ,
1024+ } ) ,
1025+ ) ,
1026+ ) ;
1027+
1028+ assert . equal ( ( yield * Ref . get ( state ) ) . closeCount , 1 ) ;
1029+ } ) ,
1030+ ) ;
1031+
8951032it . effect ( "ProviderSessionManagerV2 releases live sessions when its layer shuts down" , ( ) =>
8961033 Effect . gen ( function * ( ) {
8971034 const state = yield * Ref . make ( emptyState ) ;
0 commit comments