langfun 0.1.2.dev202509210803__py3-none-any.whl → 0.1.2.dev202509230805__py3-none-any.whl

langfun/env/base_sandbox.py

@@ -27,7 +27,6 @@ import functools
 import threading
 import time
 from typing import Annotated, Any, Callable, Iterator, Sequence, Type
-import uuid
 
 from langfun.env import interface
 import pyglove as pg
@@ -59,6 +58,16 @@ class BaseSandbox(interface.Sandbox):
       'Interval to ping the sandbox for keeping it alive..'
   ] = 60.0
 
+  proactive_session_setup: Annotated[
+      bool,
+      (
+          'If True, the sandbox will perform setup work before a user session '
+          'is started. This is useful for sandboxes that need to perform '
+          'heavy setup work, which could block the user thread for a long '
+          'time. Applicable only when `reusable` is True.'
+      )
+  ] = True
+
   #
   # There is no required methods that subclasses must implement.
   # Subclasses can override the following methods:
@@ -78,37 +87,92 @@ class BaseSandbox(interface.Sandbox):
       interface.SandboxStateError: If the sandbox is in a bad state.
     """
 
+  def _set_status(self, status: interface.Sandbox.Status) -> None:
+    """Sets the status of the sandbox."""
+    assert self._status != status, (self._status, status)
+    self.on_status_change(self._status, status)
+    self._status = status
+
+  def _maybe_report_state_error(self, e: BaseException | None) -> None:
+    """Reports sandbox state errors."""
+    if isinstance(e, interface.SandboxStateError):
+      self._state_errors.append(e)
+
   def _setup_features(self) -> None:
     """Starts the features in the sandbox."""
+    # We keep track of the features that have setup called so we can teardown
+    # the features when the sandbox is shutdown.
+    self._features_with_setup_called.clear()
+
     for feature in self._features.values():
+      self._features_with_setup_called.add(feature.name)
       feature.setup(self)
 
-  def _teardown_features(self) -> None:
-    """Tears down the features in the sandbox."""
+  def _setup_session(self) -> None:
+    """Sets up a new session for the sandbox."""
+    # We keep track of the features that have setup_session called so we can
+    # teardown the session for them when the session ends.
+    self._features_with_setup_session_called.clear()
+
     for feature in self._features.values():
-      feature.teardown()
+      self._features_with_setup_session_called.add(feature.name)
+      feature.setup_session()
 
-  def _start_session(self, session_id: str) -> None:
-    """Starts a user session."""
-    self._session_id = session_id
-    self._session_start_time = time.time()
+  def _teardown_features(self) -> interface.FeatureTeardownError | None:
+    """Tears down the features in the sandbox.
 
+    IMPORTANT: This method shall not raise any exceptions.
+
+    Returns:
+      FeatureTeardownError: If feature teardown failed with errors.
+        Otherwise None.
+    """
+    errors = {}
     for feature in self._features.values():
-      feature.setup_session(session_id)
+      if feature.name in self._features_with_setup_called:
+        try:
+          feature.teardown()
+        except BaseException as e:  # pylint: disable=broad-except
+          self._maybe_report_state_error(e)
+          errors[feature.name] = e
+    if errors:
+      return interface.FeatureTeardownError(sandbox=self, errors=errors)
+    return None
+
+  def _start_session(self) -> None:
+    """Starts a user session.
 
-  def _end_session(self) -> None:
-    try:
-      for feature in self._features.values():
-        feature.teardown_session(self._session_id)
-    finally:
-      pg.logging.info(
-          '[%s]: User session %s ended. (lifetime: %.2f seconds).',
-          self.id,
-          self._session_id,
-          time.time() - self._session_start_time
-      )
-      self._session_id = None
-      self._session_start_time = None
+    Raises:
+      BaseException: If feature setup failed with user-defined errors.
+      SandboxStateError: If sandbox or feature setup fail due sandbox state
+        errors.
+    """
+    # When pre-session setup is enabled, the session setup is performed
+    # before the session is started. Otherwise we setup the session here.
+    if not self._enable_pre_session_setup:
+      self._setup_session()
+
+  def _end_session(self) -> interface.SessionTeardownError | None:
+    """Ends a user session.
+
+    IMPORTANT: This method shall not raise any exceptions.
+
+    Returns:
+      SessionTeardownError: If session teardown failed with errors.
+        Otherwise None.
+    """
+    feature_teardown_errors = {}
+    for name, feature in self._features.items():
+      if name in self._features_with_setup_session_called:
+        try:
+          feature.teardown_session()
+        except BaseException as e:  # pylint: disable=broad-except
+          self._maybe_report_state_error(e)
+          feature_teardown_errors[name] = e
+
+    return interface.SessionTeardownError(
+        sandbox=self, errors=feature_teardown_errors
+    ) if feature_teardown_errors else None
 
   def _ping(self) -> None:
     """Implementation of ping for health checking."""
@@ -120,16 +184,17 @@ class BaseSandbox(interface.Sandbox):
   def _on_bound(self) -> None:
     """Called when the sandbox is bound."""
     super()._on_bound()
+
     self._features = pg.Dict({
         name: pg.clone(feature)
         for name, feature in self.environment.features.items()
     })
-    self._session_id = None
-    self._session_start_time = None
-    self._alive = False
-    self._start_time = None
+    self._event_handlers = []
 
-    self._needs_housekeep = (
+    self._enable_pre_session_setup = (
+        self.reusable and self.proactive_session_setup
+    )
+    self._enables_housekeep = (
         self.keepalive_interval is not None
         or any(
             feature.housekeep_interval is not None
@@ -139,20 +204,48 @@ class BaseSandbox(interface.Sandbox):
     self._housekeep_thread = None
     self._housekeep_count = 0
 
+    # Runtime state.
+    self._status = self.Status.CREATED
+    self._start_time = None
+    self._state_errors = []
+    self._features_with_setup_called = set()
+    self._features_with_setup_session_called = set()
+
+    self._session_id = None
+    self._session_start_time = None
+
   @functools.cached_property
   def working_dir(self) -> str | None:
     """Returns the working directory for the sandbox."""
     return self.id.working_dir(self.environment.root_dir)
 
   @property
-  def is_alive(self) -> bool:
-    """Returns whether the sandbox is alive."""
-    return self._alive
+  def status(self) -> interface.Sandbox.Status:
+    """Returns the state of the sandbox."""
+    return self._status
+
+  def set_acquired(self) -> None:
+    """Marks the sandbox as acquired."""
+    self._set_status(self.Status.ACQUIRED)
+
+  def add_event_handler(
+      self,
+      event_handler: interface.EnvironmentEventHandler | None
+  ) -> None:
+    """Sets the event handler for the sandbox."""
+    self._event_handlers.append(event_handler)
+
+  def remove_event_handler(
+      self,
+      event_handler: interface.EnvironmentEventHandler | None
+  ) -> None:
+    """Removes the event handler for the sandbox."""
+    self._event_handlers.remove(event_handler)
 
   @property
-  def is_busy(self) -> bool:
-    """Returns whether the sandbox is busy."""
-    return self._session_id not in (None, 'pending')
+  def state_errors(self) -> list[interface.SandboxStateError]:
+    """Returns all errors encountered during sandbox lifecycle."""
+    return self._state_errors
 
   @property
   def features(self) -> dict[str, interface.Feature]:
@@ -166,74 +259,172 @@ class BaseSandbox(interface.Sandbox):
   def start(self) -> None:
     """Starts the sandbox.
 
+    State transitions:
+      CREATED -> SETTING_UP -> READY: When all sandbox and feature setup
+        succeeds.
+      CREATED -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When sandbox or feature
+        setup fails.
+
+    `start` and `shutdown` should be called in pairs, even when the sandbox
+    fails to start. This ensures proper cleanup.
+
+    Start may fail with two sources of errors:
+
+    1. SandboxStateError: If sandbox or feature setup fail due to enviroment
+       outage or sandbox state errors.
+    2. BaseException: If feature setup failed with user-defined errors, this
+       could happen when there is bug in the user code or non-environment code
+       failure.
+
+    In both cases, the sandbox will be shutdown automatically, and the error
+    will be added to `errors`. The sandbox is considered dead and will not be
+    further used.
+
     Raises:
-      interface.SandboxStateError: If the sandbox fails to start.
+      SandboxStateError: If the sandbox is in a bad state.
+      BaseException: If feature setup failed with user-defined errors.
     """
-    assert not self._alive, 'Sandbox is already alive.'
+    assert self._status == self.Status.CREATED, (
+        f'Sandbox {self.id} cannot be started because '
+        f'it is in {self._status} status.'
+    )
 
-    def start_impl():
-      t = time.time()
+    t = time.time()
+    self._state = self.Status.SETTING_UP
+
+    try:
+      # Start the sandbox.
       self._start()
+
+      # Setup the features.
       self._setup_features()
 
-      # We mark the sandbox as alive after the setup before the maintenance
-      # thread is started. This is to avoid the maintenance thread from
-      # immediately shutting down the sandbox because it's not alive yet.
-      self._alive = True
-      self._start_time = time.time()
+      # Setup the first session if pre-session setup is enabled.
+      if self._enable_pre_session_setup:
+        self._setup_session()
 
-      if self._needs_housekeep:
+      if self._enables_housekeep:
         self._housekeep_thread = threading.Thread(
             target=self._housekeep_loop, daemon=True
         )
         self._housekeep_thread.start()
 
+      self._start_time = time.time()
+
+      # Mark the sandbox as ready when all setup succeeds.
+      self._set_status(self.Status.READY)
+
+      self.on_start()
       pg.logging.info(
           '[%s]: Sandbox started in %.2f seconds.',
           self.id, time.time() - t
       )
-
-    interface.call_with_event(
-        action=start_impl,
-        event_handler=self.on_start,
-    )
+    except BaseException as e:  # pylint: disable=broad-except
+      pg.logging.error(
+          '[%s]: Sandbox failed to start: %s',
+          self.id, e
+      )
+      self._maybe_report_state_error(e)
+      self.on_start(e)
+      self.shutdown()
+      raise e
 
   def shutdown(self) -> None:
     """Shuts down the sandbox.
 
+    State transitions:
+      SHUTTING_DOWN -> SHUTTING_DOWN: No operation.
+      OFFLINE -> OFFLINE: No operation.
+      SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When sandbox and feature
+        setup fails.
+      IN_SESSION -> SHUTTING_DOWN -> OFFLINE: When user session exits while
+        sandbox is set not to reuse, or session teardown fails.
+      FREE -> SHUTTING_DOWN -> OFFLINE: When sandbox is shutdown when the
+        environment is shutting down, or housekeeping loop shuts down the
+        sandbox due to housekeeping failures.
+
+    Please be aware that `shutdown` will be called whenever an operation on the
+    sandbox encounters a critical error. This means, `shutdown` should not make
+    the assumption that the sandbox is in a healthy state, even `start` could
+    fail. As a result, `shutdown` must allow re-entry and be thread-safe with
+    other sandbox operations.
+
+    Shutdown may fail with two sources of errors:
+
+    1. SandboxStateError: If the sandbox is in a bad state, and feature teardown
+       logic depending on a healthy sandbox may fail. In such case, we do not
+       raise error to the user as the user session is considered completed. The
+       sandbox is abandoned and new user sessions will be served on other
+       sandboxes.
+
+    2. BaseException: The sandbox is in good state, but user code raises error
+       due to bug or non-environment code failure. In such case, errors will be
+       raised to the user so the error could be surfaced and handled properly.
+       The sandbox is treated as shutdown and will not be further used.
+
     Raises:
-      interface.SandboxStateError: If the sandbox is in a bad state.
+      BaseException: If feature teardown failed with user-defined errors.
     """
-    if not self._alive:
+
+    # Allow re-entry.
+    if self._status in (
+        interface.Sandbox.Status.SHUTTING_DOWN,
+        interface.Sandbox.Status.OFFLINE
+    ):
+      return
+
+    # End current session and shutdown the sandbox if the sandbox is in session.
+    if self._status == self.Status.IN_SESSION:
+      self.end_session(shutdown_sandbox=True)
       return
 
-    self._alive = False
+    self._set_status(interface.Sandbox.Status.SHUTTING_DOWN)
     shutdown_start_time = time.time()
-    def shutdown_impl():
-      self._teardown_features()
+
+    if (self._housekeep_thread is not None
+        and threading.current_thread() is not self._housekeep_thread):
+      self._housekeep_thread.join()
+      self._housekeep_thread = None
+
+    teardown_error = self._teardown_features()
+    try:
       self._shutdown()
-      if (self._housekeep_thread is not None
-          and threading.current_thread() is not self._housekeep_thread):
-        self._housekeep_thread.join()
-        self._housekeep_thread = None
+      self._set_status(interface.Sandbox.Status.OFFLINE)
+
       pg.logging.info(
-          '[%s]: Sandbox shutdown in %.2f seconds. (lifetime: %.2f seconds)',
+          '[%s]: Sandbox shutdown in %.2f seconds. '
+          '(lifetime: %.2f seconds, teardown errors: %s)',
           self.id,
           time.time() - shutdown_start_time,
-          time.time() - self._start_time if self._start_time else 0
+          time.time() - self._start_time if self._start_time else 0,
+          teardown_error
+      )
+      self.on_shutdown(teardown_error)
+      shutdown_error = None
+    except BaseException as e:  # pylint: disable=broad-except
+      shutdown_error = e
+      self._maybe_report_state_error(e)
+      self._set_status(interface.Sandbox.Status.OFFLINE)
+      pg.logging.error(
+          '[%s]: Sandbox shutdown with error: %s',
+          self.id, e
       )
+      self.on_shutdown(teardown_error or shutdown_error)
 
-    interface.call_with_event(
-        action=shutdown_impl,
-        event_handler=self.on_shutdown,
-    )
+    # We raise non-state errors to the user following timely order, so the user
+    # code could be surfaced and handled properly.
+    if (teardown_error is not None
+        and teardown_error.has_non_sandbox_state_error):
+      raise teardown_error
+
+    if shutdown_error is not None and not isinstance(
+        shutdown_error, interface.SandboxStateError
+    ):
+      raise shutdown_error
 
   def ping(self) -> None:
     """Pings the sandbox to check if it is alive."""
-    interface.call_with_event(
-        action=self._ping,
-        event_handler=self.on_ping,
-    )
+    self._ping()
 
   #
   # API related to a user session.
@@ -254,49 +445,183 @@ class BaseSandbox(interface.Sandbox):
     """
     return self._session_id
 
-  def set_pending(self) -> None:
-    """Marks the sandbox as pending for new session."""
-    self._session_id = 'pending'
+  def start_session(
+      self,
+      session_id: str,
+  ) -> None:
+    """Begins a user session with the sandbox.
+
+    State transitions:
+      ACQUIRED -> SETTING_UP -> IN_SESSION: When session setup succeeds.
+      ACQUIRED -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When session setup
+        fails.
 
-  @property
-  def is_pending(self) -> bool:
-    """Returns whether the sandbox is pending for new session."""
-    return self._session_id == 'pending'
+    A session is a sequence of stateful interactions with the sandbox.
+    Across different sessions the sandbox are considered stateless.
+    `start_session` and `end_session` should always be called in pairs, even
+    when the session fails to start. `Sandbox.new_session` context manager is
+    the recommended way to use `start_session` and `end_session` in pairs.
 
-  def start_session(self, session_id: str) -> None:
-    """Begins a user session with the sandbox.
+    Starting a session may fail with two sources of errors:
+
+    1. SandboxStateError: If the sandbox is in a bad state or session setup
+       failed.
+
+    2. BaseException: If session setup failed with user-defined errors.
+
+    In both cases, the sandbox will be shutdown automatically and the
+    session will be considered ended. The error will be added to `errors`.
+    Future session will be served on other sandboxes.
 
     Args:
       session_id: The identifier for the user session.
 
     Raises:
-      interface.SandboxError: If the sandbox already has a user session
-        or the session cannot be started.
+      SandboxStateError: If the sandbox is already in a bad state or session
+        setup failed.
+      BaseException: If session setup failed with user-defined errors.
     """
-    assert self._session_id in (None, 'pending'), (
-        'A user session is already active for this sandbox.'
+    assert self._status == self.Status.ACQUIRED, (
+        f'Sandbox {self.id} is not in acquired state (status={self._status}).'
     )
-    interface.call_with_event(
-        action=self._start_session,
-        event_handler=self.on_session_start,
-        action_kwargs={'session_id': session_id},
-        event_handler_kwargs={'session_id': session_id},
+    assert self._session_id is None, (
+        f'A user session {self._session_id} is already active '
+        f'for sandbox {self.id}.'
     )
+    self._set_status(self.Status.SETTING_UP)
+
+    self._session_id = session_id
+    self._session_start_time = time.time()
 
-  def end_session(self) -> None:
-    """Ends the user session with the sandbox."""
-    assert self._session_id not in (None, 'pending'), (
+    try:
+      self._start_session()
+      self._set_status(self.Status.IN_SESSION)
+      self.on_session_start(session_id)
+    except BaseException as e:  # pylint: disable=broad-except
+      self._maybe_report_state_error(e)
+      self.on_session_start(session_id, e)
+      self.shutdown()
+      raise e
+
+  def end_session(self, shutdown_sandbox: bool = False) -> None:
+    """Ends the user session with the sandbox.
+
+    State transitions:
+      IN_SESSION -> READY: When user session exits normally, and sandbox is set
+        to reuse.
+      IN_SESSION -> SHUTTING_DOWN -> OFFLINE: When user session exits while
+        sandbox is set not to reuse, or session teardown fails.
+      IN_SESSION -> SETTING_UP -> READY: When user session exits normally, and
+        sandbox is set to reuse, and proactive session setup is enabled.
+      IN_SESSION -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When user session
+        exits normally, and proactive session setup is enabled but fails.
+      not (IN_SESSION) -> same state: No operation
+
+    `end_session` should always be called for each `start_session` call, even
+    when the session fails to start, to ensure proper cleanup.
+
+    `end_session` may fail with two sources of errors:
+
+    1. SandboxStateError: If the sandbox is in a bad state or session teardown
+       failed.
+
+    2. BaseException: If session teardown failed with user-defined errors.
+
+    In both cases, the sandbox will be shutdown automatically and the
+    session will be considered ended. The error will be added to `errors`.
+    Future session will be served on other sandboxes.
+
+    However, SandboxStateError encountered during `end_session` will NOT be
+    raised to the user as the user session is considered completed.
+
+    Args:
+      shutdown_sandbox: If True, the sandbox will be shutdown after session
+        teardown.
+
+    Raises:
+      BaseException: If session teardown failed with user-defined errors.
+    """
+    if self._status not in (
+        self.Status.IN_SESSION,
+    ):
+      return
+
+    assert self._session_id is not None, (
         'No user session is active for this sandbox'
     )
-    try:
-      interface.call_with_event(
-          action=self._end_session,
-          event_handler=self.on_session_end,
-          event_handler_kwargs={'session_id': self._session_id},
-      )
-    finally:
-      if not self.reusable:
-        self.shutdown()
+    shutdown_sandbox = shutdown_sandbox or not self.reusable
+
+    # Teardown features for the current session.
+    end_session_error = self._end_session()
+    previous_session_id = self._session_id
+    self._session_id = None
+    self._features_with_setup_session_called.clear()
+
+    # If there is no state error, and proactive session setup is enabled,
+    # set up the next session proactively.
+    if not self.state_errors:
+      if not shutdown_sandbox and self._enable_pre_session_setup:
+        def _setup_next_session():
+          try:
+            self._setup_session()
+            self._set_status(interface.Sandbox.Status.READY)
+          except BaseException as e:  # pylint: disable=broad-except
+            pg.logging.error(
+                '[%s]: Shutting down sandbox due to practively setting up '
+                'next session failed: %s',
+                self.id,
+                e
+            )
+            self._maybe_report_state_error(e)
+            self.shutdown()
+
+        # End session before setting up the next session.
+        self.on_session_end(previous_session_id)
+
+        # Mark the sandbox as setting up to prevent it from being acquired by
+        # other threads.
+        self._set_status(interface.Sandbox.Status.SETTING_UP)
+
+        # TODO(daiyip): Consider using a thread pool to perform next session
+        # setup.
+        threading.Thread(target=_setup_next_session).start()
+      else:
+        # End session before reporting sandbox status change.
+        self.on_session_end(previous_session_id)
+
+        # If shutdown is requested, mark the sandbox as acquired to prevent it
+        # from being acquired by other threads.
+        self._set_status(
+            interface.Sandbox.Status.ACQUIRED if shutdown_sandbox else
+            interface.Sandbox.Status.READY
+        )
+
+    # Otherwise, shutdown the sandbox.
+    else:
+      self.on_session_end(previous_session_id, self.state_errors[0])
+      self._set_status(interface.Sandbox.Status.ACQUIRED)
+      shutdown_sandbox = True
+
+    pg.logging.info(
+        '[%s]: User session %s ended. '
+        '(lifetime: %.2f seconds, teardown errors: %s).',
+        self.id,
+        self._session_id,
+        time.time() - self._session_start_time,
+        end_session_error
+    )
+    self._session_start_time = None
+    self._session_event_handler = None
+
+    if shutdown_sandbox:
+      self.shutdown()
+
+    # We only raise errors if teardown error contains non-sandbox-state error,
+    # meaning that the user code may have bug or other non-environment
+    # failures.
+    if (end_session_error is not None
+        and end_session_error.has_non_sandbox_state_error):
+      raise end_session_error  # pylint: disable=raising-bad-type
 
   #
   # Housekeeping.
@@ -308,7 +633,7 @@ class BaseSandbox(interface.Sandbox):
     last_ping = now
     last_housekeep_time = {name: now for name in self._features.keys()}
 
-    while self._alive:
+    while self._status not in (self.Status.SHUTTING_DOWN, self.Status.OFFLINE):
       if self.keepalive_interval is not None:
         if time.time() - last_ping > self.keepalive_interval:
           try:
@@ -321,6 +646,7 @@ class BaseSandbox(interface.Sandbox):
                 str(e)
             )
             self._housekeep_count += 1
+            self._maybe_report_state_error(e)
             self.shutdown()
             break
           last_ping = time.time()
@@ -339,13 +665,141 @@ class BaseSandbox(interface.Sandbox):
                 'Shutting down sandbox...',
                 self.id,
                 feature.name,
-                str(e)
+                e,
             )
+            self._maybe_report_state_error(e)
             self.shutdown()
             break
       self._housekeep_count += 1
       time.sleep(1)
 
+  #
+  # Event handlers subclasses can override.
+  #
+
+  def on_start(self, error: BaseException | None = None) -> None:
+    """Called when the sandbox is started."""
+    for handler in self._event_handlers:
+      handler.on_sandbox_start(self.environment, self, error)
+
+  def on_status_change(
+      self,
+      old_status: interface.Sandbox.Status,
+      new_status: interface.Sandbox.Status,
+  ) -> None:
+    """Called when the sandbox status changes."""
+    for handler in self._event_handlers:
+      handler.on_sandbox_status_change(
+          self.environment, self, old_status, new_status
+      )
+
+  def on_shutdown(self, error: BaseException | None = None) -> None:
+    """Called when the sandbox is shutdown."""
+    for handler in self._event_handlers:
+      handler.on_sandbox_shutdown(self.environment, self, error)
+
+  def on_feature_setup(
+      self,
+      feature: interface.Feature,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when a feature is setup."""
+    for handler in self._event_handlers:
+      handler.on_feature_setup(
+          self.environment, self, feature, error
+      )
+
+  def on_feature_teardown(
+      self,
+      feature: interface.Feature,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when a feature is teardown."""
+    for handler in self._event_handlers:
+      handler.on_feature_teardown(
+          self.environment, self, feature, error
+      )
+
+  def on_feature_setup_session(
+      self,
+      feature: interface.Feature,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when a feature is setup for a user session."""
+    for handler in self._event_handlers:
+      handler.on_feature_setup_session(
+          self.environment, self, feature, self.session_id, error
+      )
+
+  def on_feature_teardown_session(
+      self,
+      feature: interface.Feature,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when a feature is teardown for a user session."""
+    for handler in self._event_handlers:
+      handler.on_feature_teardown_session(
+          self.environment, self, feature, self.session_id, error
+      )
+
+  def on_feature_housekeep(
+      self,
+      feature: interface.Feature,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when a feature is housekeeping."""
+    for handler in self._event_handlers:
+      handler.on_feature_housekeep(
+          self.environment, self, feature, error
+      )
+
+  def on_session_start(
+      self,
+      session_id: str,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when the user session starts."""
+    for handler in self._event_handlers:
+      handler.on_session_start(
+          self.environment, self, session_id, error
+      )
+
+  def on_session_activity(
+      self,
+      session_id: str,
+      name: str,
+      feature: interface.Feature | None = None,
+      error: BaseException | None = None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed."""
+    for handler in self._event_handlers:
+      handler.on_session_activity(
+          session_id=session_id,
+          name=name,
+          environment=self.environment,
+          sandbox=self,
+          feature=feature,
+          error=error,
+          **kwargs
+      )
+
+  def on_session_end(
+      self,
+      session_id: str,
+      error: BaseException | None = None
+  ) -> None:
+    """Called when the user session ends."""
+    for handler in self._event_handlers:
+      handler.on_session_end(
+          self.environment, self, session_id, error
+      )
+
+
+#
+# Sandbox service decorator.
+#
+
 
 def sandbox_service(
     critical_errors: Sequence[
@@ -358,9 +812,38 @@ def sandbox_service(
   error matches any of the critical errors. Consequently, the sandbox will be
   shutdown automatically when the error is raised.
 
-  if the decorated method returns a context manager, a wrapper context manager
+  Example:
+
+  ```
+    with env:
+      with env.sandbox() as sb:
+        try:
+          sb.test_feature.do_something_with_non_state_error()
+        except ValueError:
+          # sandbox will not be shutdown.
+          pass
+
+        try:
+          sb.test_feature.do_something_with_state_error()
+        except ValueError:
+          assert sb.state == sb.Status.OFFLINE
+  ```
+
+  If the decorated method returns a context manager, a wrapper context manager
   will be returned, which will end the session when exiting the context.
 
+  Example:
+
+  ```
+    with env:
+      with env.test_feature.do_something_with_context_manager() as result:
+        # sandbox will be alive during the whole context manager cycle.
+  ```
+
+  For sandbox service methods, an optional `session_id` argument can be passed
+  to create a new session for the service call, even its signature does not
+  contain a `session_id` argument.
+
   Args:
     critical_errors: A sequence of exception types or tuples of exception type
       and error messages (described in regular expression), when matched, treat
@@ -379,31 +862,45 @@ def sandbox_service(
           'service method. Please use `self.session_id` instead.'
       )
 
+    def to_kwargs(*args, **kwargs):
+      num_non_self_args = len(signature.arg_names) - 1
+      if len(args) > num_non_self_args:
+        assert signature.varargs is not None, (signature, args)
+        kwargs[signature.varargs.name] = tuple(args[num_non_self_args:])
+        args = args[:num_non_self_args]
+      for i in range(len(args)):
+        # The first argument is `self`.
+        kwargs[signature.arg_names[i + 1]] = args[i]
+      return kwargs
+
     @functools.wraps(func)
     def method_wrapper(self, *args, **kwargs) -> Any:
       """Helper function to safely execute logics in the sandbox."""
-      assert isinstance(self, (interface.Sandbox, interface.Feature)), self
+      assert isinstance(self, (BaseSandbox, interface.Feature)), self
       sandbox = self.sandbox if isinstance(self, interface.Feature) else self
 
       # When a capability is directly accessed from the environment,
-      # we scope the function call within a short-lived sandbox session. This
+      # we create a new session for the capability call. This
       # prevents the sandbox from being reused for other feature calls.
-      if sandbox.is_pending:
+      if sandbox.status == interface.Sandbox.Status.ACQUIRED:
         new_session = True
-        session_id = kwargs.get('session_id', f'session-{uuid.uuid4().hex[:7]}')
+        new_session_id = kwargs.get('session_id')
+        if new_session_id is None:
+          new_session_id = sandbox.environment.new_session_id()
+
+        # If it's a feature method called from the environment, start a new
+        # session for the feature call.
+        sandbox.start_session(new_session_id)
       else:
         new_session = False
-        session_id = sandbox.session_id
 
       kwargs.pop('session_id', None)
+      session_id = sandbox.session_id
       result = None
+      state_error = None
       error = None
-      try:
-        # If it's a feature method called from the environment, start a new
-        # session for the feature call.
-        if new_session:
-          sandbox.start_session(session_id)
 
+      try:
         # Execute the service function.
         result = func(self, *args, **kwargs)
 
@@ -417,30 +914,39 @@ def sandbox_service(
         # Otherwise, return the result and end the session in the finally block.
         return result
       except interface.SandboxStateError as e:
+        sandbox._maybe_report_state_error(e)  # pylint: disable=protected-access
+        state_error = e
         error = e
         raise
       except BaseException as e:
+        error = e
         if pg.match_error(e, critical_errors):
-          error = e
-          raise interface.SandboxStateError(
+          state_error = interface.SandboxStateError(
               'Sandbox encountered an unexpected error executing '
               f'`{func.__name__}` (args={args!r}, kwargs={kwargs!r}): {e}',
               sandbox=self
-          ) from e
+          )
+          sandbox._maybe_report_state_error(state_error)  # pylint: disable=protected-access
+          raise state_error from e
         raise
       finally:
-        if error is not None:
+        if session_id is not None:
+          self.on_session_activity(
+              name=func.__name__,
+              session_id=session_id,
+              error=error,
+              **to_kwargs(*args, **kwargs),
+          )
+
+        if state_error is not None:
           sandbox.shutdown()
-
-        # End the session if it's from a feature method and the result is not a
-        # context manager.
-        if (new_session
-            and not isinstance(result, contextlib.AbstractContextManager)):
-          sandbox.end_session()
-
-        self.on_session_activity(
-            session_id=session_id, error=error, args=args, **kwargs
-        )
+        elif (new_session
+              and not isinstance(result, contextlib.AbstractContextManager)):
+          # End the session if it's from a feature method and the result is not
+          # a context manager.
+          sandbox.end_session(
+              shutdown_sandbox=isinstance(error, interface.SandboxStateError)
+          )
     return method_wrapper
   return decorator