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

langfun/core/agentic/action.py

@@ -538,6 +538,18 @@ class ExecutionTrace(pg.Object, pg.views.html.HtmlTreeView.Extension):
           remove_class=['not-started'],
       )
 
+  def remove(self, item: TracedItem) -> None:
+    """Removes an item from the sequence."""
+    index = self.items.index(item)
+    if index == -1:
+      raise ValueError(f'Item not found in execution trace: {item!r}')
+
+    with pg.notify_on_change(False):
+      self.items.pop(index)
+
+    if self._tab_control is not None:
+      self._tab_control.remove(index)
+
   def extend(self, items: Iterable[TracedItem]) -> None:
     """Extends the sequence with a list of items."""
     for item in items:
@@ -1648,13 +1660,20 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
   @contextlib.contextmanager
   def track_queries(
       self,
-      phase: str | None = None
+      phase: str | None = None,
+      track_if: Callable[
+          [lf_structured.QueryInvocation],
+          bool
+      ] | None = None,
   ) -> Iterator[list[lf_structured.QueryInvocation]]:
     """Tracks `lf.query` made within the context.
 
     Args:
       phase: The name of a new phase to track the queries in. If not provided,
         the queries will be tracked in the parent phase.
+      track_if: A function that takes a `lf_structured.QueryInvocation` and
+        returns True if the query should be included in the result. If None,
+        all queries (including failed queries) will be included.
 
     Yields:
       A list of `lf.QueryInvocation` objects, each for a single `lf.query`
@@ -1673,6 +1692,11 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       self.event_handler.on_query_start(self, self._current_action, invocation)
 
     def _query_end(invocation: lf_structured.QueryInvocation):
+      if track_if is not None and not track_if(invocation):
+        self._current_execution.remove(invocation)
+      # Even if the query is not included in the execution trace, we still
+      # count the usage summary to the current execution and trigger the
+      # event handler to log the query.
       self._current_execution.merge_usage_summary(invocation.usage_summary)
       self.event_handler.on_query_end(self, self._current_action, invocation)
 
@@ -1705,8 +1729,9 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       *,
       lm: lf.LanguageModel,
       examples: list[lf_structured.MappingExample] | None = None,
+      track_if: Callable[[lf_structured.QueryInvocation], bool] | None = None,
       **kwargs
-      ) -> Any:
+  ) -> Any:
     """Calls `lf.query` and associates it with the current invocation.
 
     The following code are equivalent:
@@ -1731,12 +1756,15 @@ class Session(pg.Object, pg.views.html.HtmlTreeView.Extension):
       default: The default value to return if the query fails.
       lm: The language model to use for the query.
       examples: The examples to use for the query.
+      track_if: A function that takes a `lf_structured.QueryInvocation`
+        and returns True if the query should be tracked.
+        If None, all queries (including failed queries) will be tracked.
       **kwargs: Additional keyword arguments to pass to `lf.query`.
 
     Returns:
       The result of the query.
     """
-    with self.track_queries():
+    with self.track_queries(track_if=track_if):
       return lf_structured.query(
           prompt,
           schema=schema,

langfun/core/agentic/action_test.py

@@ -530,6 +530,31 @@ class SessionTest(unittest.TestCase):
     self.assertIn('agent@', session.id)
     self.assertIsInstance(session.as_message(), lf.AIMessage)
 
+  def test_query_with_track_if(self):
+    lm = fake.StaticResponse('lm response')
+    session = action_lib.Session()
+
+    # Render session to trigger javascript updates to the HTML when
+    # operating on the session.
+    _ = session.to_html()
+    with session:
+      # This query will succeed.
+      session.query(
+          'prompt1',
+          schema=None,
+          lm=lm,
+          track_if=lambda q: not q.has_error,
+          default=None)
+      # This query will fail during parsing.
+      session.query(
+          'prompt2',
+          schema=int,
+          lm=lm,
+          track_if=lambda q: not q.has_error,
+          default=None)
+    self.assertEqual(len(session.root.queries), 1)
+    self.assertIsNone(session.root.queries[0].error)
+
 
 if __name__ == '__main__':
   unittest.main()

langfun/core/llms/gemini.py

@@ -830,9 +830,9 @@ class Gemini(rest.REST):
     )
 
   def _error(self, status_code: int, content: str) -> lf.LMError:
-    if (
-        status_code == 400
-        and b'exceeds the maximum number of tokens' in content
+    if status_code == 400 and (
+        b'exceeds the maximum number of tokens' in content
+        or b'Reduce the input token count and try again.' in content
     ):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)

langfun/env/__init__.py

@@ -24,6 +24,10 @@ from langfun.env.interface import Environment
 from langfun.env.interface import Sandbox
 from langfun.env.interface import Feature
 
+# Decorators for sandbox/feature methods.
+from langfun.env.interface import treat_as_sandbox_state_error
+from langfun.env.interface import log_sandbox_activity
+
 from langfun.env.base_environment import BaseEnvironment
 from langfun.env.base_sandbox import BaseSandbox
 from langfun.env.base_feature import BaseFeature

langfun/env/base_sandbox.py

@@ -26,7 +26,7 @@ import contextlib
 import functools
 import threading
 import time
-from typing import Annotated, Any, Callable, Iterator, Sequence, Type
+from typing import Annotated, Any, Iterator
 
 from langfun.env import interface
 from langfun.env.event_handlers import base as event_handler_base
@@ -100,10 +100,9 @@ class BaseSandbox(interface.Sandbox):
     self._status = status
     self._status_start_time = time.time()
 
-  def report_maybe_state_error(self, e: BaseException | None) -> None:
+  def report_state_error(self, e: interface.SandboxStateError) -> None:
     """Reports sandbox state errors."""
-    if (isinstance(e, interface.SandboxStateError)
-        and e not in self._state_errors):
+    if e not in self._state_errors:
       self._state_errors.append(e)
 
   def _setup_features(self) -> None:
@@ -141,7 +140,8 @@ class BaseSandbox(interface.Sandbox):
         try:
           feature.teardown()
         except BaseException as e:  # pylint: disable=broad-except
-          self.report_maybe_state_error(e)
+          if isinstance(e, interface.SandboxStateError):
+            self.report_state_error(e)
           errors[feature.name] = e
     if errors:
       return interface.FeatureTeardownError(sandbox=self, errors=errors)
@@ -175,7 +175,8 @@ class BaseSandbox(interface.Sandbox):
         try:
           feature.teardown_session()
         except BaseException as e:  # pylint: disable=broad-except
-          self.report_maybe_state_error(e)
+          if isinstance(e, interface.SandboxStateError):
+            self.report_state_error(e)
           feature_teardown_errors[name] = e
 
     return interface.SessionTeardownError(
@@ -277,28 +278,6 @@ class BaseSandbox(interface.Sandbox):
     """Returns the features in the sandbox."""
     return self._features
 
-  def _enter_service_call(self) -> bool:
-    """Enters a service call.
-
-    Returns:
-      True if the service call is at the top of the call stack.
-    """
-    v = getattr(self._tls_state, 'service_call_depth', None)
-    if v is None:
-      v = 0
-    setattr(self._tls_state, 'service_call_depth', v + 1)
-    return v == 0
-
-  def _exit_service_call(self) -> bool:
-    """Exits a service call.
-
-    Returns:
-      True if the service call is at the top of the call stack.
-    """
-    v = getattr(self._tls_state, 'service_call_depth')
-    setattr(self._tls_state, 'service_call_depth', v - 1)
-    return v == 1
-
   #
   # Sandbox start/shutdown.
   #
@@ -369,7 +348,8 @@ class BaseSandbox(interface.Sandbox):
           '[%s]: Sandbox failed to start in %.2f seconds: %s',
           self.id, duration, e
       )
-      self.report_maybe_state_error(e)
+      if isinstance(e, interface.SandboxStateError):
+        self.report_state_error(e)
       self.on_start(duration, e)
       self.shutdown()
       raise e
@@ -442,7 +422,8 @@ class BaseSandbox(interface.Sandbox):
       shutdown_error = None
     except BaseException as e:  # pylint: disable=broad-except
       shutdown_error = e
-      self.report_maybe_state_error(e)
+      if isinstance(e, interface.SandboxStateError):
+        self.report_state_error(e)
       self._set_status(interface.Sandbox.Status.OFFLINE)
       pg.logging.error(
           '[%s]: Sandbox shutdown with error: %s',
@@ -540,7 +521,8 @@ class BaseSandbox(interface.Sandbox):
       self._set_status(self.Status.IN_SESSION)
       self.on_session_start(session_id, time.time() - self._session_start_time)
     except BaseException as e:  # pylint: disable=broad-except
-      self.report_maybe_state_error(e)
+      if isinstance(e, interface.SandboxStateError):
+        self.report_state_error(e)
       self.on_session_start(
           session_id, time.time() - self._session_start_time, e
       )
@@ -626,7 +608,8 @@ class BaseSandbox(interface.Sandbox):
                 self.id,
                 e
             )
-            self.report_maybe_state_error(e)
+            if isinstance(e, interface.SandboxStateError):
+              self.report_state_error(e)
             self.shutdown()
 
         # End session before setting up the next session.
@@ -677,6 +660,30 @@ class BaseSandbox(interface.Sandbox):
         and end_session_error.has_non_sandbox_state_error):
       raise end_session_error  # pylint: disable=raising-bad-type
 
+  @contextlib.contextmanager
+  def track_activity(
+      self,
+      name: str,
+      feature: interface.Feature | None = None,
+      **kwargs: Any
+  ) -> Iterator[None]:
+    """Tracks an activity for the sandbox."""
+    start_time = time.time()
+    error = None
+    try:
+      yield None
+    except BaseException as e:  # pylint: disable=broad-except
+      error = e
+      raise
+    finally:
+      self.on_activity(
+          name=name,
+          feature=feature,
+          duration=time.time() - start_time,
+          error=error,
+          **kwargs
+      )
+
   #
   # Housekeeping.
   #
@@ -722,7 +729,7 @@ class BaseSandbox(interface.Sandbox):
                 str(e)
             )
             self._housekeep_counter += 1
-            self.report_maybe_state_error(e)
+            self.report_state_error(e)
             self.on_housekeep(time.time() - housekeep_start, e)
             self.shutdown()
             break
@@ -744,7 +751,7 @@ class BaseSandbox(interface.Sandbox):
                 feature.name,
                 e,
             )
-            self.report_maybe_state_error(e)
+            self.report_state_error(e)
             self._housekeep_counter += 1
             self.on_housekeep(time.time() - housekeep_start, e)
             self.shutdown()
@@ -917,208 +924,3 @@ class BaseSandbox(interface.Sandbox):
       handler.on_session_end(
           self.environment, self, session_id, duration, lifetime, error
       )
-
-
-#
-# Sandbox service decorator.
-#
-
-
-def sandbox_service(
-    critical_errors: Sequence[
-        Type[BaseException] | tuple[Type[BaseException], str]
-    ] | None = None
-) -> Callable[..., Any]:
-  """Decorator for Sandbox/Feature methods exposed as sandbox services.
-
-  This decorator will catch errors and map to `SandboxStateError` if the
-  error matches any of the critical errors. Consequently, the sandbox will be
-  shutdown automatically when the error is raised.
-
-  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
-      the sandbox as in a bad state, which will trigger a shutdown.
-
-  Returns:
-    The decorator function.
-  """
-  critical_errors = critical_errors or []
-
-  def decorator(func):
-    signature = pg.typing.get_signature(func)
-    if 'session_id' in signature.arg_names:
-      raise ValueError(
-          '`session_id` should not be used as argument for sandbox '
-          '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, (BaseSandbox, interface.Feature)), self
-      sandbox = self.sandbox if isinstance(self, interface.Feature) else self
-
-      # We count the service call stack depth so we could shutdown the sandbox
-      # at the top upon sandbox state error.
-      sandbox._enter_service_call()  # pylint: disable=protected-access
-
-      # When a capability is directly accessed from the environment,
-      # we create a new session for the capability call. This
-      # prevents the sandbox from being reused for other feature calls.
-      if sandbox.status == interface.Sandbox.Status.ACQUIRED:
-        new_session = True
-        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
-
-      kwargs.pop('session_id', None)
-      result = None
-      error = None
-      start_time = time.time()
-
-      try:
-        # Execute the service function.
-        result = func(self, *args, **kwargs)
-
-        # If the result is a context manager, wrap it with a context manager
-        # to end the session when exiting.
-        if isinstance(result, contextlib.AbstractContextManager):
-          return _service_context_manager_wrapper(
-              service=result,
-              sandbox_or_feature=self,
-              sandbox=sandbox,
-              name=func.__name__,
-              kwargs=to_kwargs(*args, **kwargs),
-              start_time=start_time,
-              new_session=new_session
-          )
-
-        # Otherwise, return the result and end the session in the finally block.
-        return result
-      except BaseException as e:
-        error = e
-        sandbox.report_maybe_state_error(e)
-        if pg.match_error(e, critical_errors):
-          state_error = interface.SandboxStateError(
-              'Sandbox encountered an unexpected error executing '
-              f'`{func.__name__}` (args={args!r}, kwargs={kwargs!r}): {e}',
-              sandbox=self
-          )
-          sandbox.report_maybe_state_error(state_error)
-          raise state_error from e
-        raise
-      finally:
-        is_topmost_call = sandbox._exit_service_call()  # pylint: disable=protected-access
-        if not isinstance(result, contextlib.AbstractContextManager):
-          self.on_activity(
-              name=func.__name__,
-              duration=time.time() - start_time,
-              error=error,
-              **to_kwargs(*args, **kwargs),
-          )
-          if new_session:
-            assert is_topmost_call
-
-            # End the session if it's from a feature method and the result
-            # is not a context manager.
-            sandbox.end_session()
-
-        # Shutdown the sandbox if it is at the top of the service call stack and
-        # has state errors.
-        if (is_topmost_call
-            and sandbox.state_errors
-            # Sandbox service method might be called during shutting down, in
-            # that case we don't want to shutdown the sandbox again.
-            and not sandbox.is_shutting_down):
-          sandbox.shutdown()
-
-    return method_wrapper
-  return decorator
-
-
-@contextlib.contextmanager
-def _service_context_manager_wrapper(
-    service: contextlib.AbstractContextManager[Any],
-    sandbox_or_feature: BaseSandbox | interface.Feature,
-    sandbox: interface.Sandbox,
-    name: str,
-    kwargs: dict[str, Any],
-    new_session: bool,
-    start_time: float,
-) -> Iterator[Any]:
-  """Context manager wrapper for ending a sandbox session when exiting."""
-  error = None
-  sandbox._enter_service_call()  # pylint: disable=protected-access
-
-  try:
-    with service as result:
-      yield result
-  except BaseException as e:
-    error = e
-    sandbox.report_maybe_state_error(error)
-    raise
-  finally:
-    sandbox_or_feature.on_activity(
-        name=name,
-        error=error,
-        duration=time.time() - start_time,
-        **kwargs,
-    )
-    is_topmost_call = sandbox._exit_service_call()  # pylint: disable=protected-access
-
-    if new_session:
-      assert is_topmost_call
-      sandbox.end_session()
-    elif isinstance(error, interface.SandboxStateError):
-      sandbox.shutdown()

langfun/env/base_test.py

@@ -15,7 +15,6 @@ import time
 from typing import Any
 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
@@ -460,10 +459,11 @@ class EnvironmentTests(unittest.TestCase):
           simulate_setup_error=interface.SandboxStateError,
           skip_notification=True
       )
-      with env.sandbox() as sb:
-        with self.assertRaises(interface.SandboxStateError):
+      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:
@@ -1250,12 +1250,11 @@ class SandboxStatusTests(unittest.TestCase):
         },
     )
     with env:
-      with env.sandbox(session_id='session1') as sb:
-        with self.assertRaises(interface.SandboxStateError):
+      with self.assertRaises(interface.SandboxStateError):
+        with env.sandbox(session_id='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)
-        sb.shell('echo bar')
+      self.assertEqual(len(sb.state_errors), 1)
+      self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
       self.assertEqual(
           self.event_handler.logs,
           [
@@ -1282,7 +1281,6 @@ class SandboxStatusTests(unittest.TestCase):
               '[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',
-              '[testing-env/test_image:0:0] shell: echo bar',
               # pylint: enable=line-too-long
           ]
       )
@@ -1306,11 +1304,6 @@ class SandboxActivityTests(unittest.TestCase):
             r'test_feature-session-[0-9a-f]{7}'
         )
 
-    with self.assertRaisesRegex(ValueError, '`session_id` should not be used'):
-      @base_sandbox.sandbox_service()
-      def foo(session_id: str):
-        del session_id
-
   def test_ping_error(self):
     env = TestingEnvironment(
         features={'test_feature': TestingFeature(housekeep_interval=0)},
@@ -1487,10 +1480,11 @@ class SandboxServiceTests(unittest.TestCase):
 
   def test_service_call_from_feature_with_error(self):
     with self.env:
-      with self.env.sandbox(session_id='session1') as sb:
-        with self.assertRaises(interface.SandboxStateError):
+      with self.assertRaises(interface.SandboxStateError):
+        with self.env.sandbox(session_id='session1') as sb:
           sb.test_feature.bad_shell_call()
-        self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
+      self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
+      self.assertEqual(len(sb.state_errors), 1)
 
     self.assertEqual(
         self.event_handler.logs,
@@ -1584,7 +1578,6 @@ class SandboxServiceTests(unittest.TestCase):
             '[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/test_feature] test_feature.my_service: None',
             '[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",
@@ -1598,11 +1591,12 @@ class SandboxServiceTests(unittest.TestCase):
 
   def test_service_context_manager_from_feature_with_error(self):
     with self.env:
-      with self.env.sandbox(session_id='session1') as sb:
-        with self.assertRaises(interface.SandboxStateError):
+      with self.assertRaises(interface.SandboxStateError):
+        with self.env.sandbox(session_id='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)
+      self.assertEqual(sb.status, interface.Sandbox.Status.OFFLINE)
+      self.assertEqual(len(sb.state_errors), 1)
     self.assertEqual(
         self.event_handler.logs,
         [
@@ -1614,7 +1608,6 @@ class SandboxServiceTests(unittest.TestCase):
             '[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/test_feature] test_feature.my_service: 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',
@@ -1645,7 +1638,6 @@ class SandboxServiceTests(unittest.TestCase):
             '[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/test_feature] test_feature.my_service: 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',
@@ -1657,7 +1649,6 @@ class SandboxServiceTests(unittest.TestCase):
             '[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/test_feature] test_feature.my_service: None',
             '[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',
@@ -1685,7 +1676,6 @@ class SandboxServiceTests(unittest.TestCase):
             '[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/test_feature] test_feature.my_service: None 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',

langfun/env/interface.py

@@ -17,8 +17,9 @@ import abc
 import contextlib
 import dataclasses
 import enum
+import functools
 import os
-from typing import Annotated, Any, ContextManager, ClassVar, Iterator, Optional
+from typing import Annotated, Any, Callable, ContextManager, ClassVar, Iterator, Optional, Sequence, Type
 
 import pyglove as pg
 
@@ -522,10 +523,17 @@ class Sandbox(pg.Object):
     """Returns True if the sandbox is online."""
     return self.status.is_online
 
-  @property
   @abc.abstractmethod
-  def state_errors(self) -> list[SandboxStateError]:
-    """Returns state errors encountered during sandbox's lifecycle."""
+  def report_state_error(self, error: SandboxStateError) -> None:
+    """Reports state error the sandbox.
+
+    If state errors are reported, the sandbox will be forcefully shutdown when
+    `Sandbox.end_session()` is called, even if the sandbox is set to be
+    reusable.
+
+    Args:
+      error: SandboxStateError to report.
+    """
 
   @abc.abstractmethod
   def start(self) -> None:
@@ -656,6 +664,9 @@ class Sandbox(pg.Object):
     `end_session` should always be called for each `start_session` call, even
     when the session fails to start, to ensure proper cleanup.
 
+    When `end_session` is called with state errors reported, the sandbox will be
+    forcefully shutdown even if the sandbox is set to be reusable.
+
     `end_session` may fail with two sources of errors:
 
     1. SandboxStateError: If the sandbox is in a bad state or session teardown
@@ -679,6 +690,24 @@ class Sandbox(pg.Object):
   def session_id(self) -> str | None:
     """Returns the current user session identifier."""
 
+  @abc.abstractmethod
+  def track_activity(
+      self,
+      name: str,
+      feature: Optional['Feature'] = None,
+      **kwargs: Any
+  ) -> ContextManager[None]:
+    """Context manager that tracks a sandbox activity.
+
+    Args:
+      name: The name of the activity.
+      feature: The feature that the activity is associated with.
+      **kwargs: Additional keyword arguments to pass to the activity handler.
+
+    Returns:
+      A context manager that tracks the activity, including duration and error.
+    """
+
   #
   # API related to a user session.
   # A sandbox could be reused across different user sessions.
@@ -719,6 +748,9 @@ class Sandbox(pg.Object):
     self.start_session(session_id)
     try:
       yield self
+    except SandboxStateError as e:
+      self.report_state_error(e)
+      raise
     finally:
       self.end_session()
 
@@ -887,6 +919,12 @@ class Feature(pg.Object):
     assert self.sandbox is not None
     return self.sandbox.session_id
 
+  def track_activity(self, name: str, **kwargs: Any) -> ContextManager[None]:
+    """Context manager that tracks a feature activity."""
+    return self.sandbox.track_activity(
+        f'{self.name}.{name}', feature=self, **kwargs
+    )
+
 
 def _make_path_compatible(id_str: str) -> str:
   """Makes a path compatible with CNS."""
@@ -900,3 +938,73 @@ def _make_path_compatible(id_str: str) -> str:
           '>': '',
       })
   )
+
+
+def treat_as_sandbox_state_error(
+    errors: Sequence[
+        Type[BaseException] | tuple[Type[BaseException], str]
+    ] | None = None
+) -> Callable[..., Any]:
+  """Decorator for Sandbox/Feature methods to convert errors to SandboxStateError.
+
+  Args:
+    errors: A sequence of exception types or tuples of (error_type, msg_regex).
+      when matched, treat the error as SandboxStateError, which will lead to
+      a sandbox shutdown when caught by `Sandbox.new_session()` context manager.
+
+  Returns:
+    The decorator function.
+  """
+
+  def decorator(func):
+    @functools.wraps(func)
+    def method_wrapper(self, *args, **kwargs) -> Any:
+      """Helper function to safely execute logics in the sandbox."""
+
+      assert isinstance(self, (Sandbox, Feature)), self
+      sandbox = self.sandbox if isinstance(self, Feature) else self
+
+      try:
+        # Execute the service function.
+        return func(self, *args, **kwargs)
+      except BaseException as e:
+        if pg.match_error(e, errors):
+          state_error = SandboxStateError(
+              'Sandbox encountered an unexpected error executing '
+              f'`{func.__name__}` (args={args!r}, kwargs={kwargs!r}): {e}',
+              sandbox=sandbox
+          )
+          raise state_error from e
+        raise
+    return method_wrapper
+  return decorator
+
+
+def log_sandbox_activity(name: str | None = None):
+  """Decorator for Sandbox/Feature methods to log sandbox activity."""
+
+  def decorator(func):
+    signature = pg.typing.get_signature(func)
+    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, (Sandbox, Feature)), self
+      with self.track_activity(
+          name or func.__name__,
+          **to_kwargs(*args, **kwargs)
+      ):
+        return func(self, *args, **kwargs)
+    return method_wrapper
+  return decorator

langfun/env/interface_test.py

@@ -11,6 +11,8 @@
 # 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.
+import contextlib
+from typing import Iterator
 import unittest
 from langfun.env import interface
 
@@ -40,5 +42,81 @@ class IdTest(unittest.TestCase):
     self.assertIsNone(sandbox_id.working_dir(root_dir=None))
 
 
+class TestingSandbox(interface.Sandbox):
+
+  id: interface.Sandbox.Id = interface.Sandbox.Id(
+      environment_id=interface.Environment.Id('env'),
+      image_id='test_image',
+      sandbox_id='0:0'
+  )
+  image_id: str = 'test_image'
+  features: dict[str, interface.Feature] = {}
+  status: interface.Sandbox.Status = interface.Sandbox.Status.READY
+  session_id: str | None = None
+
+  def environment(self) -> interface.Environment:
+    pass
+
+  def _on_bound(self) -> None:
+    self.activities = []
+
+  def report_state_error(self, error: interface.SandboxStateError) -> None:
+    pass
+
+  def start(self) -> None:
+    pass
+
+  def shutdown(self) -> None:
+    pass
+
+  def start_session(self, session_id: str) -> None:
+    pass
+
+  def end_session(self, shutdown_sandbox: bool = False) -> None:
+    pass
+
+  @contextlib.contextmanager
+  def track_activity(
+      self,
+      name: str,
+      feature: interface.Feature | None = None,
+      **kwargs
+  ) -> Iterator[None]:
+    error = None
+    try:
+      yield
+    except BaseException as e:
+      error = e
+      raise
+    finally:
+      self.activities.append((name, error, kwargs))
+
+
+class DecoratorTest(unittest.TestCase):
+
+  def test_treat_as_sandbox_state_error(self):
+
+    class SandboxA(TestingSandbox):
+
+      @interface.treat_as_sandbox_state_error(errors=(ValueError,))
+      def foo(self, bar: str) -> None:
+        raise ValueError(bar)
+
+    with self.assertRaises(interface.SandboxStateError):
+      SandboxA().foo('foo')
+
+  def test_log_sandbox_activity(self):
+
+    class SandboxB(TestingSandbox):
+
+      @interface.log_sandbox_activity()
+      def bar(self, x: str) -> None:
+        pass
+
+    sb = SandboxB()
+    sb.bar('foo')
+    self.assertEqual(sb.activities, [('bar', None, {'x': 'foo'})])
+
+
 if __name__ == '__main__':
   unittest.main()

langfun/env/load_balancers_test.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import concurrent.futures
+import contextlib
 import time
+from typing import Any, Iterator
 import unittest
 
 from langfun.env import interface
@@ -48,6 +50,9 @@ class TestingSandbox(interface.Sandbox):
   def state_errors(self) -> list[interface.SandboxStateError]:
     return []
 
+  def report_state_error(self, error: interface.SandboxStateError) -> None:
+    pass
+
   def set_status(self, status: interface.Sandbox.Status) -> None:
     self.rebind(status=status, skip_notification=True)
 
@@ -70,6 +75,18 @@ class TestingSandbox(interface.Sandbox):
   def session_id(self) -> str | None:
     return self._session_id
 
+  @contextlib.contextmanager
+  def track_activity(
+      self,
+      name: str,
+      feature: interface.Feature | None = None,
+      **kwargs: Any
+  ) -> Iterator[None]:
+    try:
+      yield
+    finally:
+      pass
+
 
 class RoundRobinTest(unittest.TestCase):
 

langfun/env/test_utils.py

@@ -112,7 +112,8 @@ class TestingSandbox(base_sandbox.BaseSandbox):
       self._raise_error('Sandbox shutdown error', self.simulate_shutdown_error)
     super()._shutdown()
 
-  @base_sandbox.sandbox_service(critical_errors=(RuntimeError,))
+  @interface.treat_as_sandbox_state_error(errors=(RuntimeError,))
+  @interface.log_sandbox_activity()
   def shell(
       self,
       code: str,
@@ -184,19 +185,19 @@ class TestingFeature(base_feature.BaseFeature):
     if self.call_end_session_on_teardown_session:
       self.sandbox.end_session()
 
-  @base_sandbox.sandbox_service()
+  @interface.log_sandbox_activity()
   def num_shell_calls(self) -> int:
     return len(self.sandbox._shell_history)  # pylint: disable=protected-access
 
-  @base_sandbox.sandbox_service()
+  @interface.log_sandbox_activity()
   def bad_shell_call(self) -> None:
     self.sandbox.shell('bad command', raise_error=RuntimeError)
 
-  @base_sandbox.sandbox_service()
+  @interface.log_sandbox_activity()
   def show_session_id(self):
     return self.session_id
 
-  @base_sandbox.sandbox_service()
+  @interface.log_sandbox_activity()
   def call_with_varargs(self, code: str, *args, **kwargs):
     del code, args, kwargs
     return 0
@@ -205,7 +206,6 @@ class TestingFeature(base_feature.BaseFeature):
     super()._on_bound()
     self._service = None
 
-  @base_sandbox.sandbox_service()
   @contextlib.contextmanager
   def my_service(self) -> Iterator[Service]:
     try:

{langfun-0.1.2.dev202510290805.dist-info → langfun-0.1.2.dev202510310805.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfun
-Version: 0.1.2.dev202510290805
+Version: 0.1.2.dev202510310805
 Summary: Langfun: Language as Functions.
 Home-page: https://github.com/google/langfun
 Author: Langfun Authors

{langfun-0.1.2.dev202510290805.dist-info → langfun-0.1.2.dev202510310805.dist-info}/RECORD

@@ -35,10 +35,10 @@ langfun/core/subscription_test.py,sha256=Y4ZdbZEwm83YNZBxHff0QR4QUa4rdaNXA3_jfIc
 langfun/core/template.py,sha256=CQOEA1Lq0gU_uk43K1FJjpBSwL8o5fglFr2tGCbOxpI,26008
 langfun/core/template_test.py,sha256=vQZvFVS4VHk-6GUdOEQ-UA_4tlVBbpRPqB1Bw7DFJJg,19052
 langfun/core/agentic/__init__.py,sha256=vsnuvjaz9-nysBjdihGf43JC8AyLPhPJwIOevyONyAQ,1517
-langfun/core/agentic/action.py,sha256=ojwaPIV_a_khKPR6x1Fk5i2dsUTSe3VjKaxnZ92b0nE,58243
+langfun/core/agentic/action.py,sha256=Zz-17gegJxzCZ9hcOhv-atJX1EfIN-vHonS6z99Bllc,59523
 langfun/core/agentic/action_eval.py,sha256=YTilyUEkJl_8FVMgdfO17PurWWaEJ6oA15CuefJJRLk,4887
 langfun/core/agentic/action_eval_test.py,sha256=7AkOwNbUX-ZgR1R0a7bvUZ5abNTUV7blf_8Mnrwb-II,2811
-langfun/core/agentic/action_test.py,sha256=a2D4FOuob7MviuPZR2Wy6xNKnjlTLxhK8HUy8WIyt08,19076
+langfun/core/agentic/action_test.py,sha256=4X11bNffNZG8WfjB1Z-Sm33Z2hs0YPzz-6ct77V8M7Q,19824
 langfun/core/coding/__init__.py,sha256=5utju_fwEsImaiftx4oXKl9FAM8p281k8-Esdh_-m1w,835
 langfun/core/coding/python/__init__.py,sha256=yTXm92oLpQb4A-fZ2qy-bzfhPYND7B-oidtbv1PNaX0,1678
 langfun/core/coding/python/correction.py,sha256=4PD76Xfv36Xrm8Ji3-GgGDNImtcDqWfMw3z6ircJMlM,7285
@@ -101,7 +101,7 @@ langfun/core/llms/deepseek.py,sha256=Pyv2Pmviu91HfNGR994Mh7AKNvWHAAlue7Xfb9MZaPo
 langfun/core/llms/deepseek_test.py,sha256=DvROWPlDuow5E1lfoSkhyGt_ELA19JoQoDsTnRgDtTg,1847
 langfun/core/llms/fake.py,sha256=bDk_4u7V2LmYUotyOaicwzi0-lnWOIIBbR3-Bil1P3o,3481
 langfun/core/llms/fake_test.py,sha256=lC-C2TpEsnf2kmZpa3OiH2H944I4hMWTAaHEXzRj1DU,7855
-langfun/core/llms/gemini.py,sha256=-DL5PebzaTjz7rTFw_1RC5O1aE4EYSv3oNsM65YKCoo,30143
+langfun/core/llms/gemini.py,sha256=Bsd9UmI-Z6j_ZJposExuh7ChCt1ZF2QE1nWPWHaE39k,30204
 langfun/core/llms/gemini_test.py,sha256=y1s0W65SrdepbZxzgIeoTB2MI7sXnfBDf1NsGn57LbM,7617
 langfun/core/llms/google_genai.py,sha256=ogyoOUK4s1OcSFKun0YK5xBRDVyxmvz9WsYNKAwuB0g,5918
 langfun/core/llms/google_genai_test.py,sha256=NKNtpebArQ9ZR7Qsnhd2prFIpMjleojy6o6VMXkJ1zY,1502
@@ -174,24 +174,24 @@ langfun/core/templates/demonstration.py,sha256=vCrgYubdZM5Umqcgp8NUVGXgr4P_c-fik
 langfun/core/templates/demonstration_test.py,sha256=SafcDQ0WgI7pw05EmPI2S4v1t3ABKzup8jReCljHeK4,2162
 langfun/core/templates/selfplay.py,sha256=yhgrJbiYwq47TgzThmHrDQTF4nDrTI09CWGhuQPNv-s,2273
 langfun/core/templates/selfplay_test.py,sha256=Ot__1P1M8oJfoTp-M9-PQ6HUXqZKyMwvZ5f7yQ3yfyM,2326
-langfun/env/__init__.py,sha256=m6Y8y16ms9lytvO_r-Br8HmTp2rNDhb3R6JJaH5dEEk,1491
+langfun/env/__init__.py,sha256=VTjLmS_SkxtkBCmr6hBb4iACeLLPmcbWJkK7MEXsATA,1652
 langfun/env/base_environment.py,sha256=8Cpwb4D37zHj5AK4qgmGIPyxIMEDxTIPHXqxTSNkE_M,25030
 langfun/env/base_feature.py,sha256=JDEhL9LkbBHB0c603guchry7cy_zaIReg5vqExyQQgg,6902
-langfun/env/base_sandbox.py,sha256=aBzPWJGszWuwv8BObvBoIVtl8aiOp41YjTEfVbYYDKQ,37425
-langfun/env/base_test.py,sha256=PqFLqT-IwOe3zH4JWeCtUdDjSqG9KQ4RBQebo67AIqc,74983
-langfun/env/interface.py,sha256=_XguM6xjPQuiTUUrs-4_3ejm8rnolO5cfxpuYE6yoX8,28520
-langfun/env/interface_test.py,sha256=1ZeLPsB9QUyervNrg2LdkJeYx-rQYQ1Zq_IzIKJtS5M,1496
+langfun/env/base_sandbox.py,sha256=0dduNozBtPqawvSX-pv6sEweDNTi89lrT_5qKk5FWak,30656
+langfun/env/base_test.py,sha256=6yHtExP3Soa-l0ug73ZSV9ZfS0B4vu-uaS_K_xvsDwk,74261
+langfun/env/interface.py,sha256=2X-gGD41cnejde5DIQ7QxgW3r85oKP4Gp0iWP-ABmLM,32126
+langfun/env/interface_test.py,sha256=d8vdXL1PkrNQGwzfEI2r6esd6SBnTMgc-3GNArsnuj4,3295
 langfun/env/load_balancers.py,sha256=qRhCthqzjZIQBwta8qC1C0s0J-VQAVomJQqI7Nqv-r4,1948
-langfun/env/load_balancers_test.py,sha256=E1ksJ5Xsoj5aJTM4i9EKDiWKDFljs75FpoeJpbtCKL8,3771
-langfun/env/test_utils.py,sha256=5_lCkTvXC1J9unPdgmqn5PxaucU7GCcro0iFtjehJo0,14160
+langfun/env/load_balancers_test.py,sha256=Bg0h-AL7LpWb_aixFXs_FpgQp4dWLvodcsQj-mys6zs,4125
+langfun/env/test_utils.py,sha256=9fqhxxKERiyAte9bbtwqXqZ1ihNGyuOdCiPPY8G0I8w,14171
 langfun/env/event_handlers/__init__.py,sha256=H34n-TbKSgtxqBhE-yAti8fY6weF2_v3yw59M9_zmGM,443
 langfun/env/event_handlers/base.py,sha256=eGdJ6N5em9kX-c9wzm1TdnRP5_5IAltX5JTHILdjzLM,10124
 langfun/env/event_handlers/event_logger.py,sha256=3dbPjBe53dBgntYHlyLlj_77hVecPSXkmKeiGXMxlO0,12699
 langfun/env/event_handlers/event_logger_test.py,sha256=ryupxaEP9D8wdtSsSwZRSZwqFaHCaSD-bFSea_T7QJo,9133
 langfun/env/event_handlers/metric_writer.py,sha256=F_Gk1lpJX5SZ6-Hyrf_-utf4gvSKvMmcov8VkKogZCU,19618
 langfun/env/event_handlers/metric_writer_test.py,sha256=sntUifTPHGixUshIgVBX4Q9vJL-xbeS0Cpd5X5hSQyQ,5955
-langfun-0.1.2.dev202510290805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-langfun-0.1.2.dev202510290805.dist-info/METADATA,sha256=0QBqKpkm1e7B23GJpMUX5XQWfJEWwn_6nl-Qw_m2tos,7522
-langfun-0.1.2.dev202510290805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-langfun-0.1.2.dev202510290805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
-langfun-0.1.2.dev202510290805.dist-info/RECORD,,
+langfun-0.1.2.dev202510310805.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+langfun-0.1.2.dev202510310805.dist-info/METADATA,sha256=IKnFNK0k4M96hXuSz3sVo6QOFf9z36tKYsRaeJNvKqg,7522
+langfun-0.1.2.dev202510310805.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+langfun-0.1.2.dev202510310805.dist-info/top_level.txt,sha256=RhlEkHxs1qtzmmtWSwYoLVJAc1YrbPtxQ52uh8Z9VvY,8
+langfun-0.1.2.dev202510310805.dist-info/RECORD,,