fix: clean up AUT launches after session deletion

WebDriverAgentLib/Routing/FBSession.m

@@ -50,7 +50,7 @@ @interface FBSession ()
 @property (nonatomic, nullable) FBAlertsMonitor *alertsMonitor;
 @property (nonatomic, readwrite) NSMutableDictionary<NSNumber *, NSMutableDictionary<NSString *, NSNumber *> *> *elementsVisibilityCache;
 
-- (BOOL)fb_isTestedApplicationSameAsSystemAppWithTimeout:(NSTimeInterval)timeout;
+- (BOOL)fb_isApplicationSameAsSystemApp:(XCUIApplication *)application timeout:(NSTimeInterval)timeout;
 - (void)fb_terminateTestedApplicationWithTimeout:(NSTimeInterval)timeout generation:(NSUInteger)generation;
 @end
 
@@ -333,7 +333,7 @@ - (void)kill
     if (nil != self.testedApplication
         && FBConfiguration.sharedInstance.shouldTerminateApp
         && self.testedApplication.running
-        && ![self fb_isTestedApplicationSameAsSystemAppWithTimeout:FB_IS_SYSTEM_APP_CHECK_TIMEOUT_SEC]) {
+        && ![self fb_isApplicationSameAsSystemApp:self.testedApplication timeout:FB_IS_SYSTEM_APP_CHECK_TIMEOUT_SEC]) {
       // Blocks until the app is either actually terminated or durably given up on (never left
       // pending) - see -fb_terminateTestedApplicationWithTimeout:generation: - so it's safe to
       // report this teardown as finished as soon as this returns.
@@ -401,6 +401,22 @@ - (XCUIApplication *)launchApplicationWithBundleId:(NSString *)bundleIdentifier
                                        environment:(nullable NSDictionary <NSString *, NSString *> *)environment
                             asApplicationUnderTest:(BOOL)asApplicationUnderTest
 {
+  NSCondition *condition = self.class.teardownCondition;
+  NSUInteger generation = 0;
+  BOOL shouldTerminateApp = NO;
+  if (asApplicationUnderTest) {
+    [condition lock];
+    BOOL isActive = self == _activeSession;
+    generation = _sessionGeneration;
+    shouldTerminateApp = FBConfiguration.sharedInstance.shouldTerminateApp;
+    [condition unlock];
+    if (!isActive) {
+      @throw [NSException exceptionWithName:FBSessionDoesNotExistException
+                                    reason:@"Session does not exist"
+                                  userInfo:nil];
+    }
+  }
+
   XCUIApplication *app = [self makeApplicationWithBundleId:bundleIdentifier];
   if (nil == shouldWaitForQuiescence) {
     // Iherit the quiescence check setting from the main app under test by default.
@@ -419,9 +435,36 @@ - (XCUIApplication *)launchApplicationWithBundleId:(NSString *)bundleIdentifier
     [app activate];
   }
   if (asApplicationUnderTest) {
+    BOOL waitForQuiescence = app.fb_shouldWaitForQuiescence;
+    // A DELETE can finish while launch/activate is still waiting for XCTest.
+    // Publish the AUT under the same lock that -kill uses to clear the session.
+    [condition lock];
+    BOOL isActive = self == _activeSession && generation == _sessionGeneration;
+    if (isActive) {
+      self.testedApplication = app;
+      self.shouldAppsWaitForQuiescence = waitForQuiescence;
+    }
+    [condition unlock];
+    if (!isActive) {
+      // The HTTP request may already have been abandoned, but its handler keeps
+      // running. Clean up only while this launch still owns the generation;
+      // a replacement session may already be using the same bundle identifier.
+      if (shouldTerminateApp
+          && ![self fb_isApplicationSameAsSystemApp:app timeout:FB_IS_SYSTEM_APP_CHECK_TIMEOUT_SEC]
+          && [self.class beginTerminationForGeneration:generation]) {
+        @try {
+          [app terminate];
+        } @catch (NSException *e) {
+          [FBLogger logFmt:@"%@", e.description];
+        } @finally {
+          [self.class endTermination];
+        }
+      }
+      @throw [NSException exceptionWithName:FBSessionDoesNotExistException
+                                    reason:@"Session was deleted during application launch"
+                                  userInfo:nil];
+    }
     [FBLogger logFmt:@"The application '%@' is set as the application under test", bundleIdentifier];
-    self.testedApplication = app;
-    self.shouldAppsWaitForQuiescence = app.fb_shouldWaitForQuiescence;
   }
   if ([app fb_isSameAppAs:self.testedApplication]) {
     self.isTestedApplicationExpectedToRun = YES;
@@ -462,9 +505,9 @@ - (XCUIApplication *)makeApplicationWithBundleId:(NSString *)bundleIdentifier
 }
 
 // Has no async variant and can block on the shared accessibility channel. Run off-thread and give
-// up after `timeout`, assuming the tested app IS the system app - safer, since it means skipping
+// up after `timeout`, assuming the supplied app IS the system app - safer, since it means skipping
 // termination rather than risking terminating springboard.
-- (BOOL)fb_isTestedApplicationSameAsSystemAppWithTimeout:(NSTimeInterval)timeout
+- (BOOL)fb_isApplicationSameAsSystemApp:(XCUIApplication *)application timeout:(NSTimeInterval)timeout
 {
   __block XCUIApplication *systemApp = nil;
   __block NSException *caughtException = nil;
@@ -473,18 +516,18 @@ - (BOOL)fb_isTestedApplicationSameAsSystemAppWithTimeout:(NSTimeInterval)timeout
     // Undocumented private API; guard in case it hard-asserts off-main like -terminate does on
     // some Xcode/iOS version - uncaught, that would crash the whole process.
     @try {
-      systemApp = XCUIApplication.fb_systemApplication;
+      systemApp = [[application class] fb_systemApplication];
     } @catch (NSException *e) {
       caughtException = e;
     }
     dispatch_semaphore_signal(sem);
   });
   int64_t timeoutNs = (int64_t)(timeout * NSEC_PER_SEC);
   if (0 != dispatch_semaphore_wait(sem, dispatch_time(DISPATCH_TIME_NOW, timeoutNs)) || nil != caughtException) {
-    [FBLogger logFmt:@"Could not determine the system application within %@ seconds%@; assuming '%@' might be it and skipping its termination", @(timeout), nil == caughtException ? @"" : [NSString stringWithFormat:@" (%@)", caughtException.description], self.testedApplication.bundleID];
+    [FBLogger logFmt:@"Could not determine the system application within %@ seconds%@; assuming '%@' might be it and skipping its termination", @(timeout), nil == caughtException ? @"" : [NSString stringWithFormat:@" (%@)", caughtException.description], application.bundleID];
     return YES;
   }
-  return [self.testedApplication fb_isSameAppAs:systemApp];
+  return [application fb_isSameAppAs:systemApp];
 }
 
 // -terminate hard-asserts off-main, but -kill can now run on a background queue. Dispatching to

WebDriverAgentTests/IntegrationTests/FBSessionIntegrationTests.m

@@ -172,6 +172,32 @@ - (void)testApplicationIsNotLaunchedAsApplicationUnderTestByDefault
   XCTAssertNoThrow(self.session.activeApplication);
 }
 
+// DELETE /session while /wda/apps/launch is still launching the app as the application under
+// test: whichever of them finishes first, the app must not outlive the session. Settings takes
+// long enough to launch for the deletion to usually land in the middle.
+- (void)testApplicationLaunchedAsApplicationUnderTestDoesNotOutliveConcurrentSessionDeletion
+{
+  [self startSessionWithoutApplicationUnderTest];
+  XCUIApplication *settingsApp = [[XCUIApplication alloc] initWithBundleIdentifier:SETTINGS_BUNDLE_ID];
+  [settingsApp terminate];
+  FBAssertWaitTillBecomesTrue(settingsApp.state == XCUIApplicationStateNotRunning);
+  FBSession *session = self.session;
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.5 * NSEC_PER_SEC)),
+                 dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0), ^{
+    [session kill];
+  });
+  @try {
+    [session launchApplicationWithBundleId:SETTINGS_BUNDLE_ID
+                   shouldWaitForQuiescence:nil
+                                 arguments:nil
+                               environment:nil
+                    asApplicationUnderTest:YES];
+  } @catch (NSException *e) {
+    XCTAssertEqualObjects(e.name, FBSessionDoesNotExistException);
+  }
+  FBAssertWaitTillBecomesTrue(settingsApp.state == XCUIApplicationStateNotRunning);
+}
+
 - (void)testLaunchUnattachedApp
 {
   [FBUnattachedAppLauncher launchAppWithBundleId:SETTINGS_BUNDLE_ID];

WebDriverAgentTests/UnitTests/FBSessionTests.m

@@ -9,9 +9,67 @@
 #import <XCTest/XCTest.h>
 
 #import "FBSession.h"
+#import "FBSession-Private.h"
+#import "FBExceptions.h"
 #import "FBConfiguration.h"
 #import "XCUIApplicationDouble.h"
 
+@interface FBSession (SystemApplicationTests)
+- (BOOL)fb_isApplicationSameAsSystemApp:(XCUIApplication *)application timeout:(NSTimeInterval)timeout;
+@end
+
+static id (^FBSystemApplicationLookup)(void);
+
+// Run teardown at the exact point where launch has started but the AUT has
+// not yet been published. No wall-clock sleeps or simulator launch timing.
+@interface FBAUTLaunchApplicationDouble : XCUIApplicationDouble
+@property (nonatomic, copy) void (^duringLaunch)(void);
+@property (nonatomic, copy) NSArray *launchArguments;
+@property (nonatomic, copy) NSDictionary *launchEnvironment;
+@property (nonatomic) BOOL didLaunch;
+@end
+
+@implementation FBAUTLaunchApplicationDouble
++ (id)fb_systemApplication
+{
+  return FBSystemApplicationLookup ? FBSystemApplicationLookup() : nil;
+}
+- (void)launch
+{
+  self.didLaunch = YES;
+  if (self.duringLaunch) {
+    self.duringLaunch();
+  }
+}
+- (BOOL)fb_isSameAppAs:(id)application
+{
+  return self == application;
+}
+@end
+
+@interface FBAUTTerminationFailureApplicationDouble : FBAUTLaunchApplicationDouble
+@property (nonatomic) BOOL didAttemptTermination;
+@end
+
+@implementation FBAUTTerminationFailureApplicationDouble
+- (void)terminate
+{
+  self.didAttemptTermination = YES;
+  @throw [NSException exceptionWithName:@"TerminationFailure" reason:nil userInfo:nil];
+}
+@end
+
+@interface FBAUTLaunchSession : FBSession
+@property (nonatomic, strong) FBAUTLaunchApplicationDouble *launchApplication;
+@end
+
+@implementation FBAUTLaunchSession
+- (XCUIApplication *)makeApplicationWithBundleId:(NSString *)bundleIdentifier
+{
+  return (id)self.launchApplication;
+}
+@end
+
 @interface FBSessionTests : XCTestCase
 @property (nonatomic, strong) FBSession *session;
 @property (nonatomic, strong) XCUIApplication *testedApplication;
@@ -31,6 +89,10 @@ - (void)setUp
 
 - (void)tearDown
 {
+  if ([self.session isKindOfClass:FBAUTLaunchSession.class]) {
+    ((FBAUTLaunchSession *)self.session).launchApplication.duringLaunch = nil;
+  }
+  FBSystemApplicationLookup = nil;
   [self.session kill];
   FBConfiguration.sharedInstance.shouldTerminateApp = self.shouldTerminateAppValue;
   [super tearDown];
@@ -64,4 +126,143 @@ - (void)testActiveSessionIsNilAfterKilling
   XCTAssertNil([FBSession activeSession]);
 }
 
+- (FBAUTLaunchApplicationDouble *)prepareAUTLaunch
+{
+  FBAUTLaunchSession *session = [FBAUTLaunchSession new];
+  session.identifier = NSUUID.UUID.UUIDString;
+  session.launchApplication = [FBAUTLaunchApplicationDouble new];
+  [FBSession markSessionActive:session];
+  self.session = session;
+  return session.launchApplication;
+}
+
+- (void)launchAsAUT
+{
+  [self.session launchApplicationWithBundleId:@"com.facebook.awesome"
+                    shouldWaitForQuiescence:nil
+                                  arguments:nil
+                                environment:nil
+                     asApplicationUnderTest:YES];
+}
+
+- (void)testCompletedLaunchRegistersAUT
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  XCTAssertNoThrow([self launchAsAUT]);
+  XCTAssertEqual([self.session valueForKey:@"testedApplication"], app);
+  XCTAssertFalse(app.didTerminate);
+}
+
+- (void)testDeletedSessionCleansUpInFlightAUTLaunch
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  FBConfiguration.sharedInstance.shouldTerminateApp = YES;
+  FBSession *session = self.session;
+  app.duringLaunch = ^{ [session kill]; };
+  XCTAssertThrowsSpecificNamed([self launchAsAUT], NSException,
+                              FBSessionDoesNotExistException);
+  XCTAssertTrue(app.didTerminate);
+  XCTAssertNil([self.session valueForKey:@"testedApplication"]);
+  XCTAssertNil(FBSession.activeSession);
+}
+
+- (void)testTerminationFailurePreservesDeletedSessionErrorAndReleasesClaim
+{
+  [self prepareAUTLaunch];
+  FBAUTTerminationFailureApplicationDouble *app = [FBAUTTerminationFailureApplicationDouble new];
+  ((FBAUTLaunchSession *)self.session).launchApplication = app;
+  FBConfiguration.sharedInstance.shouldTerminateApp = YES;
+  FBSession *session = self.session;
+  app.duringLaunch = ^{ [session kill]; };
+
+  XCTAssertThrowsSpecificNamed([self launchAsAUT], NSException, FBSessionDoesNotExistException);
+  XCTAssertTrue(app.didAttemptTermination);
+  XCTAssertNil([self.session valueForKey:@"testedApplication"]);
+  XCTAssertNil(FBSession.activeSession);
+
+  // Creating the next session must succeed even when cleanup termination throws.
+  // A leaked generation claim would make this fail with a session creation error.
+  FBSession *replacement = nil;
+  XCTAssertNoThrow(replacement = [FBSession initWithApplication:nil]);
+  XCTAssertNotNil(replacement);
+  XCTAssertEqual(FBSession.activeSession, replacement);
+  [replacement kill];
+}
+
+- (void)testDeletedSessionPreservesAppWhenTerminationIsDisabled
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  FBSession *session = self.session;
+  app.duringLaunch = ^{ [session kill]; };
+  XCTAssertThrowsSpecificNamed([self launchAsAUT], NSException,
+                              FBSessionDoesNotExistException);
+  XCTAssertFalse(app.didTerminate);
+}
+
+- (void)testAbandonedAUTLaunchDoesNotTerminateReplacementSessionApp
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  FBConfiguration.sharedInstance.shouldTerminateApp = YES;
+  FBSession *session = self.session;
+  __block FBSession *replacement;
+  app.duringLaunch = ^{
+    [session kill];
+    replacement = [FBSession initWithApplication:(id)app];
+  };
+  XCTAssertThrowsSpecificNamed([self launchAsAUT], NSException,
+                              FBSessionDoesNotExistException);
+  XCTAssertFalse(app.didTerminate);
+  XCTAssertEqual(FBSession.activeSession, replacement);
+  [replacement kill];
+}
+
+- (void)testAUTLaunchRejectsAnAlreadyDeletedSession
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  [self.session kill];
+  XCTAssertThrowsSpecificNamed([self launchAsAUT], NSException,
+                              FBSessionDoesNotExistException);
+  XCTAssertFalse(app.didLaunch);
+}
+
+- (void)testSystemApplicationCheckUsesSuppliedApplication
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  FBSystemApplicationLookup = ^id { return app; };
+  XCTAssertTrue([self.session fb_isApplicationSameAsSystemApp:(id)app timeout:1.]);
+  FBSystemApplicationLookup = ^id { return nil; };
+  XCTAssertFalse([self.session fb_isApplicationSameAsSystemApp:(id)app timeout:1.]);
+}
+
+- (void)testSystemApplicationLookupExceptionSkipsTermination
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  FBConfiguration.sharedInstance.shouldTerminateApp = YES;
+  FBSession *session = self.session;
+  app.duringLaunch = ^{ [session kill]; };
+  FBSystemApplicationLookup = ^id {
+    @throw [NSException exceptionWithName:@"LookupFailure" reason:nil userInfo:nil];
+  };
+  XCTAssertThrowsSpecificNamed([self launchAsAUT], NSException, FBSessionDoesNotExistException);
+  XCTAssertFalse(app.didTerminate);
+}
+
+- (void)testSystemApplicationLookupTimeoutSkipsTermination
+{
+  FBAUTLaunchApplicationDouble *app = [self prepareAUTLaunch];
+  dispatch_semaphore_t releaseLookup = dispatch_semaphore_create(0);
+  XCTestExpectation *finished = [self expectationWithDescription:@"System application lookup finished"];
+  FBSystemApplicationLookup = ^id {
+    dispatch_semaphore_wait(releaseLookup, DISPATCH_TIME_FOREVER);
+    [finished fulfill];
+    return nil;
+  };
+  @try {
+    XCTAssertTrue([self.session fb_isApplicationSameAsSystemApp:(id)app timeout:0.01]);
+  } @finally {
+    dispatch_semaphore_signal(releaseLookup);
+    [self waitForExpectationsWithTimeout:5. handler:nil];
+  }
+}
+
 @end