langfun 0.1.2.dev202510310805__py3-none-any.whl → 0.1.2.dev202511020804__py3-none-any.whl

langfun/env/{base_test.py → base_sandbox_test.py}

@@ -17,8 +17,6 @@ import unittest
 
 from langfun.env import interface
 from langfun.env import test_utils
-from langfun.env.event_handlers import base as event_handler_base
-
 
 TestingEnvironment = test_utils.TestingEnvironment
 TestingSandbox = test_utils.TestingSandbox
@@ -26,453 +24,7 @@ TestingFeature = test_utils.TestingFeature
 TestingEventHandler = test_utils.TestingEventHandler
 
 
-class EnvironmentTests(unittest.TestCase):
-
-  def test_basics(self):
-    env = TestingEnvironment(
-        image_ids=['test_image'],
-        root_dir='/tmp',
-        pool_size=0,
-        features={'test_feature': TestingFeature()},
-        outage_grace_period=1,
-        outage_retry_interval=0,
-    )
-    self.assertIsNone(interface.Environment.current())
-    self.assertEqual(env.image_ids, ['test_image'])
-    self.assertFalse(env.supports_dynamic_image_loading)
-    self.assertEqual(env.status, interface.Environment.Status.CREATED)
-    self.assertFalse(env.is_online)
-    self.assertEqual(env.min_pool_size('test_image'), 0)
-    self.assertEqual(env.max_pool_size('test_image'), 0)
-    self.assertEqual(env.sandbox_pool, {})
-    self.assertEqual(env.id, interface.Environment.Id('testing-env'))
-    self.assertEqual(env.outage_grace_period, 1)
-    self.assertEqual(env.features['test_feature'].name, 'test_feature')
-
-    self.assertIsNone(env.start_time)
-
-    with env:
-      self.assertEqual(env.status, interface.Environment.Status.ONLINE)
-      self.assertIs(interface.Environment.current(), env)
-      self.assertTrue(env.is_online)
-      self.assertIsNotNone(env.start_time)
-      self.assertEqual(env.offline_duration, 0.0)
-      self.assertEqual(env.sandbox_pool, {})
-      self.assertEqual(env.working_dir, '/tmp/testing-env')
-
-      with env.sandbox(session_id='session1') as sb:
-        self.assertEqual(
-            sb.id, interface.Sandbox.Id(
-                environment_id=env.id,
-                image_id=sb.image_id,
-                sandbox_id='0')
-        )
-        self.assertEqual(sb.session_id, 'session1')
-        self.assertEqual(sb.working_dir, '/tmp/testing-env/test_image/0')
-        self.assertTrue(sb.is_online)
-        self.assertIs(sb.test_feature, sb.features['test_feature'])
-        self.assertEqual(
-            sb.test_feature.working_dir,
-            '/tmp/testing-env/test_image/0/test_feature'
-        )
-        with self.assertRaises(AttributeError):
-          _ = sb.test_feature2
-      self.assertFalse(sb.is_online)
-
-      with self.assertRaisesRegex(
-          ValueError, 'Environment .* does not serve image ID .*'
-      ):
-        env.sandbox('test_image2')
-
-      with env.test_feature() as feature:
-        self.assertIsInstance(feature, TestingFeature)
-        self.assertEqual(
-            feature.sandbox.status, interface.Sandbox.Status.IN_SESSION
-        )
-        self.assertTrue(
-            feature.sandbox.session_id.startswith('test_feature-session')
-        )
-
-      with self.assertRaises(AttributeError):
-        _ = env.test_feature2
-
-  def test_dynamic_image_loading(self):
-    env = TestingEnvironment(
-        image_ids=[],
-        supports_dynamic_image_loading=True,
-        pool_size=0,
-        features={'test_feature': TestingFeature()},
-        outage_grace_period=1,
-        outage_retry_interval=0,
-    )
-    with env:
-      with env.sandbox(image_id='test_image2') as sb:
-        self.assertEqual(sb.image_id, 'test_image2')
-
-      with self.assertRaisesRegex(
-          ValueError, 'Environment .* does not have a default image ID.'
-      ):
-        env.sandbox()
-
-  def test_dynamic_image_loading_with_pooling(self):
-    env = TestingEnvironment(
-        image_ids=[],
-        supports_dynamic_image_loading=True,
-        pool_size=2,
-        features={'test_feature': TestingFeature()},
-        outage_grace_period=1,
-        outage_retry_interval=0,
-    )
-    with env:
-      with env.sandbox(image_id='test_image'):
-        self.assertEqual(len(env.sandbox_pool['test_image']), 1)
-
-        with env.sandbox(image_id='test_image'):
-          self.assertEqual(len(env.sandbox_pool['test_image']), 2)
-
-          with self.assertRaises(interface.EnvironmentOverloadError):
-            with env.sandbox(image_id='test_image'):
-              pass
-        self.assertEqual(len(env.sandbox_pool['test_image']), 2)
-
-        with env.sandbox(image_id='test_image'):
-          self.assertEqual(len(env.sandbox_pool['test_image']), 2)
-
-  def test_image_feature_mappings(self):
-    env = TestingEnvironment(
-        image_ids=[
-            'test_image1',
-            'test_image2',
-        ],
-        features={
-            'test_feature': TestingFeature(
-                applicable_images=['test_image1.*']
-            ),
-            'test_feature2': TestingFeature(
-                applicable_images=['test_image2.*']
-            ),
-            'test_feature3': TestingFeature(
-                applicable_images=['test_image.*']
-            ),
-        },
-        pool_size=0,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      with env.sandbox(image_id='test_image1') as sb:
-        self.assertIn('test_feature', sb.features)
-        self.assertNotIn('test_feature2', sb.features)
-        self.assertIn('test_feature3', sb.features)
-
-      with env.sandbox(image_id='test_image2') as sb:
-        self.assertNotIn('test_feature', sb.features)
-        self.assertIn('test_feature2', sb.features)
-        self.assertIn('test_feature3', sb.features)
-
-      with env.test_feature() as feature:
-        self.assertEqual(feature.sandbox.image_id, 'test_image1')
-
-      with self.assertRaisesRegex(
-          ValueError, 'Feature .* is not applicable to .*'
-      ):
-        with env.test_feature('test_image2'):
-          pass
-
-      with env.test_feature2() as feature:
-        self.assertEqual(feature.sandbox.image_id, 'test_image2')
-
-      with env.test_feature3() as feature:
-        self.assertEqual(feature.sandbox.image_id, 'test_image1')
-
-      with env.test_feature3('test_image2') as feature:
-        self.assertEqual(feature.sandbox.image_id, 'test_image2')
-
-  def test_feature_applicability_check(self):
-    with self.assertRaisesRegex(
-        ValueError, 'Feature .* is not applicable to .*'
-    ):
-      TestingEnvironment(
-          image_ids=[
-              'test_image1',
-          ],
-          features={
-              'test_feature2': TestingFeature(
-                  applicable_images=['test_image2.*']
-              ),
-          },
-      )
-    env = TestingEnvironment(
-        image_ids=[],
-        supports_dynamic_image_loading=True,
-        features={
-            'test_feature2': TestingFeature(
-                applicable_images=['test_image2.*']
-            ),
-        },
-        pool_size=0
-    )
-    with env:
-      with self.assertRaisesRegex(
-          ValueError, 'No image ID found for feature .*'
-      ):
-        with env.test_feature2():
-          pass
-
-      # Dynamically loaded IDs.
-      with env.test_feature2('test_image2') as feature:
-        self.assertEqual(feature.sandbox.image_id, 'test_image2')
-
-  def test_pool_size(self):
-    env = TestingEnvironment(
-        image_ids=['test_image'],
-        pool_size=1,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-    )
-    self.assertEqual(env.min_pool_size('test_image'), 1)
-    self.assertEqual(env.max_pool_size('test_image'), 1)
-
-    env = TestingEnvironment(
-        image_ids=['test_image'],
-        pool_size=(0, 256),
-        outage_grace_period=1,
-        outage_retry_interval=0,
-    )
-    self.assertEqual(env.min_pool_size('test_image'), 0)
-    self.assertEqual(env.max_pool_size('test_image'), 256)
-
-    env = TestingEnvironment(
-        image_ids=['test_image'],
-        pool_size={
-            'test_.*': (0, 128),
-            'my.*': (5, 64),
-            'exact_image_name': 10,
-        },
-        outage_grace_period=1,
-        outage_retry_interval=0,
-    )
-    self.assertEqual(env.min_pool_size('test_image'), 0)
-    self.assertEqual(env.max_pool_size('test_image'), 128)
-    self.assertEqual(env.min_pool_size('my_image'), 5)
-    self.assertEqual(env.max_pool_size('my_image'), 64)
-    self.assertEqual(env.min_pool_size('exact_image_name'), 10)
-    self.assertEqual(env.max_pool_size('exact_image_name'), 10)
-    self.assertEqual(env.min_pool_size('some_image'), 0)  # default
-    self.assertEqual(env.max_pool_size('some_image'), 256)  # default
-
-  def test_del(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=0,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    env.start()
-    sb = env.acquire()
-    del sb
-    del env
-
-  def test_acquire_env_offline(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=0,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with self.assertRaises(interface.EnvironmentOutageError):
-      env.acquire()
-
-  def test_acquire_no_pooling(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=0,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      sb = env.acquire()
-      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
-      self.assertIsNone(env.working_dir)
-      self.assertIsNone(sb.working_dir)
-      self.assertIsNone(sb.test_feature.working_dir)
-
-  def test_acquire_no_pooling_with_error(self):
-    env = TestingEnvironment(
-        features={
-            'test_feature': TestingFeature(
-                simulate_setup_error=interface.SandboxStateError
-            )
-        },
-        pool_size=0,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      with self.assertRaises(interface.EnvironmentOutageError):
-        env.acquire()
-
-  def test_acquire_with_pooling(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=1,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      sb = env.acquire()
-      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
-
-  def test_acquire_with_pooling_overload(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=1,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      sb = env.acquire()
-      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
-      with self.assertRaises(interface.EnvironmentOverloadError):
-        env.acquire()
-
-  def test_acquire_with_growing_pool(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=(1, 3),
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      self.assertEqual(len(env.sandbox_pool['test_image']), 1)
-      self.assertEqual(
-          env.stats(),
-          {
-              'sandbox': {
-                  'test_image': {
-                      'created': 0,
-                      'setting_up': 0,
-                      'ready': 1,
-                      'acquired': 0,
-                      'in_session': 0,
-                      'exiting_session': 0,
-                      'shutting_down': 0,
-                      'offline': 0,
-                  }
-              }
-          }
-      )
-      sb = env.acquire()
-      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
-      self.assertEqual(
-          env.stats(),
-          {
-              'sandbox': {
-                  'test_image': {
-                      'created': 0,
-                      'setting_up': 0,
-                      'ready': 0,
-                      'acquired': 1,
-                      'in_session': 0,
-                      'exiting_session': 0,
-                      'shutting_down': 0,
-                      'offline': 0,
-                  }
-              }
-          }
-      )
-      self.assertEqual(len(env.sandbox_pool['test_image']), 1)
-      sb2 = env.acquire()
-      self.assertEqual(sb2.status, interface.Sandbox.Status.ACQUIRED)
-      self.assertEqual(len(env.sandbox_pool['test_image']), 2)
-      self.assertEqual(
-          env.stats(),
-          {
-              'sandbox': {
-                  'test_image': {
-                      'created': 0,
-                      'setting_up': 0,
-                      'ready': 0,
-                      'acquired': 2,
-                      'in_session': 0,
-                      'exiting_session': 0,
-                      'shutting_down': 0,
-                      'offline': 0,
-                  }
-              }
-          }
-      )
-    self.assertEqual(
-        env.stats(),
-        {
-            'sandbox': {}
-        }
-    )
-
-  def test_acquire_with_growing_pool_failure(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=(1, 3),
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      self.assertEqual(len(env.sandbox_pool), 1)
-      sb = env.acquire()
-      self.assertEqual(sb.status, interface.Sandbox.Status.ACQUIRED)
-
-      # Make future sandbox setup to fail.
-      env.features.test_feature.rebind(
-          simulate_setup_error=interface.SandboxStateError,
-          skip_notification=True
-      )
-      with self.assertRaises(interface.EnvironmentOutageError):
-        env.acquire()
-
-  def test_housekeep_error(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature()},
-        pool_size=1,
-        proactive_session_setup=True,
-        outage_grace_period=1,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    with env:
-      self.assertEqual(len(env.sandbox_pool), 1)
-      self.assertIn('test_image', env.sandbox_pool)
-      self.assertEqual(
-          env.sandbox_pool['test_image'][0].status,
-          interface.Sandbox.Status.READY
-      )
-      # Make future sandbox setup to fail.
-      env.features.test_feature.rebind(
-          simulate_setup_error=interface.SandboxStateError,
-          skip_notification=True
-      )
-      with self.assertRaises(interface.SandboxStateError):
-        with env.sandbox() as sb:
-          sb.shell('bad command', raise_error=interface.SandboxStateError)
-      self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
-      self.assertEqual(len(sb.state_errors), 1)
-      sb_offline_time = time.time()
-      while time.time() - sb_offline_time < 10:
-        if not env.is_online:
-          break
-        time.sleep(0.5)
-      self.assertFalse(env.is_online)
-
-
-class SandboxStatusTests(unittest.TestCase):
+class SandboxStateTests(unittest.TestCase):
 
   def setUp(self):
     super().setUp()
@@ -494,7 +46,7 @@ class SandboxStatusTests(unittest.TestCase):
         pool_size=pool_size,
         features=features,
         outage_grace_period=0,
-        event_handlers=[self.event_handler],
+        event_handler=self.event_handler,
         outage_retry_interval=0,
         **kwargs
     )
@@ -508,39 +60,39 @@ class SandboxStatusTests(unittest.TestCase):
     )
     self.assertFalse(env.enable_pooling('test_image'))
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
       self.assertEqual(
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-              '[testing-env/test_image:0/feature2] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session',
               '[testing-env/test_image:0] setting_up -> in_session',
               "[testing-env/test_image:0] session 'session1' started",
-              '[testing-env/test_image:0/session1] shell: echo "hello"',
+              '[testing-env/test_image:0@session1] shell: echo "hello"',
               '[testing-env/test_image:0] in_session -> exiting_session',
-              '[testing-env/test_image:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0/feature1] feature teardown session',
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0/feature2] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0/feature1@session1] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0] session 'session1' ended",
               '[testing-env/test_image:0] exiting_session -> acquired',
               '[testing-env/test_image:0] acquired -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
-              '[testing-env/test_image:0] shell: "feature2" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
               '[testing-env/test_image:0] sandbox shutdown'
@@ -559,7 +111,7 @@ class SandboxStatusTests(unittest.TestCase):
     )
     self.assertTrue(env.enable_pooling('test_image'))
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
       sb.wait_until_not(
           (
@@ -572,14 +124,14 @@ class SandboxStatusTests(unittest.TestCase):
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
-              '[testing-env/test_image:0:0] shell: "feature1" setup',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0:0/feature1] feature setup',
-              '[testing-env/test_image:0:0] shell: "feature2" setup',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0:0/feature2] feature setup',
-              '[testing-env/test_image:0:0] shell: "feature1" setup session',
-              '[testing-env/test_image:0:0/feature1] feature setup session',
-              '[testing-env/test_image:0:0] shell: "feature2" setup session',
-              '[testing-env/test_image:0:0/feature2] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup session',
+              '[testing-env/test_image:0:0/feature1@<idle>] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature2" setup session',
+              '[testing-env/test_image:0:0/feature2@<idle>] feature setup session',
               '[testing-env/test_image:0:0] created -> ready',
               '[testing-env/test_image:0:0] sandbox started',
               '[testing-env] environment started',
@@ -587,36 +139,31 @@ class SandboxStatusTests(unittest.TestCase):
               '[testing-env/test_image:0:0] acquired -> setting_up',
               '[testing-env/test_image:0:0] setting_up -> in_session',
               "[testing-env/test_image:0:0] session 'session1' started",
-              '[testing-env/test_image:0:0/session1] shell: echo "hello"',
+              '[testing-env/test_image:0:0@session1] shell: echo "hello"',
               '[testing-env/test_image:0:0] in_session -> exiting_session',
-              '[testing-env/test_image:0:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0:0/feature1] feature teardown session',
-              '[testing-env/test_image:0:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0:0/feature2] feature teardown session',
+              '[testing-env/test_image:0:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0:0/feature1@session1] feature teardown session',
+              '[testing-env/test_image:0:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0:0] session 'session1' ended",
               '[testing-env/test_image:0:0] exiting_session -> setting_up',
-              '[testing-env/test_image:0:0] shell: "feature1" setup session',
-              '[testing-env/test_image:0:0/feature1] feature setup session',
-              '[testing-env/test_image:0:0] shell: "feature2" setup session',
-              '[testing-env/test_image:0:0/feature2] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup session',
+              '[testing-env/test_image:0:0/feature1@<idle>] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature2" setup session',
+              '[testing-env/test_image:0:0/feature2@<idle>] feature setup session',
               '[testing-env/test_image:0:0] setting_up -> ready'
               # pylint: enable=line-too-long
           ]
       )
 
-  def test_practive_session_setup_with_setup_session_error(self):
+  def test_proactive_session_setup_with_setup_session_error(self):
     env = self._create_env(
         features={'test_feature': TestingFeature(setup_session_delay=0.5)},
         pool_size=1,
-    )
-    event_handler = TestingEventHandler(
-        log_sandbox_status=True,
-        log_feature_setup=True,
-        log_session_setup=True,
+        housekeep_interval=10.0,
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
-        sb.add_event_handler(event_handler)
+      with env.sandbox('session1') as sb:
         sb.test_feature.rebind(
             simulate_setup_session_error=interface.SandboxStateError,
             skip_notification=True
@@ -630,17 +177,28 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(len(sb.state_errors), 1)
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
       self.assertEqual(
-          event_handler.logs,
+          self.event_handler.logs,
           [
               # pylint: disable=line-too-long
+              '[testing-env/test_image:0:0@<idle>] shell: "test_feature" setup',
+              '[testing-env/test_image:0:0/test_feature] feature setup',
+              '[testing-env/test_image:0:0@<idle>] shell: "test_feature" setup session',
+              '[testing-env/test_image:0:0/test_feature@<idle>] feature setup session',
+              '[testing-env/test_image:0:0] created -> ready',
+              '[testing-env/test_image:0:0] sandbox started',
+              '[testing-env] environment started',
+              '[testing-env/test_image:0:0] ready -> acquired',
+              '[testing-env/test_image:0:0] acquired -> setting_up',
+              '[testing-env/test_image:0:0] setting_up -> in_session',
+              "[testing-env/test_image:0:0] session 'session1' started",
               '[testing-env/test_image:0:0] in_session -> exiting_session',
-              '[testing-env/test_image:0:0/session1] shell: "test_feature" teardown session',
-              '[testing-env/test_image:0:0/test_feature] feature teardown session',
+              '[testing-env/test_image:0:0@session1] shell: "test_feature" teardown session',
+              '[testing-env/test_image:0:0/test_feature@session1] feature teardown session',
               "[testing-env/test_image:0:0] session 'session1' ended",
               '[testing-env/test_image:0:0] exiting_session -> setting_up',
-              '[testing-env/test_image:0:0/test_feature] feature setup session with SandboxStateError',  # pylint: disable=line-too-long
+              '[testing-env/test_image:0:0/test_feature@<idle>] feature setup session with SandboxStateError',
               '[testing-env/test_image:0:0] setting_up -> shutting_down',
-              '[testing-env/test_image:0:0] shell: "test_feature" teardown',
+              '[testing-env/test_image:0:0@<idle>] shell: "test_feature" teardown',
               '[testing-env/test_image:0:0/test_feature] feature teardown',
               '[testing-env/test_image:0:0] shutting_down -> offline',
               '[testing-env/test_image:0:0] sandbox shutdown'
@@ -658,7 +216,7 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(ValueError):
-        with env.sandbox(session_id='session1'):
+        with env.sandbox('session1'):
           pass
       self.assertTrue(env.is_online)
       self.assertEqual(
@@ -708,45 +266,44 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(ValueError):
-        with env.sandbox(session_id='session1') as sb:
+        with env.sandbox('session1') as sb:
           sb.shell('echo "hello"')
       self.assertEqual(len(sb.state_errors), 0)
-
     self.assertEqual(
         self.event_handler.logs,
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "feature1" setup',
+            '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
             '[testing-env/test_image:0/feature1] feature setup',
-            '[testing-env/test_image:0] shell: "feature2" setup',
+            '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
             '[testing-env/test_image:0/feature2] feature setup',
             '[testing-env/test_image:0] created -> ready',
             '[testing-env/test_image:0] sandbox started',
             '[testing-env/test_image:0] ready -> acquired',
             '[testing-env/test_image:0] acquired -> setting_up',
-            '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-            '[testing-env/test_image:0/feature1] feature setup session',
-            '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-            '[testing-env/test_image:0/feature2] feature setup session',
+            '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+            '[testing-env/test_image:0/feature1@session1] feature setup session',
+            '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+            '[testing-env/test_image:0/feature2@session1] feature setup session',
             '[testing-env/test_image:0] setting_up -> in_session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: echo "hello"',
+            '[testing-env/test_image:0@session1] shell: echo "hello"',
             '[testing-env/test_image:0] in_session -> exiting_session',
-            '[testing-env/test_image:0/session1] shell: "feature1" teardown session',
-            '[testing-env/test_image:0/feature1] feature teardown session',
-            '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-            '[testing-env/test_image:0/feature2] feature teardown session',
+            '[testing-env/test_image:0@session1] shell: "feature1" teardown session',
+            '[testing-env/test_image:0/feature1@session1] feature teardown session',
+            '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+            '[testing-env/test_image:0/feature2@session1] feature teardown session',
             "[testing-env/test_image:0] session 'session1' ended",
             '[testing-env/test_image:0] exiting_session -> acquired',
             '[testing-env/test_image:0] acquired -> shutting_down',
-            '[testing-env/test_image:0] shell: "feature1" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
             '[testing-env/test_image:0/feature1] feature teardown',
-            '[testing-env/test_image:0] shell: "feature2" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
             '[testing-env/test_image:0/feature2] feature teardown',
             '[testing-env/test_image:0] shutting_down -> offline',
             '[testing-env/test_image:0] sandbox shutdown with ValueError',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
@@ -768,21 +325,21 @@ class SandboxStatusTests(unittest.TestCase):
         self.event_handler.logs,
         [
             # pylint: disable=line-too-long
-            '[testing-env/test_image:0:0] shell: "feature1" setup',
+            '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup',
             '[testing-env/test_image:0:0/feature1] feature setup',
-            '[testing-env/test_image:0:0] shell: "feature2" setup',
+            '[testing-env/test_image:0:0@<idle>] shell: "feature2" setup',
             '[testing-env/test_image:0:0/feature2] feature setup',
-            '[testing-env/test_image:0:0] shell: "feature1" setup session',
-            '[testing-env/test_image:0:0/feature1] feature setup session',
-            '[testing-env/test_image:0:0] shell: "feature2" setup session',
-            '[testing-env/test_image:0:0/feature2] feature setup session',
+            '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup session',
+            '[testing-env/test_image:0:0/feature1@<idle>] feature setup session',
+            '[testing-env/test_image:0:0@<idle>] shell: "feature2" setup session',
+            '[testing-env/test_image:0:0/feature2@<idle>] feature setup session',
             '[testing-env/test_image:0:0] created -> ready',
             '[testing-env/test_image:0:0] sandbox started',
             '[testing-env] environment started',
             '[testing-env/test_image:0:0] ready -> shutting_down',
-            '[testing-env/test_image:0:0] shell: "feature1" teardown',
+            '[testing-env/test_image:0:0@<idle>] shell: "feature1" teardown',
             '[testing-env/test_image:0:0/feature1] feature teardown',
-            '[testing-env/test_image:0:0] shell: "feature2" teardown',
+            '[testing-env/test_image:0:0@<idle>] shell: "feature2" teardown',
             '[testing-env/test_image:0:0/feature2] feature teardown',
             '[testing-env/test_image:0:0] shutting_down -> offline',
             '[testing-env/test_image:0:0] sandbox shutdown with ValueError',
@@ -800,44 +357,45 @@ class SandboxStatusTests(unittest.TestCase):
         simulate_shutdown_error=interface.SandboxStateError,
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
       self.assertEqual(len(sb.state_errors), 1)
+
     self.assertEqual(
         self.event_handler.logs,
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "feature1" setup',
+            '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
             '[testing-env/test_image:0/feature1] feature setup',
-            '[testing-env/test_image:0] shell: "feature2" setup',
+            '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
             '[testing-env/test_image:0/feature2] feature setup',
             '[testing-env/test_image:0] created -> ready',
             '[testing-env/test_image:0] sandbox started',
             '[testing-env/test_image:0] ready -> acquired',
             '[testing-env/test_image:0] acquired -> setting_up',
-            '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-            '[testing-env/test_image:0/feature1] feature setup session',
-            '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-            '[testing-env/test_image:0/feature2] feature setup session',
+            '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+            '[testing-env/test_image:0/feature1@session1] feature setup session',
+            '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+            '[testing-env/test_image:0/feature2@session1] feature setup session',
             '[testing-env/test_image:0] setting_up -> in_session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: echo "hello"',
+            '[testing-env/test_image:0@session1] shell: echo "hello"',
             '[testing-env/test_image:0] in_session -> exiting_session',
-            '[testing-env/test_image:0/session1] shell: "feature1" teardown session',
-            '[testing-env/test_image:0/feature1] feature teardown session',
-            '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-            '[testing-env/test_image:0/feature2] feature teardown session',
+            '[testing-env/test_image:0@session1] shell: "feature1" teardown session',
+            '[testing-env/test_image:0/feature1@session1] feature teardown session',
+            '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+            '[testing-env/test_image:0/feature2@session1] feature teardown session',
             "[testing-env/test_image:0] session 'session1' ended",
             '[testing-env/test_image:0] exiting_session -> acquired',
             '[testing-env/test_image:0] acquired -> shutting_down',
-            '[testing-env/test_image:0] shell: "feature1" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
             '[testing-env/test_image:0/feature1] feature teardown',
-            '[testing-env/test_image:0] shell: "feature2" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
             '[testing-env/test_image:0/feature2] feature teardown',
             '[testing-env/test_image:0] shutting_down -> offline',
             '[testing-env/test_image:0] sandbox shutdown with SandboxStateError',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
@@ -853,24 +411,24 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(ValueError):
-        with env.sandbox(session_id='session1'):
+        with env.sandbox('session1'):
           pass
       self.assertEqual(
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
               '[testing-env/test_image:0/feature2] feature setup with ValueError',
               '[testing-env/test_image:0] sandbox started with ValueError',
               '[testing-env/test_image:0] created -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
-              '[testing-env/test_image:0] shell: "feature2" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
-              '[testing-env/test_image:0] sandbox shutdown',
+              '[testing-env/test_image:0] sandbox shutdown'
               # pylint: enable=line-too-long
           ]
       )
@@ -886,7 +444,7 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(interface.EnvironmentOutageError):
-        with env.sandbox(session_id='session1'):
+        with env.sandbox('session1'):
           pass
       self.assertEqual(
           self.event_handler.logs,
@@ -896,7 +454,7 @@ class SandboxStatusTests(unittest.TestCase):
               '[testing-env/test_image:0/feature1] feature setup with SandboxStateError',
               '[testing-env/test_image:0] sandbox started with SandboxStateError',
               '[testing-env/test_image:0] created -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
               '[testing-env/test_image:0] sandbox shutdown',
@@ -916,36 +474,36 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(interface.FeatureTeardownError):
-        with env.sandbox(session_id='session1'):
+        with env.sandbox('session1'):
           pass
       self.assertEqual(
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-              '[testing-env/test_image:0/feature2] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session',
               '[testing-env/test_image:0] setting_up -> in_session',
               "[testing-env/test_image:0] session 'session1' started",
               '[testing-env/test_image:0] in_session -> exiting_session',
-              '[testing-env/test_image:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0/feature1] feature teardown session',
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0/feature2] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0/feature1@session1] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0] session 'session1' ended",
               '[testing-env/test_image:0] exiting_session -> acquired',
               '[testing-env/test_image:0] acquired -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
               '[testing-env/test_image:0/feature2] feature teardown with ValueError',
               '[testing-env/test_image:0] shutting_down -> offline',
@@ -965,7 +523,7 @@ class SandboxStatusTests(unittest.TestCase):
         },
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         pass
       self.assertEqual(len(sb.state_errors), 1)
       self.assertEqual(
@@ -973,30 +531,30 @@ class SandboxStatusTests(unittest.TestCase):
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-              '[testing-env/test_image:0/feature2] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session',
               '[testing-env/test_image:0] setting_up -> in_session',
               "[testing-env/test_image:0] session 'session1' started",
               '[testing-env/test_image:0] in_session -> exiting_session',
-              '[testing-env/test_image:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0/feature1] feature teardown session',
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0/feature2] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0/feature1@session1] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0] session 'session1' ended",
               '[testing-env/test_image:0] exiting_session -> acquired',
               '[testing-env/test_image:0] acquired -> shutting_down',
-              '[testing-env/test_image:0/feature1] feature teardown with SandboxStateError',  # pylint: disable=line-too-long
-              '[testing-env/test_image:0] shell: "feature2" teardown',
+              '[testing-env/test_image:0/feature1] feature teardown with SandboxStateError',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
               '[testing-env/test_image:0] sandbox shutdown with FeatureTeardownError',
@@ -1015,32 +573,32 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(ValueError):
-        with env.sandbox(session_id='session1') as sb:
+        with env.sandbox('session1') as sb:
           sb.shell('echo "hello"')
       self.assertEqual(
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/feature2] feature setup session with ValueError',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session with ValueError',
               "[testing-env/test_image:0] session 'session1' started with ValueError",
               '[testing-env/test_image:0] setting_up -> shutting_down',
-              '[testing-env/test_image:0/session1] shell: "feature1" teardown',
+              '[testing-env/test_image:0@session1] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
-              '[testing-env/test_image:0] sandbox shutdown',
+              '[testing-env/test_image:0] sandbox shutdown'
               # pylint: enable=line-too-long
           ]
       )
@@ -1056,7 +614,7 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(interface.SessionTeardownError):
-        with env.sandbox(session_id='session1') as sb:
+        with env.sandbox('session1') as sb:
           sb.shell('echo "hello"')
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
       self.assertEqual(
@@ -1064,34 +622,34 @@ class SandboxStatusTests(unittest.TestCase):
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-              '[testing-env/test_image:0/feature2] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session',
               '[testing-env/test_image:0] setting_up -> in_session',
               "[testing-env/test_image:0] session 'session1' started",
-              '[testing-env/test_image:0/session1] shell: echo "hello"',
+              '[testing-env/test_image:0@session1] shell: echo "hello"',
               '[testing-env/test_image:0] in_session -> exiting_session',
-              '[testing-env/test_image:0/feature1] feature teardown session with ValueError',  # pylint: disable=line-too-long
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0/feature2] feature teardown session',
+              '[testing-env/test_image:0/feature1@session1] feature teardown session with ValueError',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0] session 'session1' ended",
               '[testing-env/test_image:0] exiting_session -> acquired',
               '[testing-env/test_image:0] acquired -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
-              '[testing-env/test_image:0] shell: "feature2" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
-              '[testing-env/test_image:0] sandbox shutdown'
+              '[testing-env/test_image:0] sandbox shutdown',
               # pylint: enable=line-too-long
           ]
       )
@@ -1106,7 +664,7 @@ class SandboxStatusTests(unittest.TestCase):
         },
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
       self.assertEqual(len(sb.state_errors), 1)
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
@@ -1115,34 +673,34 @@ class SandboxStatusTests(unittest.TestCase):
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-              '[testing-env/test_image:0/feature2] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session',
               '[testing-env/test_image:0] setting_up -> in_session',
               "[testing-env/test_image:0] session 'session1' started",
-              '[testing-env/test_image:0/session1] shell: echo "hello"',
+              '[testing-env/test_image:0@session1] shell: echo "hello"',
               '[testing-env/test_image:0] in_session -> exiting_session',
-              '[testing-env/test_image:0/feature1] feature teardown session with SandboxStateError',  # pylint: disable=line-too-long
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0/feature2] feature teardown session',
+              '[testing-env/test_image:0/feature1@session1] feature teardown session with SandboxStateError',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0] session 'session1' ended with SandboxStateError",
               '[testing-env/test_image:0] exiting_session -> acquired',
               '[testing-env/test_image:0] acquired -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
-              '[testing-env/test_image:0] shell: "feature2" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
-              '[testing-env/test_image:0] sandbox shutdown'
+              '[testing-env/test_image:0] sandbox shutdown',
               # pylint: enable=line-too-long
           ]
       )
@@ -1157,39 +715,39 @@ class SandboxStatusTests(unittest.TestCase):
         },
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
       self.assertEqual(
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
               '[testing-env] environment started',
-              '[testing-env/test_image:0] shell: "feature1" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0/feature1] feature setup',
-              '[testing-env/test_image:0] shell: "feature2" setup',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" setup',
               '[testing-env/test_image:0/feature2] feature setup',
               '[testing-env/test_image:0] created -> ready',
               '[testing-env/test_image:0] sandbox started',
               '[testing-env/test_image:0] ready -> acquired',
               '[testing-env/test_image:0] acquired -> setting_up',
-              '[testing-env/test_image:0/session1] shell: "feature1" setup session',
-              '[testing-env/test_image:0/feature1] feature setup session',
-              '[testing-env/test_image:0/session1] shell: "feature2" setup session',
-              '[testing-env/test_image:0/feature2] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature1" setup session',
+              '[testing-env/test_image:0/feature1@session1] feature setup session',
+              '[testing-env/test_image:0@session1] shell: "feature2" setup session',
+              '[testing-env/test_image:0/feature2@session1] feature setup session',
               '[testing-env/test_image:0] setting_up -> in_session',
               "[testing-env/test_image:0] session 'session1' started",
-              '[testing-env/test_image:0/session1] shell: echo "hello"',
+              '[testing-env/test_image:0@session1] shell: echo "hello"',
               '[testing-env/test_image:0] in_session -> exiting_session',
-              '[testing-env/test_image:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0/feature1] feature teardown session',
-              '[testing-env/test_image:0/session1] shell: "feature2" teardown session',
-              '[testing-env/test_image:0/feature2] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0/feature1@session1] feature teardown session',
+              '[testing-env/test_image:0@session1] shell: "feature2" teardown session',
+              '[testing-env/test_image:0/feature2@session1] feature teardown session',
               "[testing-env/test_image:0] session 'session1' ended",
               '[testing-env/test_image:0] exiting_session -> acquired',
               '[testing-env/test_image:0] acquired -> shutting_down',
-              '[testing-env/test_image:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0/feature1] feature teardown',
-              '[testing-env/test_image:0] shell: "feature2" teardown',
+              '[testing-env/test_image:0@<idle>] shell: "feature2" teardown',
               '[testing-env/test_image:0/feature2] feature teardown',
               '[testing-env/test_image:0] shutting_down -> offline',
               '[testing-env/test_image:0] sandbox shutdown'
@@ -1205,7 +763,7 @@ class SandboxStatusTests(unittest.TestCase):
         },
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         with self.assertRaises(ValueError):
           sb.shell('echo foo', raise_error=ValueError)
         self.assertEqual(len(sb.state_errors), 0)
@@ -1217,10 +775,10 @@ class SandboxStatusTests(unittest.TestCase):
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
-              '[testing-env/test_image:0:0] shell: "feature1" setup',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0:0/feature1] feature setup',
-              '[testing-env/test_image:0:0] shell: "feature1" setup session',
-              '[testing-env/test_image:0:0/feature1] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup session',
+              '[testing-env/test_image:0:0/feature1@<idle>] feature setup session',
               '[testing-env/test_image:0:0] created -> ready',
               '[testing-env/test_image:0:0] sandbox started',
               '[testing-env] environment started',
@@ -1228,15 +786,15 @@ class SandboxStatusTests(unittest.TestCase):
               '[testing-env/test_image:0:0] acquired -> setting_up',
               '[testing-env/test_image:0:0] setting_up -> in_session',
               "[testing-env/test_image:0:0] session 'session1' started",
-              '[testing-env/test_image:0:0/session1] shell: echo foo with ValueError',
-              '[testing-env/test_image:0:0/session1] shell: echo bar',
+              '[testing-env/test_image:0:0@session1] shell: echo foo with ValueError',
+              '[testing-env/test_image:0:0@session1] shell: echo bar',
               '[testing-env/test_image:0:0] in_session -> exiting_session',
-              '[testing-env/test_image:0:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0:0/feature1] feature teardown session',
+              '[testing-env/test_image:0:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0:0/feature1@session1] feature teardown session',
               "[testing-env/test_image:0:0] session 'session1' ended",
               '[testing-env/test_image:0:0] exiting_session -> setting_up',
-              '[testing-env/test_image:0:0] shell: "feature1" setup session',
-              '[testing-env/test_image:0:0/feature1] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup session',
+              '[testing-env/test_image:0:0/feature1@<idle>] feature setup session',
               '[testing-env/test_image:0:0] setting_up -> ready',
               # pylint: enable=line-too-long
           ]
@@ -1251,7 +809,7 @@ class SandboxStatusTests(unittest.TestCase):
     )
     with env:
       with self.assertRaises(interface.SandboxStateError):
-        with env.sandbox(session_id='session1') as sb:
+        with env.sandbox('session1') as sb:
           sb.shell('echo foo', raise_error=RuntimeError)
       self.assertEqual(len(sb.state_errors), 1)
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
@@ -1259,10 +817,10 @@ class SandboxStatusTests(unittest.TestCase):
           self.event_handler.logs,
           [
               # pylint: disable=line-too-long
-              '[testing-env/test_image:0:0] shell: "feature1" setup',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup',
               '[testing-env/test_image:0:0/feature1] feature setup',
-              '[testing-env/test_image:0:0] shell: "feature1" setup session',
-              '[testing-env/test_image:0:0/feature1] feature setup session',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" setup session',
+              '[testing-env/test_image:0:0/feature1@<idle>] feature setup session',
               '[testing-env/test_image:0:0] created -> ready',
               '[testing-env/test_image:0:0] sandbox started',
               '[testing-env] environment started',
@@ -1270,14 +828,14 @@ class SandboxStatusTests(unittest.TestCase):
               '[testing-env/test_image:0:0] acquired -> setting_up',
               '[testing-env/test_image:0:0] setting_up -> in_session',
               "[testing-env/test_image:0:0] session 'session1' started",
-              '[testing-env/test_image:0:0/session1] shell: echo foo with RuntimeError',
+              '[testing-env/test_image:0:0@session1] shell: echo foo with RuntimeError',
               '[testing-env/test_image:0:0] in_session -> exiting_session',
-              '[testing-env/test_image:0:0/session1] shell: "feature1" teardown session',
-              '[testing-env/test_image:0:0/feature1] feature teardown session',
+              '[testing-env/test_image:0:0@session1] shell: "feature1" teardown session',
+              '[testing-env/test_image:0:0/feature1@session1] feature teardown session',
               "[testing-env/test_image:0:0] session 'session1' ended with SandboxStateError",
               '[testing-env/test_image:0:0] exiting_session -> acquired',
               '[testing-env/test_image:0:0] acquired -> shutting_down',
-              '[testing-env/test_image:0:0] shell: "feature1" teardown',
+              '[testing-env/test_image:0:0@<idle>] shell: "feature1" teardown',
               '[testing-env/test_image:0:0/feature1] feature teardown',
               '[testing-env/test_image:0:0] shutting_down -> offline',
               '[testing-env/test_image:0:0] sandbox shutdown',
@@ -1311,7 +869,7 @@ class SandboxActivityTests(unittest.TestCase):
         sandbox_keepalive_interval=0,
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         sb.rebind(
             simulate_ping_error=interface.SandboxStateError,
             skip_notification=True
@@ -1328,10 +886,10 @@ class SandboxActivityTests(unittest.TestCase):
         outage_grace_period=0,
         outage_retry_interval=0.1,
         sandbox_keepalive_interval=0,
-        event_handlers=[event_handler],
+        event_handler=event_handler,
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
+      with env.sandbox('session1') as sb:
         self.assertEqual(len(env.sandbox_pool), 1)
         self.assertEqual(sb.status, interface.Sandbox.Status.IN_SESSION)
         self.assertEqual(sb.session_id, 'session1')
@@ -1351,22 +909,22 @@ class SandboxActivityTests(unittest.TestCase):
         event_handler.logs,
         [
             # pylint: disable=line-too-long
-            '[testing-env/test_image:0:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0:0/test_feature] feature setup',
-            '[testing-env/test_image:0:0] shell: "test_feature" setup session',
+            '[testing-env/test_image:0:0@<idle>] shell: "test_feature" setup session',
             '[testing-env/test_image:0:0] sandbox started',
             '[testing-env] environment started',
             "[testing-env/test_image:0:0] session 'session1' started",
-            '[testing-env/test_image:0:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0:0] session 'session1' ended with SandboxStateError",
-            '[testing-env/test_image:0:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0:0/test_feature] feature teardown',
             '[testing-env/test_image:0:0] sandbox shutdown',
-            '[testing-env/test_image:0:1] shell: "test_feature" setup',
+            '[testing-env/test_image:0:1@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0:1/test_feature] feature setup',
-            '[testing-env/test_image:0:1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0:1@<idle>] shell: "test_feature" setup session',
             '[testing-env/test_image:0:1] sandbox started',
-            '[testing-env/test_image:0:1] shell: "test_feature" teardown',
+            '[testing-env/test_image:0:1@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0:1/test_feature] feature teardown',
             '[testing-env/test_image:0:1] sandbox shutdown',
             '[testing-env] environment shutdown'
@@ -1374,27 +932,6 @@ class SandboxActivityTests(unittest.TestCase):
         ]
     )
 
-  def test_remove_event_handler(self):
-    env = TestingEnvironment(
-        features={'test_feature': TestingFeature(housekeep_interval=0)},
-        pool_size=1,
-        outage_grace_period=0,
-        outage_retry_interval=0,
-        sandbox_keepalive_interval=0,
-    )
-    event_handler = TestingEventHandler()
-    with env:
-      with env.sandbox(session_id='session1') as sb:
-        sb.add_event_handler(event_handler)
-        sb.shell('test_feature')
-        sb.remove_event_handler(event_handler)
-      events = list(event_handler.logs)
-      sb.wait_until_not(interface.Sandbox.Status.SETTING_UP)
-      self.assertGreater(len(events), 0)
-      with env.sandbox(session_id='session2') as sb:
-        sb.shell('test_feature')
-      self.assertEqual(len(events), len(event_handler.logs))
-
 
 class SandboxServiceTests(unittest.TestCase):
 
@@ -1408,13 +945,13 @@ class SandboxServiceTests(unittest.TestCase):
         outage_grace_period=0,
         outage_retry_interval=0,
         sandbox_keepalive_interval=0,
-        event_handlers=[self.event_handler],
+        event_handler=self.event_handler,
         random_seed=1,
     )
 
   def test_service_call_activity_log(self):
 
-    class CustomEventHandler(event_handler_base.EventHandler):
+    class CustomEventHandler(interface.EventHandler):
 
       def __init__(self):
         self.calls = []
@@ -1422,9 +959,17 @@ class SandboxServiceTests(unittest.TestCase):
       def on_sandbox_activity(
           self,
           name: str,
-          environment: interface.Environment,
           sandbox: interface.Sandbox,
-          feature: interface.Feature | None,
+          session_id: str | None,
+          duration: float,
+          error: BaseException | None,
+          **kwargs: Any):
+        self.calls.append((session_id, name, kwargs))
+
+      def on_feature_activity(
+          self,
+          name: str,
+          feature: interface.Feature,
           session_id: str | None,
           duration: float,
           error: BaseException | None,
@@ -1435,7 +980,7 @@ class SandboxServiceTests(unittest.TestCase):
     env = TestingEnvironment(
         features={'test_feature': TestingFeature()},
         pool_size=0,
-        event_handlers=[event_handler],
+        event_handler=event_handler,
     )
     with env:
       with env.test_feature(session_id='session1') as test_feature:
@@ -1453,7 +998,7 @@ class SandboxServiceTests(unittest.TestCase):
 
   def test_service_call_from_feature(self):
     with self.env:
-      with self.env.sandbox(session_id='session1') as sb:
+      with self.env.sandbox('session1') as sb:
         self.assertEqual(sb.test_feature.num_shell_calls(), 2)
         self.assertEqual(sb.test_feature.num_shell_calls(), 2)
     self.assertEqual(
@@ -1461,19 +1006,19 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/session1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@session1] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1/test_feature] test_feature.num_shell_calls: None',
-            '[testing-env/test_image:0/session1/test_feature] test_feature.num_shell_calls: None',
-            '[testing-env/test_image:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0/test_feature@session1] test_feature.num_shell_calls: None',
+            '[testing-env/test_image:0/test_feature@session1] test_feature.num_shell_calls: None',
+            '[testing-env/test_image:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'session1' ended",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
@@ -1481,26 +1026,25 @@ class SandboxServiceTests(unittest.TestCase):
   def test_service_call_from_feature_with_error(self):
     with self.env:
       with self.assertRaises(interface.SandboxStateError):
-        with self.env.sandbox(session_id='session1') as sb:
+        with self.env.sandbox('session1') as sb:
           sb.test_feature.bad_shell_call()
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
       self.assertEqual(len(sb.state_errors), 1)
-
     self.assertEqual(
         self.event_handler.logs,
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/session1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@session1] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: bad command with RuntimeError',
-            '[testing-env/test_image:0/session1/test_feature] test_feature.bad_shell_call: None with SandboxStateError',
-            '[testing-env/test_image:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0@session1] shell: bad command with RuntimeError',
+            '[testing-env/test_image:0/test_feature@session1] test_feature.bad_shell_call: None with SandboxStateError',
+            '[testing-env/test_image:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'session1' ended with SandboxStateError",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
             '[testing-env] environment shutdown'
@@ -1517,15 +1061,15 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/test_feature-session-2291d8c] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@test_feature-session-2291d8c] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'test_feature-session-2291d8c' started",
-            '[testing-env/test_image:0/test_feature-session-2291d8c/test_feature] test_feature.num_shell_calls: None',
-            '[testing-env/test_image:0/test_feature-session-2291d8c] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0/test_feature@test_feature-session-2291d8c] test_feature.num_shell_calls: None',
+            '[testing-env/test_image:0@test_feature-session-2291d8c] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'test_feature-session-2291d8c' ended",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
             '[testing-env] environment shutdown'
@@ -1543,26 +1087,26 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/session1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@session1] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: bad command with RuntimeError',
-            '[testing-env/test_image:0/session1/test_feature] test_feature.bad_shell_call: None with SandboxStateError',
-            '[testing-env/test_image:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0@session1] shell: bad command with RuntimeError',
+            '[testing-env/test_image:0/test_feature@session1] test_feature.bad_shell_call: None with SandboxStateError',
+            '[testing-env/test_image:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'session1' ended with SandboxStateError",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
 
   def test_service_context_manager_from_feature(self):
     with self.env:
-      with self.env.sandbox(session_id='session1') as sb:
+      with self.env.sandbox('session1') as sb:
         with sb.test_feature.my_service() as service:
           service.do('hello')
         sb.shell('foo')
@@ -1572,19 +1116,19 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/session1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@session1] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: hello',
-            '[testing-env/test_image:0/session1] shell: foo',
-            '[testing-env/test_image:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0@session1] shell: hello',
+            '[testing-env/test_image:0@session1] shell: foo',
+            '[testing-env/test_image:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'session1' ended",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
@@ -1592,7 +1136,7 @@ class SandboxServiceTests(unittest.TestCase):
   def test_service_context_manager_from_feature_with_error(self):
     with self.env:
       with self.assertRaises(interface.SandboxStateError):
-        with self.env.sandbox(session_id='session1') as sb:
+        with self.env.sandbox('session1') as sb:
           with sb.test_feature.my_service() as service:
             service.do('hello', raise_error=interface.SandboxStateError)
       self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
@@ -1602,18 +1146,18 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/session1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@session1] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: hello with SandboxStateError',
-            '[testing-env/test_image:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0@session1] shell: hello with SandboxStateError',
+            '[testing-env/test_image:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'session1' ended with SandboxStateError",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
@@ -1632,29 +1176,29 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/session1] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@session1] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'session1' started",
-            '[testing-env/test_image:0/session1] shell: foo',
-            '[testing-env/test_image:0/session1] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0@session1] shell: foo',
+            '[testing-env/test_image:0@session1] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'session1' ended",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
-            '[testing-env/test_image:1] shell: "test_feature" setup',
+            '[testing-env/test_image:1@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:1/test_feature] feature setup',
             '[testing-env/test_image:1] sandbox started',
-            '[testing-env/test_image:1/test_feature-session-2291d8c] shell: "test_feature" setup session',
+            '[testing-env/test_image:1@test_feature-session-2291d8c] shell: "test_feature" setup session',
             "[testing-env/test_image:1] session 'test_feature-session-2291d8c' started",
-            '[testing-env/test_image:1/test_feature-session-2291d8c] shell: bar',
-            '[testing-env/test_image:1/test_feature-session-2291d8c] shell: "test_feature" teardown session',
+            '[testing-env/test_image:1@test_feature-session-2291d8c] shell: bar',
+            '[testing-env/test_image:1@test_feature-session-2291d8c] shell: "test_feature" teardown session',
             "[testing-env/test_image:1] session 'test_feature-session-2291d8c' ended",
-            '[testing-env/test_image:1] shell: "test_feature" teardown',
+            '[testing-env/test_image:1@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:1/test_feature] feature teardown',
             '[testing-env/test_image:1] sandbox shutdown',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )
@@ -1670,18 +1214,18 @@ class SandboxServiceTests(unittest.TestCase):
         [
             # pylint: disable=line-too-long
             '[testing-env] environment started',
-            '[testing-env/test_image:0] shell: "test_feature" setup',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" setup',
             '[testing-env/test_image:0/test_feature] feature setup',
             '[testing-env/test_image:0] sandbox started',
-            '[testing-env/test_image:0/test_feature-session-2291d8c] shell: "test_feature" setup session',
+            '[testing-env/test_image:0@test_feature-session-2291d8c] shell: "test_feature" setup session',
             "[testing-env/test_image:0] session 'test_feature-session-2291d8c' started",
-            '[testing-env/test_image:0/test_feature-session-2291d8c] shell: hello with SandboxStateError',
-            '[testing-env/test_image:0/test_feature-session-2291d8c] shell: "test_feature" teardown session',
+            '[testing-env/test_image:0@test_feature-session-2291d8c] shell: hello with SandboxStateError',
+            '[testing-env/test_image:0@test_feature-session-2291d8c] shell: "test_feature" teardown session',
             "[testing-env/test_image:0] session 'test_feature-session-2291d8c' ended with SandboxStateError",
-            '[testing-env/test_image:0] shell: "test_feature" teardown',
+            '[testing-env/test_image:0@<idle>] shell: "test_feature" teardown',
             '[testing-env/test_image:0/test_feature] feature teardown',
             '[testing-env/test_image:0] sandbox shutdown',
-            '[testing-env] environment shutdown'
+            '[testing-env] environment shutdown',
             # pylint: enable=line-too-long
         ]
     )