langfun 0.1.2.dev202509280803__py3-none-any.whl → 0.1.2.dev202509300805__py3-none-any.whl

langfun/env/base_environment.py

@@ -221,7 +221,7 @@ class BaseEnvironment(interface.Environment):
       self._sandbox_pool = [None] * self.min_pool_size
       for i, sandbox, _ in lf.concurrent_map(
           lambda i: self._bring_up_sandbox_with_retry(
-              sandbox_id=str(i), shutdown_env_upon_outage=False
+              sandbox_id=f'{i}:0', shutdown_env_upon_outage=False
           ),
           range(self.min_pool_size),
           silence_on_errors=None,
@@ -335,6 +335,7 @@ class BaseEnvironment(interface.Environment):
         f'it is in {self._status.value!r} status.'
     )
 
+    self.on_starting()
     starting_time = time.time()
     try:
       self._start()
@@ -359,11 +360,12 @@ class BaseEnvironment(interface.Environment):
 
     self._set_status(self.Status.SHUTTING_DOWN)
 
+    shutting_down_time = time.time()
     try:
       self._shutdown()
-      self.on_shutdown()
+      self.on_shutdown(duration=time.time() - shutting_down_time)
     except BaseException as e:  # pylint: disable=broad-except
-      self.on_shutdown(error=e)
+      self.on_shutdown(duration=time.time() - shutting_down_time, error=e)
       raise e
 
   #
@@ -412,7 +414,8 @@ class BaseEnvironment(interface.Environment):
         else:
           try:
             sandbox = self._bring_up_sandbox(
-                sandbox_id=str(self._increment_sandbox_id()), set_acquired=True,
+                sandbox_id=f'{self._increment_sandbox_id()}:0',
+                set_acquired=True,
             )
             # Append is atomic and does not require locking.
             self._sandbox_pool.append(sandbox)
@@ -544,8 +547,9 @@ class BaseEnvironment(interface.Environment):
         len(dead_pool_indices),
     )
     def _replace(i: int):
+      generation = int(self._sandbox_pool[i].id.sandbox_id.split(':')[1])
       self._sandbox_pool[i] = self._bring_up_sandbox_with_retry(
-          str(i), shutdown_env_upon_outage=False
+          f'{i}:{generation + 1}', shutdown_env_upon_outage=False
       )
 
     # TODO(daiyip): Consider to loose the condition to allow some dead
@@ -564,6 +568,11 @@ class BaseEnvironment(interface.Environment):
   # Event handlers subclasses can override.
   #
 
+  def on_starting(self) -> None:
+    """Called when the environment is getting started."""
+    for handler in self.event_handlers:
+      handler.on_environment_starting(self)
+
   def on_start(
       self,
       duration: float, error: BaseException | None = None
@@ -582,8 +591,11 @@ class BaseEnvironment(interface.Environment):
     for handler in self.event_handlers:
       handler.on_environment_housekeep(self, housekeep_counter, duration, error)
 
-  def on_shutdown(self, error: BaseException | None = None) -> None:
+  def on_shutdown(
+      self,
+      duration: float,
+      error: BaseException | None = None) -> None:
     """Called when the environment is shutdown."""
     lifetime = (time.time() - self.start_time) if self.start_time else 0.0
     for handler in self.event_handlers:
-      handler.on_environment_shutdown(self, lifetime, error)
+      handler.on_environment_shutdown(self, duration, lifetime, error)

langfun/env/base_sandbox.py

@@ -418,6 +418,7 @@ class BaseSandbox(interface.Sandbox):
       self.end_session(shutdown_sandbox=True)
       return
 
+    shutting_down_time = time.time()
     self._set_status(interface.Sandbox.Status.SHUTTING_DOWN)
 
     if (self._housekeep_thread is not None
@@ -429,7 +430,10 @@ class BaseSandbox(interface.Sandbox):
     try:
       self._shutdown()
       self._set_status(interface.Sandbox.Status.OFFLINE)
-      self.on_shutdown(teardown_error)
+      self.on_shutdown(
+          duration=time.time() - shutting_down_time,
+          error=teardown_error
+      )
       shutdown_error = None
     except BaseException as e:  # pylint: disable=broad-except
       shutdown_error = e
@@ -439,7 +443,10 @@ class BaseSandbox(interface.Sandbox):
           '[%s]: Sandbox shutdown with error: %s',
           self.id, e
       )
-      self.on_shutdown(teardown_error or shutdown_error)
+      self.on_shutdown(
+          duration=time.time() - shutting_down_time,
+          error=teardown_error or shutdown_error
+      )
 
     # We raise non-state errors to the user following timely order, so the user
     # code could be surfaced and handled properly.
@@ -591,6 +598,7 @@ class BaseSandbox(interface.Sandbox):
     # Set sandbox status to EXITING_SESSION to avoid re-entry.
     self._set_status(self.Status.EXITING_SESSION)
     shutdown_sandbox = shutdown_sandbox or not self.reusable
+    ending_time = time.time()
 
     # Teardown features for the current session.
     end_session_error = self._end_session()
@@ -617,7 +625,9 @@ class BaseSandbox(interface.Sandbox):
             self.shutdown()
 
         # End session before setting up the next session.
-        self.on_session_end(previous_session_id)
+        self.on_session_end(
+            previous_session_id, duration=time.time() - ending_time
+        )
 
         # Mark the sandbox as setting up to prevent it from being acquired by
         # other threads.
@@ -628,7 +638,9 @@ class BaseSandbox(interface.Sandbox):
         threading.Thread(target=_setup_next_session).start()
       else:
         # End session before reporting sandbox status change.
-        self.on_session_end(previous_session_id)
+        self.on_session_end(
+            previous_session_id, duration=time.time() - ending_time
+        )
 
         # If shutdown is requested, mark the sandbox as acquired to prevent it
         # from being acquired by other threads.
@@ -639,7 +651,11 @@ class BaseSandbox(interface.Sandbox):
 
     # Otherwise, shutdown the sandbox.
     else:
-      self.on_session_end(previous_session_id, self.state_errors[0])
+      self.on_session_end(
+          previous_session_id,
+          duration=time.time() - ending_time,
+          error=self.state_errors[0]
+      )
       self._set_status(interface.Sandbox.Status.ACQUIRED)
       shutdown_sandbox = True
 
@@ -762,14 +778,20 @@ class BaseSandbox(interface.Sandbox):
           time.time() - self._status_start_time
       )
 
-  def on_shutdown(self, error: BaseException | None = None) -> None:
+  def on_shutdown(
+      self,
+      duration: float,
+      error: BaseException | None = None
+  ) -> None:
     """Called when the sandbox is shutdown."""
     if self._start_time is None:
       lifetime = 0.0
     else:
       lifetime = time.time() - self._start_time
     for handler in self._event_handlers:
-      handler.on_sandbox_shutdown(self.environment, self, lifetime, error)
+      handler.on_sandbox_shutdown(
+          self.environment, self, duration, lifetime, error
+      )
 
   def on_housekeep(
       self,
@@ -880,13 +902,14 @@ class BaseSandbox(interface.Sandbox):
   def on_session_end(
       self,
       session_id: str,
+      duration: float,
       error: BaseException | None = None
   ) -> None:
     """Called when the user session ends."""
     lifetime = time.time() - self._session_start_time
     for handler in self._event_handlers:
       handler.on_session_end(
-          self.environment, self, session_id, lifetime, error
+          self.environment, self, session_id, duration, lifetime, error
       )
 
 

langfun/env/base_test.py

@@ -381,34 +381,34 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           self.event_handler.logs,
           [
-              '[testing-env/0] shell: "feature1" setup',
-              '[testing-env/0/feature1] feature setup',
-              '[testing-env/0] shell: "feature2" setup',
-              '[testing-env/0/feature2] feature setup',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] shell: "feature2" setup session',
-              '[testing-env/0/feature2] feature setup session',
-              '[testing-env/0] created -> ready',
-              '[testing-env/0] sandbox started',
+              '[testing-env/0:0] shell: "feature1" setup',
+              '[testing-env/0:0/feature1] feature setup',
+              '[testing-env/0:0] shell: "feature2" setup',
+              '[testing-env/0:0/feature2] feature setup',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] shell: "feature2" setup session',
+              '[testing-env/0:0/feature2] feature setup session',
+              '[testing-env/0:0] created -> ready',
+              '[testing-env/0:0] sandbox started',
               '[testing-env] environment started',
-              '[testing-env/0] ready -> acquired',
-              '[testing-env/0] acquired -> setting_up',
-              '[testing-env/0] setting_up -> in_session',
-              "[testing-env/0] session 'session1' started",
-              '[testing-env/0/session1] shell: echo "hello"',
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "feature1" teardown session',
-              '[testing-env/0/feature1] feature teardown session',
-              '[testing-env/0/session1] shell: "feature2" teardown session',
-              '[testing-env/0/feature2] feature teardown session',
-              "[testing-env/0] session 'session1' ended",
-              '[testing-env/0] exiting_session -> setting_up',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] shell: "feature2" setup session',
-              '[testing-env/0/feature2] feature setup session',
-              '[testing-env/0] setting_up -> ready'
+              '[testing-env/0:0] ready -> acquired',
+              '[testing-env/0:0] acquired -> setting_up',
+              '[testing-env/0:0] setting_up -> in_session',
+              "[testing-env/0:0] session 'session1' started",
+              '[testing-env/0:0/session1] shell: echo "hello"',
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "feature1" teardown session',
+              '[testing-env/0:0/feature1] feature teardown session',
+              '[testing-env/0:0/session1] shell: "feature2" teardown session',
+              '[testing-env/0:0/feature2] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended",
+              '[testing-env/0:0] exiting_session -> setting_up',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] shell: "feature2" setup session',
+              '[testing-env/0:0/feature2] feature setup session',
+              '[testing-env/0:0] setting_up -> ready'
           ]
       )
 
@@ -440,17 +440,19 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           event_handler.logs,
           [
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "test_feature" teardown session',
-              '[testing-env/0/test_feature] feature teardown session',
-              "[testing-env/0] session 'session1' ended",
-              '[testing-env/0] exiting_session -> setting_up',
-              '[testing-env/0/test_feature] feature setup session with SandboxStateError',  # pylint: disable=line-too-long
-              '[testing-env/0] setting_up -> shutting_down',
-              '[testing-env/0] shell: "test_feature" teardown',
-              '[testing-env/0/test_feature] feature teardown',
-              '[testing-env/0] shutting_down -> offline',
-              '[testing-env/0] sandbox shutdown'
+              # pylint: disable=line-too-long
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "test_feature" teardown session',
+              '[testing-env/0:0/test_feature] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended",
+              '[testing-env/0:0] exiting_session -> setting_up',
+              '[testing-env/0:0/test_feature] feature setup session with SandboxStateError',  # pylint: disable=line-too-long
+              '[testing-env/0:0] setting_up -> shutting_down',
+              '[testing-env/0:0] shell: "test_feature" teardown',
+              '[testing-env/0:0/test_feature] feature teardown',
+              '[testing-env/0:0] shutting_down -> offline',
+              '[testing-env/0:0] sandbox shutdown'
+              # pylint: enable=line-too-long
           ]
       )
 
@@ -493,10 +495,10 @@ class SandboxStatusTests(unittest.TestCase):
     self.assertEqual(
         self.event_handler.logs,
         [
-            '[testing-env/0] sandbox started with SandboxStateError',
-            '[testing-env/0] created -> shutting_down',
-            '[testing-env/0] shutting_down -> offline',
-            '[testing-env/0] sandbox shutdown',
+            '[testing-env/0:0] sandbox started with SandboxStateError',
+            '[testing-env/0:0] created -> shutting_down',
+            '[testing-env/0:0] shutting_down -> offline',
+            '[testing-env/0:0] sandbox shutdown',
             '[testing-env] environment started with EnvironmentOutageError',
             '[testing-env] environment shutdown'
         ]
@@ -569,24 +571,24 @@ class SandboxStatusTests(unittest.TestCase):
     self.assertEqual(
         self.event_handler.logs,
         [
-            '[testing-env/0] shell: "feature1" setup',
-            '[testing-env/0/feature1] feature setup',
-            '[testing-env/0] shell: "feature2" setup',
-            '[testing-env/0/feature2] feature setup',
-            '[testing-env/0] shell: "feature1" setup session',
-            '[testing-env/0/feature1] feature setup session',
-            '[testing-env/0] shell: "feature2" setup session',
-            '[testing-env/0/feature2] feature setup session',
-            '[testing-env/0] created -> ready',
-            '[testing-env/0] sandbox started',
+            '[testing-env/0:0] shell: "feature1" setup',
+            '[testing-env/0:0/feature1] feature setup',
+            '[testing-env/0:0] shell: "feature2" setup',
+            '[testing-env/0:0/feature2] feature setup',
+            '[testing-env/0:0] shell: "feature1" setup session',
+            '[testing-env/0:0/feature1] feature setup session',
+            '[testing-env/0:0] shell: "feature2" setup session',
+            '[testing-env/0:0/feature2] feature setup session',
+            '[testing-env/0:0] created -> ready',
+            '[testing-env/0:0] sandbox started',
             '[testing-env] environment started',
-            '[testing-env/0] ready -> shutting_down',
-            '[testing-env/0] shell: "feature1" teardown',
-            '[testing-env/0/feature1] feature teardown',
-            '[testing-env/0] shell: "feature2" teardown',
-            '[testing-env/0/feature2] feature teardown',
-            '[testing-env/0] shutting_down -> offline',
-            '[testing-env/0] sandbox shutdown with ValueError',
+            '[testing-env/0:0] ready -> shutting_down',
+            '[testing-env/0:0] shell: "feature1" teardown',
+            '[testing-env/0:0/feature1] feature teardown',
+            '[testing-env/0:0] shell: "feature2" teardown',
+            '[testing-env/0:0/feature2] feature teardown',
+            '[testing-env/0:0] shutting_down -> offline',
+            '[testing-env/0:0] sandbox shutdown with ValueError',
             '[testing-env] environment shutdown with ValueError'
         ]
     )
@@ -998,27 +1000,27 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           self.event_handler.logs,
           [
-              '[testing-env/0] shell: "feature1" setup',
-              '[testing-env/0/feature1] feature setup',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] created -> ready',
-              '[testing-env/0] sandbox started',
+              '[testing-env/0:0] shell: "feature1" setup',
+              '[testing-env/0:0/feature1] feature setup',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] created -> ready',
+              '[testing-env/0:0] sandbox started',
               '[testing-env] environment started',
-              '[testing-env/0] ready -> acquired',
-              '[testing-env/0] acquired -> setting_up',
-              '[testing-env/0] setting_up -> in_session',
-              "[testing-env/0] session 'session1' started",
-              '[testing-env/0/session1] shell: echo foo with ValueError',
-              '[testing-env/0/session1] shell: echo bar',
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "feature1" teardown session',
-              '[testing-env/0/feature1] feature teardown session',
-              "[testing-env/0] session 'session1' ended",
-              '[testing-env/0] exiting_session -> setting_up',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] setting_up -> ready',
+              '[testing-env/0:0] ready -> acquired',
+              '[testing-env/0:0] acquired -> setting_up',
+              '[testing-env/0:0] setting_up -> in_session',
+              "[testing-env/0:0] session 'session1' started",
+              '[testing-env/0:0/session1] shell: echo foo with ValueError',
+              '[testing-env/0:0/session1] shell: echo bar',
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "feature1" teardown session',
+              '[testing-env/0:0/feature1] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended",
+              '[testing-env/0:0] exiting_session -> setting_up',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] setting_up -> ready',
           ]
       )
 
@@ -1039,29 +1041,31 @@ class SandboxStatusTests(unittest.TestCase):
       self.assertEqual(
           self.event_handler.logs,
           [
-              '[testing-env/0] shell: "feature1" setup',
-              '[testing-env/0/feature1] feature setup',
-              '[testing-env/0] shell: "feature1" setup session',
-              '[testing-env/0/feature1] feature setup session',
-              '[testing-env/0] created -> ready',
-              '[testing-env/0] sandbox started',
+              # pylint: disable=line-too-long
+              '[testing-env/0:0] shell: "feature1" setup',
+              '[testing-env/0:0/feature1] feature setup',
+              '[testing-env/0:0] shell: "feature1" setup session',
+              '[testing-env/0:0/feature1] feature setup session',
+              '[testing-env/0:0] created -> ready',
+              '[testing-env/0:0] sandbox started',
               '[testing-env] environment started',
-              '[testing-env/0] ready -> acquired',
-              '[testing-env/0] acquired -> setting_up',
-              '[testing-env/0] setting_up -> in_session',
-              "[testing-env/0] session 'session1' started",
-              '[testing-env/0/session1] shell: echo foo with RuntimeError',
-              '[testing-env/0] in_session -> exiting_session',
-              '[testing-env/0/session1] shell: "feature1" teardown session',
-              '[testing-env/0/feature1] feature teardown session',
-              "[testing-env/0] session 'session1' ended with SandboxStateError",
-              '[testing-env/0] exiting_session -> acquired',
-              '[testing-env/0] acquired -> shutting_down',
-              '[testing-env/0] shell: "feature1" teardown',
-              '[testing-env/0/feature1] feature teardown',
-              '[testing-env/0] shutting_down -> offline',
-              '[testing-env/0] sandbox shutdown',
-              '[testing-env/0] shell: echo bar',
+              '[testing-env/0:0] ready -> acquired',
+              '[testing-env/0:0] acquired -> setting_up',
+              '[testing-env/0:0] setting_up -> in_session',
+              "[testing-env/0:0] session 'session1' started",
+              '[testing-env/0:0/session1] shell: echo foo with RuntimeError',
+              '[testing-env/0:0] in_session -> exiting_session',
+              '[testing-env/0:0/session1] shell: "feature1" teardown session',
+              '[testing-env/0:0/feature1] feature teardown session',
+              "[testing-env/0:0] session 'session1' ended with SandboxStateError",
+              '[testing-env/0:0] exiting_session -> acquired',
+              '[testing-env/0:0] acquired -> shutting_down',
+              '[testing-env/0:0] shell: "feature1" teardown',
+              '[testing-env/0:0/feature1] feature teardown',
+              '[testing-env/0:0] shutting_down -> offline',
+              '[testing-env/0:0] sandbox shutdown',
+              '[testing-env/0:0] shell: echo bar',
+              # pylint: enable=line-too-long
           ]
       )
 
@@ -1107,12 +1111,15 @@ class SandboxActivityTests(unittest.TestCase):
         self.assertEqual(sb.status, sb.Status.OFFLINE)
 
   def test_housekeep_error(self):
+    event_handler = TestingEventHandler(log_housekeep=False)
     env = TestingEnvironment(
         features={'test_feature': TestingFeature(housekeep_interval=0)},
         pool_size=1,
+        housekeep_interval=1.0,
         outage_grace_period=0,
-        outage_retry_interval=0,
+        outage_retry_interval=0.1,
         sandbox_keepalive_interval=0,
+        event_handlers=[event_handler],
     )
     with env:
       with env.sandbox('session1') as sb:
@@ -1129,6 +1136,31 @@ class SandboxActivityTests(unittest.TestCase):
         ):
           time.sleep(0.01)
         self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
+      env.wait_for_housekeeping()
+    self.assertEqual(
+        event_handler.logs,
+        [
+            '[testing-env/0:0] shell: "test_feature" setup',
+            '[testing-env/0:0/test_feature] feature setup',
+            '[testing-env/0:0] shell: "test_feature" setup session',
+            '[testing-env/0:0] sandbox started',
+            '[testing-env] environment started',
+            "[testing-env/0:0] session 'session1' started",
+            '[testing-env/0:0/session1] shell: "test_feature" teardown session',
+            "[testing-env/0:0] session 'session1' ended with SandboxStateError",
+            '[testing-env/0:0] shell: "test_feature" teardown',
+            '[testing-env/0:0/test_feature] feature teardown',
+            '[testing-env/0:0] sandbox shutdown',
+            '[testing-env/0:1] shell: "test_feature" setup',
+            '[testing-env/0:1/test_feature] feature setup',
+            '[testing-env/0:1] shell: "test_feature" setup session',
+            '[testing-env/0:1] sandbox started',
+            '[testing-env/0:1] shell: "test_feature" teardown',
+            '[testing-env/0:1/test_feature] feature teardown',
+            '[testing-env/0:1] sandbox shutdown',
+            '[testing-env] environment shutdown'
+        ]
+    )
 
   def test_remove_event_handler(self):
     env = TestingEnvironment(

langfun/env/event_handlers/__init__.py

@@ -8,5 +8,7 @@ from langfun.env.event_handlers.base import EventHandler
 from langfun.env.event_handlers.event_logger import EventLogger
 from langfun.env.event_handlers.event_logger import ConsoleEventLogger
 
+from langfun.env.event_handlers.metric_writer import MetricWriter
+
 # pylint: enable=g-importing-member
 # pylint: enable=g-bad-import-order

langfun/env/event_handlers/base.py

@@ -48,6 +48,7 @@ class _SessionEventHandler:
       environment: Environment,
       sandbox: Sandbox,
       session_id: str,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -57,6 +58,7 @@ class _SessionEventHandler:
       environment: The environment.
       sandbox: The sandbox.
       session_id: The session ID.
+      duration: The time spent on ending the session.
       lifetime: The session lifetime in seconds.
       error: The error that caused the session to end. If None, the session
         ended normally.
@@ -162,6 +164,7 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
       self,
       environment: Environment,
       sandbox: Sandbox,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -170,6 +173,7 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
     Args:
       environment: The environment.
       sandbox: The sandbox.
+      duration: The time spent on shutting down the sandbox.
       lifetime: The sandbox lifetime in seconds.
       error: The error that caused the sandbox to shutdown. If None, the
         sandbox shutdown normally.
@@ -223,6 +227,13 @@ class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
 class EventHandler(_SandboxEventHandler):
   """Base class for event handlers of an environment."""
 
+  def on_environment_starting(self, environment: Environment) -> None:
+    """Called when the environment is getting started.
+
+    Args:
+      environment: The environment.
+    """
+
   def on_environment_start(
       self,
       environment: Environment,
@@ -258,6 +269,7 @@ class EventHandler(_SandboxEventHandler):
   def on_environment_shutdown(
       self,
       environment: Environment,
+      duration: float,
       lifetime: float,
       error: BaseException | None
   ) -> None:
@@ -265,6 +277,7 @@ class EventHandler(_SandboxEventHandler):
 
     Args:
       environment: The environment.
+      duration: The environment shutdown duration in seconds.
       lifetime: The environment lifetime in seconds.
       error: The error that caused the environment to shutdown. If None, the
         environment shutdown normally.