langfun 0.1.2.dev202509250804__py3-none-any.whl → 0.1.2.dev202509270803__py3-none-any.whl

langfun/env/base_test.py

@@ -18,12 +18,13 @@ import unittest
 from langfun.env import base_sandbox
 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
 TestingFeature = test_utils.TestingFeature
-TestingEnvironmentEventHandler = test_utils.TestingEnvironmentEventHandler
+TestingEventHandler = test_utils.TestingEventHandler
 
 
 class EnvironmentTests(unittest.TestCase):
@@ -42,9 +43,8 @@ class EnvironmentTests(unittest.TestCase):
     self.assertEqual(env.min_pool_size, 0)
     self.assertEqual(env.max_pool_size, 0)
     self.assertEqual(env.sandbox_pool, [])
-    self.assertEqual(env.id, interface.EnvironmentId('testing-env'))
+    self.assertEqual(env.id, interface.Environment.Id('testing-env'))
     self.assertEqual(env.outage_grace_period, 1)
-    self.assertEqual(env.stats_report_interval, 60)
     self.assertEqual(env.features['test_feature'].name, 'test_feature')
 
     self.assertIsNone(env.start_time)
@@ -60,12 +60,16 @@ class EnvironmentTests(unittest.TestCase):
 
       with env.sandbox('session1') as sb:
         self.assertEqual(
-            sb.id, interface.SandboxId(environment_id=env.id, sandbox_id='0')
+            sb.id, interface.Sandbox.Id(environment_id=env.id, sandbox_id='0')
         )
         self.assertEqual(sb.session_id, 'session1')
         self.assertEqual(sb.working_dir, '/tmp/testing-env/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/0/test_feature'
+        )
         with self.assertRaises(AttributeError):
           _ = sb.test_feature2
       self.assertFalse(sb.is_online)
@@ -74,14 +78,26 @@ class EnvironmentTests(unittest.TestCase):
       with self.assertRaises(AttributeError):
         _ = env.test_feature2
 
+  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,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with self.assertRaises(interface.EnvironmentOutageError):
       env.acquire()
@@ -92,12 +108,14 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=0,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        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(
@@ -109,8 +127,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=0,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       with self.assertRaises(interface.EnvironmentOutageError):
@@ -122,8 +139,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       sb = env.acquire()
@@ -135,8 +151,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       sb = env.acquire()
@@ -150,17 +165,76 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=(1, 3),
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       self.assertEqual(len(env.sandbox_pool), 1)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  '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': {
+                  '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), 1)
       sb2 = env.acquire()
       self.assertEqual(sb2.status, interface.Sandbox.Status.ACQUIRED)
       self.assertEqual(len(env.sandbox_pool), 2)
+      self.assertEqual(
+          env.stats(),
+          {
+              'sandbox': {
+                  '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': {
+                'created': 0,
+                'setting_up': 0,
+                'ready': 0,
+                'acquired': 0,
+                'in_session': 0,
+                'exiting_session': 0,
+                'shutting_down': 0,
+                'offline': 0,
+            }
+        }
+    )
 
   def test_acquire_with_growing_pool_failure(self):
     env = TestingEnvironment(
@@ -168,8 +242,7 @@ class EnvironmentTests(unittest.TestCase):
         pool_size=(1, 3),
         outage_grace_period=1,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       self.assertEqual(len(env.sandbox_pool), 1)
@@ -184,15 +257,14 @@ class EnvironmentTests(unittest.TestCase):
       with self.assertRaises(interface.EnvironmentOutageError):
         env.acquire()
 
-  def test_maintenance_error(self):
+  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,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
     with env:
       self.assertEqual(len(env.sandbox_pool), 1)
@@ -216,7 +288,7 @@ class SandboxStatusTests(unittest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.event_handler = TestingEnvironmentEventHandler(
+    self.event_handler = TestingEventHandler(
         log_sandbox_status=True,
         log_feature_setup=True,
         log_session_setup=True,
@@ -246,6 +318,7 @@ class SandboxStatusTests(unittest.TestCase):
             'feature2': TestingFeature(),
         },
     )
+    self.assertFalse(env.enable_pooling)
     with env:
       with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
@@ -294,6 +367,7 @@ class SandboxStatusTests(unittest.TestCase):
         pool_size=1,
         proactive_session_setup=True,
     )
+    self.assertTrue(env.enable_pooling)
     with env:
       with env.sandbox('session1') as sb:
         sb.shell('echo "hello"')
@@ -343,7 +417,7 @@ class SandboxStatusTests(unittest.TestCase):
         features={'test_feature': TestingFeature(setup_session_delay=0.5)},
         pool_size=1,
     )
-    event_handler = TestingEnvironmentEventHandler(
+    event_handler = TestingEventHandler(
         log_sandbox_status=True,
         log_feature_setup=True,
         log_session_setup=True,
@@ -919,6 +993,7 @@ class SandboxStatusTests(unittest.TestCase):
         self.assertEqual(len(sb.state_errors), 0)
         sb.shell('echo bar')
         self.assertEqual(sb.status, interface.Sandbox.Status.IN_SESSION)
+      sb.wait_until_not(interface.Sandbox.Status.SETTING_UP)
       self.assertEqual(sb.status, interface.Sandbox.Status.READY)
       self.assertEqual(
           self.event_handler.logs,
@@ -1020,7 +1095,7 @@ class SandboxActivityTests(unittest.TestCase):
     env = TestingEnvironment(
         features={'test_feature': TestingFeature(housekeep_interval=0)},
         pool_size=1,
-        keepalive_interval=0,
+        sandbox_keepalive_interval=0,
     )
     with env:
       with env.sandbox('session1') as sb:
@@ -1037,7 +1112,7 @@ class SandboxActivityTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=0,
         outage_retry_interval=0,
-        keepalive_interval=0,
+        sandbox_keepalive_interval=0,
     )
     with env:
       with env.sandbox('session1') as sb:
@@ -1052,7 +1127,7 @@ class SandboxActivityTests(unittest.TestCase):
         while sb.housekeep_counter == housekeep_count or (
             sb.status == interface.Sandbox.Status.IN_SESSION
         ):
-          time.sleep(0.1)
+          time.sleep(0.01)
         self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
 
   def test_remove_event_handler(self):
@@ -1061,10 +1136,9 @@ class SandboxActivityTests(unittest.TestCase):
         pool_size=1,
         outage_grace_period=0,
         outage_retry_interval=0,
-        keepalive_interval=0,
-        stats_report_interval=1,
+        sandbox_keepalive_interval=0,
     )
-    event_handler = TestingEnvironmentEventHandler()
+    event_handler = TestingEventHandler()
     with env:
       with env.sandbox('session1') as sb:
         sb.add_event_handler(event_handler)
@@ -1082,21 +1156,20 @@ class SandboxServiceTests(unittest.TestCase):
   def setUp(self):
     super().setUp()
     self.maxDiff = None
-    self.event_handler = TestingEnvironmentEventHandler()
+    self.event_handler = TestingEventHandler()
     self.env = TestingEnvironment(
         features={'test_feature': TestingFeature()},
         pool_size=0,
         outage_grace_period=0,
         outage_retry_interval=0,
-        keepalive_interval=0,
+        sandbox_keepalive_interval=0,
         event_handlers=[self.event_handler],
-        stats_report_interval=1,
         random_seed=1,
     )
 
   def test_service_call_activity_log(self):
 
-    class CustomEventHandler(interface.EnvironmentEventHandler):
+    class CustomEventHandler(event_handler_base.EventHandler):
 
       def __init__(self):
         self.calls = []

langfun/env/event_handlers/__init__.py

@@ -0,0 +1,12 @@
+"""Environment event handlers."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=g-bad-import-order
+
+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
+
+# pylint: enable=g-importing-member
+# pylint: enable=g-bad-import-order

langfun/env/event_handlers/base.py

@@ -0,0 +1,271 @@
+
+# Copyright 2025 The Langfun Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base classes for Langfun environment event handlers."""
+
+from langfun.env import interface
+
+Environment = interface.Environment
+Sandbox = interface.Sandbox
+Feature = interface.Feature
+
+
+class _SessionEventHandler:
+  """Base class for session event handlers."""
+
+  def on_session_start(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session starts.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      session_id: The session ID.
+      duration: The time spent on starting the session.
+      error: The error that caused the session to start. If None, the session
+        started normally.
+    """
+
+  def on_session_end(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      session_id: str,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox session ends.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      session_id: The session ID.
+      lifetime: The session lifetime in seconds.
+      error: The error that caused the session to end. If None, the session
+        ended normally.
+    """
+
+
+class _FeatureEventHandler:
+  """Base class for feature event handlers."""
+
+  def on_feature_setup(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      feature: Feature,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is setup."""
+
+  def on_feature_teardown(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      feature: Feature,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is teardown."""
+
+  def on_feature_teardown_session(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      feature: Feature,
+      session_id: str,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a feature is teardown with a session."""
+
+  def on_feature_setup_session(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      feature: Feature,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a feature is setup with a session."""
+
+  def on_feature_housekeep(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      feature: Feature,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox feature is housekeeping."""
+
+
+class _SandboxEventHandler(_FeatureEventHandler, _SessionEventHandler):
+  """Base class for sandbox event handlers."""
+
+  def on_sandbox_start(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox is started.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      duration: The time spent on starting the sandbox.
+      error: The error that caused the sandbox to start. If None, the sandbox
+        started normally.
+    """
+
+  def on_sandbox_status_change(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      old_status: 'Sandbox.Status',
+      new_status: 'Sandbox.Status',
+      span: float,
+  ) -> None:
+    """Called when a sandbox status changes.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      old_status: The old sandbox status.
+      new_status: The new sandbox status.
+      span: Time spent on the old status in seconds.
+    """
+
+  def on_sandbox_shutdown(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox is shutdown.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      lifetime: The sandbox lifetime in seconds.
+      error: The error that caused the sandbox to shutdown. If None, the
+        sandbox shutdown normally.
+    """
+
+  def on_sandbox_activity(
+      self,
+      name: str,
+      environment: Environment,
+      sandbox: Sandbox,
+      feature: Feature | None,
+      session_id: str | None,
+      duration: float,
+      error: BaseException | None,
+      **kwargs
+  ) -> None:
+    """Called when a sandbox activity is performed.
+
+    Args:
+      name: The name of the sandbox activity.
+      environment: The environment.
+      sandbox: The sandbox.
+      feature: The feature that is associated with the sandbox activity.
+      session_id: The session ID.
+      duration: The sandbox activity duration in seconds.
+      error: The error that caused the sandbox activity to perform. If None,
+        the sandbox activity performed normally.
+      **kwargs: The keyword arguments of the sandbox activity.
+    """
+
+  def on_sandbox_housekeep(
+      self,
+      environment: Environment,
+      sandbox: Sandbox,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when a sandbox finishes a round of housekeeping.
+
+    Args:
+      environment: The environment.
+      sandbox: The sandbox.
+      counter: Zero-based counter of the housekeeping round.
+      duration: The sandbox housekeeping duration in seconds.
+      error: The error that caused the sandbox to housekeeping. If None, the
+        sandbox housekeeping normally.
+    """
+
+
+class EventHandler(_SandboxEventHandler):
+  """Base class for event handlers of an environment."""
+
+  def on_environment_start(
+      self,
+      environment: Environment,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment is started.
+
+    Args:
+      environment: The environment.
+      duration: The environment start duration in seconds.
+      error: The error that failed the environment start. If None, the
+        environment started normally.
+    """
+
+  def on_environment_housekeep(
+      self,
+      environment: Environment,
+      counter: int,
+      duration: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment finishes a round of housekeeping.
+
+    Args:
+      environment: The environment.
+      counter: Zero-based counter of the housekeeping round.
+      duration: The environment start duration in seconds.
+      error: The error that failed the housekeeping. If None, the
+        housekeeping succeeded.
+    """
+
+  def on_environment_shutdown(
+      self,
+      environment: Environment,
+      lifetime: float,
+      error: BaseException | None
+  ) -> None:
+    """Called when the environment is shutdown.
+
+    Args:
+      environment: The environment.
+      lifetime: The environment lifetime in seconds.
+      error: The error that caused the environment to shutdown. If None, the
+        environment shutdown normally.
+    """