letta-nightly 0.7.21.dev20250522104246__py3-none-any.whl → 0.7.22.dev20250523104244__py3-none-any.whl

letta/services/per_agent_lock_manager.py

@@ -1,6 +1,8 @@
 import threading
 from collections import defaultdict
 
+from letta.tracing import trace_method
+
 
 class PerAgentLockManager:
     """Manages per-agent locks."""
@@ -8,10 +10,12 @@ class PerAgentLockManager:
     def __init__(self):
         self.locks = defaultdict(threading.Lock)
 
+    @trace_method
     def get_lock(self, agent_id: str) -> threading.Lock:
         """Retrieve the lock for a specific agent_id."""
         return self.locks[agent_id]
 
+    @trace_method
     def clear_lock(self, agent_id: str):
         """Optionally remove a lock if no longer needed (to prevent unbounded growth)."""
         if agent_id in self.locks:

letta/services/provider_manager.py

@@ -6,12 +6,14 @@ from letta.schemas.providers import Provider as PydanticProvider
 from letta.schemas.providers import ProviderCheck, ProviderCreate, ProviderUpdate
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.tracing import trace_method
 from letta.utils import enforce_types
 
 
 class ProviderManager:
 
     @enforce_types
+    @trace_method
     def create_provider(self, request: ProviderCreate, actor: PydanticUser) -> PydanticProvider:
         """Create a new provider if it doesn't already exist."""
         with db_registry.session() as session:
@@ -32,6 +34,7 @@ class ProviderManager:
             return new_provider.to_pydantic()
 
     @enforce_types
+    @trace_method
     def update_provider(self, provider_id: str, provider_update: ProviderUpdate, actor: PydanticUser) -> PydanticProvider:
         """Update provider details."""
         with db_registry.session() as session:
@@ -48,6 +51,7 @@ class ProviderManager:
             return existing_provider.to_pydantic()
 
     @enforce_types
+    @trace_method
     def delete_provider_by_id(self, provider_id: str, actor: PydanticUser):
         """Delete a provider."""
         with db_registry.session() as session:
@@ -62,6 +66,7 @@ class ProviderManager:
             session.commit()
 
     @enforce_types
+    @trace_method
     def list_providers(
         self,
         actor: PydanticUser,
@@ -87,16 +92,45 @@ class ProviderManager:
             return [provider.to_pydantic() for provider in providers]
 
     @enforce_types
+    @trace_method
+    async def list_providers_async(
+        self,
+        actor: PydanticUser,
+        name: Optional[str] = None,
+        provider_type: Optional[ProviderType] = None,
+        after: Optional[str] = None,
+        limit: Optional[int] = 50,
+    ) -> List[PydanticProvider]:
+        """List all providers with optional pagination."""
+        filter_kwargs = {}
+        if name:
+            filter_kwargs["name"] = name
+        if provider_type:
+            filter_kwargs["provider_type"] = provider_type
+        async with db_registry.async_session() as session:
+            providers = await ProviderModel.list_async(
+                db_session=session,
+                after=after,
+                limit=limit,
+                actor=actor,
+                **filter_kwargs,
+            )
+            return [provider.to_pydantic() for provider in providers]
+
+    @enforce_types
+    @trace_method
     def get_provider_id_from_name(self, provider_name: Union[str, None], actor: PydanticUser) -> Optional[str]:
         providers = self.list_providers(name=provider_name, actor=actor)
         return providers[0].id if providers else None
 
     @enforce_types
+    @trace_method
     def get_override_key(self, provider_name: Union[str, None], actor: PydanticUser) -> Optional[str]:
         providers = self.list_providers(name=provider_name, actor=actor)
         return providers[0].api_key if providers else None
 
     @enforce_types
+    @trace_method
     def check_provider_api_key(self, provider_check: ProviderCheck) -> None:
         provider = PydanticProvider(
             name=provider_check.provider_type.value,

letta/services/sandbox_config_manager.py

@@ -12,6 +12,7 @@ from letta.schemas.sandbox_config import SandboxConfig as PydanticSandboxConfig
 from letta.schemas.sandbox_config import SandboxConfigCreate, SandboxConfigUpdate, SandboxType
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.tracing import trace_method
 from letta.utils import enforce_types, printd
 
 logger = get_logger(__name__)
@@ -21,6 +22,7 @@ class SandboxConfigManager:
     """Manager class to handle business logic related to SandboxConfig and SandboxEnvironmentVariable."""
 
     @enforce_types
+    @trace_method
     def get_or_create_default_sandbox_config(self, sandbox_type: SandboxType, actor: PydanticUser) -> PydanticSandboxConfig:
         sandbox_config = self.get_sandbox_config_by_type(sandbox_type, actor=actor)
         if not sandbox_config:
@@ -38,6 +40,7 @@ class SandboxConfigManager:
         return sandbox_config
 
     @enforce_types
+    @trace_method
     def create_or_update_sandbox_config(self, sandbox_config_create: SandboxConfigCreate, actor: PydanticUser) -> PydanticSandboxConfig:
         """Create or update a sandbox configuration based on the PydanticSandboxConfig schema."""
         config = sandbox_config_create.config
@@ -71,6 +74,61 @@ class SandboxConfigManager:
                 return db_sandbox.to_pydantic()
 
     @enforce_types
+    @trace_method
+    async def get_or_create_default_sandbox_config_async(self, sandbox_type: SandboxType, actor: PydanticUser) -> PydanticSandboxConfig:
+        sandbox_config = await self.get_sandbox_config_by_type_async(sandbox_type, actor=actor)
+        if not sandbox_config:
+            logger.debug(f"Creating new sandbox config of type {sandbox_type}, none found for organization {actor.organization_id}.")
+
+            # TODO: Add more sandbox types later
+            if sandbox_type == SandboxType.E2B:
+                default_config = {}  # Empty
+            else:
+                # TODO: May want to move this to environment variables v.s. persisting in database
+                default_local_sandbox_path = LETTA_TOOL_EXECUTION_DIR
+                default_config = LocalSandboxConfig(sandbox_dir=default_local_sandbox_path).model_dump(exclude_none=True)
+
+            sandbox_config = await self.create_or_update_sandbox_config_async(SandboxConfigCreate(config=default_config), actor=actor)
+        return sandbox_config
+
+    @enforce_types
+    @trace_method
+    async def create_or_update_sandbox_config_async(
+        self, sandbox_config_create: SandboxConfigCreate, actor: PydanticUser
+    ) -> PydanticSandboxConfig:
+        """Create or update a sandbox configuration based on the PydanticSandboxConfig schema."""
+        config = sandbox_config_create.config
+        sandbox_type = config.type
+        sandbox_config = PydanticSandboxConfig(
+            type=sandbox_type, config=config.model_dump(exclude_none=True), organization_id=actor.organization_id
+        )
+
+        # Attempt to retrieve the existing sandbox configuration by type within the organization
+        db_sandbox = await self.get_sandbox_config_by_type_async(sandbox_config.type, actor=actor)
+        if db_sandbox:
+            # Prepare the update data, excluding fields that should not be reset
+            update_data = sandbox_config.model_dump(exclude_unset=True, exclude_none=True)
+            update_data = {key: value for key, value in update_data.items() if getattr(db_sandbox, key) != value}
+
+            # If there are changes, update the sandbox configuration
+            if update_data:
+                db_sandbox = await self.update_sandbox_config_async(db_sandbox.id, SandboxConfigUpdate(**update_data), actor)
+            else:
+                printd(
+                    f"`create_or_update_sandbox_config` was called with user_id={actor.id}, organization_id={actor.organization_id}, "
+                    f"type={sandbox_config.type}, but found existing configuration with nothing to update."
+                )
+
+            return db_sandbox
+        else:
+            # If the sandbox configuration doesn't exist, create a new one
+            async with db_registry.async_session() as session:
+                db_sandbox = SandboxConfigModel(**sandbox_config.model_dump(exclude_none=True))
+                await db_sandbox.create_async(session, actor=actor)
+                return db_sandbox.to_pydantic()
+
+    @enforce_types
+    @trace_method
     def update_sandbox_config(
         self, sandbox_config_id: str, sandbox_update: SandboxConfigUpdate, actor: PydanticUser
     ) -> PydanticSandboxConfig:
@@ -98,6 +156,35 @@ class SandboxConfigManager:
             return sandbox.to_pydantic()
 
     @enforce_types
+    @trace_method
+    async def update_sandbox_config_async(
+        self, sandbox_config_id: str, sandbox_update: SandboxConfigUpdate, actor: PydanticUser
+    ) -> PydanticSandboxConfig:
+        """Update an existing sandbox configuration."""
+        async with db_registry.async_session() as session:
+            sandbox = await SandboxConfigModel.read_async(db_session=session, identifier=sandbox_config_id, actor=actor)
+            # We need to check that the sandbox_update provided is the same type as the original sandbox
+            if sandbox.type != sandbox_update.config.type:
+                raise ValueError(
+                    f"Mismatched type for sandbox config update: tried to update sandbox_config of type {sandbox.type} with config of type {sandbox_update.config.type}"
+                )
+
+            update_data = sandbox_update.model_dump(exclude_unset=True, exclude_none=True)
+            update_data = {key: value for key, value in update_data.items() if getattr(sandbox, key) != value}
+
+            if update_data:
+                for key, value in update_data.items():
+                    setattr(sandbox, key, value)
+                await sandbox.update_async(db_session=session, actor=actor)
+            else:
+                printd(
+                    f"`update_sandbox_config` called with user_id={actor.id}, organization_id={actor.organization_id}, "
+                    f"name={sandbox.type}, but nothing to update."
+                )
+            return sandbox.to_pydantic()
+
+    @enforce_types
+    @trace_method
     def delete_sandbox_config(self, sandbox_config_id: str, actor: PydanticUser) -> PydanticSandboxConfig:
         """Delete a sandbox configuration by its ID."""
         with db_registry.session() as session:
@@ -106,6 +193,7 @@ class SandboxConfigManager:
             return sandbox.to_pydantic()
 
     @enforce_types
+    @trace_method
     def list_sandbox_configs(
         self,
         actor: PydanticUser,
@@ -123,6 +211,7 @@ class SandboxConfigManager:
             return [sandbox.to_pydantic() for sandbox in sandboxes]
 
     @enforce_types
+    @trace_method
     async def list_sandbox_configs_async(
         self,
         actor: PydanticUser,
@@ -140,6 +229,7 @@ class SandboxConfigManager:
             return [sandbox.to_pydantic() for sandbox in sandboxes]
 
     @enforce_types
+    @trace_method
     def get_sandbox_config_by_id(self, sandbox_config_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSandboxConfig]:
         """Retrieve a sandbox configuration by its ID."""
         with db_registry.session() as session:
@@ -150,6 +240,7 @@ class SandboxConfigManager:
                 return None
 
     @enforce_types
+    @trace_method
     def get_sandbox_config_by_type(self, type: SandboxType, actor: Optional[PydanticUser] = None) -> Optional[PydanticSandboxConfig]:
         """Retrieve a sandbox config by its type."""
         with db_registry.session() as session:
@@ -167,6 +258,27 @@ class SandboxConfigManager:
                 return None
 
     @enforce_types
+    @trace_method
+    async def get_sandbox_config_by_type_async(
+        self, type: SandboxType, actor: Optional[PydanticUser] = None
+    ) -> Optional[PydanticSandboxConfig]:
+        """Retrieve a sandbox config by its type."""
+        async with db_registry.async_session() as session:
+            try:
+                sandboxes = await SandboxConfigModel.list_async(
+                    db_session=session,
+                    type=type,
+                    organization_id=actor.organization_id,
+                    limit=1,
+                )
+                if sandboxes:
+                    return sandboxes[0].to_pydantic()
+                return None
+            except NoResultFound:
+                return None
+
+    @enforce_types
+    @trace_method
     def create_sandbox_env_var(
         self, env_var_create: SandboxEnvironmentVariableCreate, sandbox_config_id: str, actor: PydanticUser
     ) -> PydanticEnvVar:
@@ -194,6 +306,7 @@ class SandboxConfigManager:
             return env_var.to_pydantic()
 
     @enforce_types
+    @trace_method
     def update_sandbox_env_var(
         self, env_var_id: str, env_var_update: SandboxEnvironmentVariableUpdate, actor: PydanticUser
     ) -> PydanticEnvVar:
@@ -215,6 +328,7 @@ class SandboxConfigManager:
             return env_var.to_pydantic()
 
     @enforce_types
+    @trace_method
     def delete_sandbox_env_var(self, env_var_id: str, actor: PydanticUser) -> PydanticEnvVar:
         """Delete a sandbox environment variable by its ID."""
         with db_registry.session() as session:
@@ -223,6 +337,7 @@ class SandboxConfigManager:
             return env_var.to_pydantic()
 
     @enforce_types
+    @trace_method
     def list_sandbox_env_vars(
         self,
         sandbox_config_id: str,
@@ -242,6 +357,7 @@ class SandboxConfigManager:
             return [env_var.to_pydantic() for env_var in env_vars]
 
     @enforce_types
+    @trace_method
     async def list_sandbox_env_vars_async(
         self,
         sandbox_config_id: str,
@@ -261,6 +377,7 @@ class SandboxConfigManager:
             return [env_var.to_pydantic() for env_var in env_vars]
 
     @enforce_types
+    @trace_method
     def list_sandbox_env_vars_by_key(
         self, key: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
     ) -> List[PydanticEnvVar]:
@@ -276,6 +393,7 @@ class SandboxConfigManager:
             return [env_var.to_pydantic() for env_var in env_vars]
 
     @enforce_types
+    @trace_method
     def get_sandbox_env_vars_as_dict(
         self, sandbox_config_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
     ) -> Dict[str, str]:
@@ -286,6 +404,18 @@ class SandboxConfigManager:
         return result
 
     @enforce_types
+    @trace_method
+    async def get_sandbox_env_vars_as_dict_async(
+        self, sandbox_config_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
+    ) -> Dict[str, str]:
+        env_vars = await self.list_sandbox_env_vars_async(sandbox_config_id, actor, after, limit)
+        result = {}
+        for env_var in env_vars:
+            result[env_var.key] = env_var.value
+        return result
+
+    @enforce_types
+    @trace_method
     def get_sandbox_env_var_by_key_and_sandbox_config_id(
         self, key: str, sandbox_config_id: str, actor: Optional[PydanticUser] = None
     ) -> Optional[PydanticEnvVar]:

letta/services/source_manager.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import List, Optional
 
 from letta.orm.errors import NoResultFound
@@ -9,6 +10,7 @@ from letta.schemas.source import Source as PydanticSource
 from letta.schemas.source import SourceUpdate
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
+from letta.tracing import trace_method
 from letta.utils import enforce_types, printd
 
 
@@ -16,25 +18,27 @@ class SourceManager:
     """Manager class to handle business logic related to Sources."""
 
     @enforce_types
-    def create_source(self, source: PydanticSource, actor: PydanticUser) -> PydanticSource:
+    @trace_method
+    async def create_source(self, source: PydanticSource, actor: PydanticUser) -> PydanticSource:
         """Create a new source based on the PydanticSource schema."""
         # Try getting the source first by id
-        db_source = self.get_source_by_id(source.id, actor=actor)
+        db_source = await self.get_source_by_id(source.id, actor=actor)
         if db_source:
             return db_source
         else:
-            with db_registry.session() as session:
+            async with db_registry.async_session() as session:
                 # Provide default embedding config if not given
                 source.organization_id = actor.organization_id
                 source = SourceModel(**source.model_dump(to_orm=True, exclude_none=True))
-                source.create(session, actor=actor)
+                await source.create_async(session, actor=actor)
             return source.to_pydantic()
 
     @enforce_types
-    def update_source(self, source_id: str, source_update: SourceUpdate, actor: PydanticUser) -> PydanticSource:
+    @trace_method
+    async def update_source(self, source_id: str, source_update: SourceUpdate, actor: PydanticUser) -> PydanticSource:
         """Update a source by its ID with the given SourceUpdate object."""
-        with db_registry.session() as session:
-            source = SourceModel.read(db_session=session, identifier=source_id, actor=actor)
+        async with db_registry.async_session() as session:
+            source = await SourceModel.read_async(db_session=session, identifier=source_id, actor=actor)
 
             # get update dictionary
             update_data = source_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
@@ -53,18 +57,22 @@ class SourceManager:
             return source.to_pydantic()
 
     @enforce_types
-    def delete_source(self, source_id: str, actor: PydanticUser) -> PydanticSource:
+    @trace_method
+    async def delete_source(self, source_id: str, actor: PydanticUser) -> PydanticSource:
         """Delete a source by its ID."""
-        with db_registry.session() as session:
-            source = SourceModel.read(db_session=session, identifier=source_id)
-            source.hard_delete(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            source = await SourceModel.read_async(db_session=session, identifier=source_id)
+            await source.hard_delete_async(db_session=session, actor=actor)
             return source.to_pydantic()
 
     @enforce_types
-    def list_sources(self, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50, **kwargs) -> List[PydanticSource]:
+    @trace_method
+    async def list_sources(
+        self, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50, **kwargs
+    ) -> List[PydanticSource]:
         """List all sources with optional pagination."""
-        with db_registry.session() as session:
-            sources = SourceModel.list(
+        async with db_registry.async_session() as session:
+            sources = await SourceModel.list_async(
                 db_session=session,
                 after=after,
                 limit=limit,
@@ -74,18 +82,17 @@ class SourceManager:
             return [source.to_pydantic() for source in sources]
 
     @enforce_types
-    def size(
-        self,
-        actor: PydanticUser,
-    ) -> int:
+    @trace_method
+    async def size(self, actor: PydanticUser) -> int:
         """
         Get the total count of sources for the given user.
         """
-        with db_registry.session() as session:
-            return SourceModel.size(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            return await SourceModel.size_async(db_session=session, actor=actor)
 
     @enforce_types
-    def list_attached_agents(self, source_id: str, actor: Optional[PydanticUser] = None) -> List[PydanticAgentState]:
+    @trace_method
+    async def list_attached_agents(self, source_id: str, actor: Optional[PydanticUser] = None) -> List[PydanticAgentState]:
         """
         Lists all agents that have the specified source attached.
 
@@ -96,30 +103,33 @@ class SourceManager:
         Returns:
             List[PydanticAgentState]: List of agents that have this source attached
         """
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             # Verify source exists and user has permission to access it
-            source = SourceModel.read(db_session=session, identifier=source_id, actor=actor)
+            source = await SourceModel.read_async(db_session=session, identifier=source_id, actor=actor)
 
             # The agents relationship is already loaded due to lazy="selectin" in the Source model
             # and will be properly filtered by organization_id due to the OrganizationMixin
-            return [agent.to_pydantic() for agent in source.agents]
+            agents_orm = source.agents
+            return await asyncio.gather(*[agent.to_pydantic_async() for agent in agents_orm])
 
     # TODO: We make actor optional for now, but should most likely be enforced due to security reasons
     @enforce_types
-    def get_source_by_id(self, source_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSource]:
+    @trace_method
+    async def get_source_by_id(self, source_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticSource]:
         """Retrieve a source by its ID."""
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             try:
-                source = SourceModel.read(db_session=session, identifier=source_id, actor=actor)
+                source = await SourceModel.read_async(db_session=session, identifier=source_id, actor=actor)
                 return source.to_pydantic()
             except NoResultFound:
                 return None
 
     @enforce_types
-    def get_source_by_name(self, source_name: str, actor: PydanticUser) -> Optional[PydanticSource]:
+    @trace_method
+    async def get_source_by_name(self, source_name: str, actor: PydanticUser) -> Optional[PydanticSource]:
         """Retrieve a source by its name."""
-        with db_registry.session() as session:
-            sources = SourceModel.list(
+        async with db_registry.async_session() as session:
+            sources = await SourceModel.list_async(
                 db_session=session,
                 name=source_name,
                 organization_id=actor.organization_id,
@@ -131,44 +141,49 @@ class SourceManager:
                 return sources[0].to_pydantic()
 
     @enforce_types
-    def create_file(self, file_metadata: PydanticFileMetadata, actor: PydanticUser) -> PydanticFileMetadata:
+    @trace_method
+    async def create_file(self, file_metadata: PydanticFileMetadata, actor: PydanticUser) -> PydanticFileMetadata:
         """Create a new file based on the PydanticFileMetadata schema."""
-        db_file = self.get_file_by_id(file_metadata.id, actor=actor)
+        db_file = await self.get_file_by_id(file_metadata.id, actor=actor)
         if db_file:
             return db_file
         else:
-            with db_registry.session() as session:
+            async with db_registry.async_session() as session:
                 file_metadata.organization_id = actor.organization_id
                 file_metadata = FileMetadataModel(**file_metadata.model_dump(to_orm=True, exclude_none=True))
-                file_metadata.create(session, actor=actor)
+                await file_metadata.create_async(session, actor=actor)
             return file_metadata.to_pydantic()
 
     # TODO: We make actor optional for now, but should most likely be enforced due to security reasons
     @enforce_types
-    def get_file_by_id(self, file_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticFileMetadata]:
+    @trace_method
+    async def get_file_by_id(self, file_id: str, actor: Optional[PydanticUser] = None) -> Optional[PydanticFileMetadata]:
         """Retrieve a file by its ID."""
-        with db_registry.session() as session:
+        async with db_registry.async_session() as session:
             try:
-                file = FileMetadataModel.read(db_session=session, identifier=file_id, actor=actor)
+                file = await FileMetadataModel.read_async(db_session=session, identifier=file_id, actor=actor)
                 return file.to_pydantic()
             except NoResultFound:
                 return None
 
     @enforce_types
-    def list_files(
+    @trace_method
+    async def list_files(
         self, source_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50
     ) -> List[PydanticFileMetadata]:
         """List all files with optional pagination."""
-        with db_registry.session() as session:
-            files = FileMetadataModel.list(
+        async with db_registry.async_session() as session:
+            files_all = await FileMetadataModel.list_async(db_session=session, organization_id=actor.organization_id, source_id=source_id)
+            files = await FileMetadataModel.list_async(
                 db_session=session, after=after, limit=limit, organization_id=actor.organization_id, source_id=source_id
             )
             return [file.to_pydantic() for file in files]
 
     @enforce_types
-    def delete_file(self, file_id: str, actor: PydanticUser) -> PydanticFileMetadata:
+    @trace_method
+    async def delete_file(self, file_id: str, actor: PydanticUser) -> PydanticFileMetadata:
         """Delete a file by its ID."""
-        with db_registry.session() as session:
-            file = FileMetadataModel.read(db_session=session, identifier=file_id)
-            file.hard_delete(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            file = await FileMetadataModel.read_async(db_session=session, identifier=file_id)
+            await file.hard_delete_async(db_session=session, actor=actor)
             return file.to_pydantic()

letta/services/step_manager.py

@@ -14,13 +14,14 @@ from letta.schemas.step import Step as PydanticStep
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
 from letta.services.helpers.noop_helper import singleton
-from letta.tracing import get_trace_id
+from letta.tracing import get_trace_id, trace_method
 from letta.utils import enforce_types
 
 
 class StepManager:
 
     @enforce_types
+    @trace_method
     def list_steps(
         self,
         actor: PydanticUser,
@@ -54,6 +55,7 @@ class StepManager:
             return [step.to_pydantic() for step in steps]
 
     @enforce_types
+    @trace_method
     def log_step(
         self,
         actor: PydanticUser,
@@ -96,6 +98,7 @@ class StepManager:
             return new_step.to_pydantic()
 
     @enforce_types
+    @trace_method
     async def log_step_async(
         self,
         actor: PydanticUser,
@@ -138,12 +141,14 @@ class StepManager:
             return new_step.to_pydantic()
 
     @enforce_types
+    @trace_method
     def get_step(self, step_id: str, actor: PydanticUser) -> PydanticStep:
         with db_registry.session() as session:
             step = StepModel.read(db_session=session, identifier=step_id, actor=actor)
             return step.to_pydantic()
 
     @enforce_types
+    @trace_method
     def update_step_transaction_id(self, actor: PydanticUser, step_id: str, transaction_id: str) -> PydanticStep:
         """Update the transaction ID for a step.
 
@@ -236,6 +241,7 @@ class NoopStepManager(StepManager):
     """
 
     @enforce_types
+    @trace_method
     def log_step(
         self,
         actor: PydanticUser,
@@ -253,6 +259,7 @@ class NoopStepManager(StepManager):
         return
 
     @enforce_types
+    @trace_method
     async def log_step_async(
         self,
         actor: PydanticUser,