diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml
index 696465849a..3162fa3883 100644
--- a/.github/workflows/mypy.yml
+++ b/.github/workflows/mypy.yml
@@ -16,7 +16,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13', '3.14']
+        python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
 
     steps:
       - name: Checkout code
diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py
index 38ffb4b467..1b9fc8ca79 100644
--- a/vertexai/_genai/_agent_engines_utils.py
+++ b/vertexai/_genai/_agent_engines_utils.py
@@ -32,7 +32,6 @@
 from typing import (
     Any,
     AsyncIterator,
-    Awaitable,
     Callable,
     Coroutine,
     Dict,
@@ -59,6 +58,12 @@
 from . import types as genai_types
 
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+
 try:
     _BUILTIN_MODULE_NAMES: Sequence[str] = sys.builtin_module_names
 except AttributeError:
@@ -78,20 +83,30 @@
     _STDLIB_MODULE_NAMES: frozenset[str] = frozenset()  # type: ignore[no-redef]
 
 
-try:
-    from google.cloud import storage
+if typing.TYPE_CHECKING:
+    from google.cloud import storage  # type: ignore[attr-defined]
 
-    _StorageBucket: type[Any] = storage.Bucket
-except (ImportError, AttributeError):
-    _StorageBucket: type[Any] = Any  # type: ignore[no-redef]
+    _StorageBucket: TypeAlias = storage.Bucket
+else:
+    try:
+        from google.cloud import storage  # type: ignore[attr-defined]
 
+        _StorageBucket: type[Any] = storage.Bucket
+    except (ImportError, AttributeError):
+        _StorageBucket: type[Any] = Any  # type: ignore[no-redef]
 
-try:
+
+if typing.TYPE_CHECKING:
     import packaging
 
-    _SpecifierSet: type[Any] = packaging.specifiers.SpecifierSet
-except (ImportError, AttributeError):
-    _SpecifierSet: type[Any] = Any  # type: ignore[no-redef]
+    _SpecifierSet = packaging.specifiers.SpecifierSet
+else:
+    try:
+        import packaging
+
+        _SpecifierSet: type[Any] = packaging.specifiers.SpecifierSet
+    except (ImportError, AttributeError):
+        _SpecifierSet: type[Any] = Any  # type: ignore[no-redef]
 
 
 try:
@@ -258,16 +273,22 @@ class OperationRegistrable(Protocol):
     """Protocol for agents that have registered operations."""
 
     @abc.abstractmethod
-    def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:  # type: ignore[no-untyped-def]
+    def register_operations(self, **kwargs: Any) -> dict[str, list[str]]:
         """Register the user provided operations (modes and methods)."""
+        pass
 
 
-try:
+if typing.TYPE_CHECKING:
     from google.adk.agents import BaseAgent
 
-    ADKAgent: type[Any] = BaseAgent
-except (ImportError, AttributeError):
-    ADKAgent: type[Any] = Any  # type: ignore[no-redef]
+    ADKAgent: TypeAlias = BaseAgent
+else:
+    try:
+        from google.adk.agents import BaseAgent
+
+        ADKAgent: Optional[TypeAlias] = BaseAgent
+    except (ImportError, AttributeError):
+        ADKAgent = None  # type: ignore[no-redef]
 
 _AgentEngineInterface = Union[
     ADKAgent,
@@ -283,8 +304,9 @@ def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:  # type: ig
 class _ModuleAgentAttributes(TypedDict, total=False):
     module_name: str
     agent_name: str
-    register_operations: Dict[str, Sequence[str]]
+    register_operations: Dict[str, list[str]]
     sys_paths: Optional[Sequence[str]]
+    agent: _AgentEngineInterface
 
 
 class ModuleAgent(Cloneable, OperationRegistrable):
@@ -300,7 +322,7 @@ def __init__(
         *,
         module_name: str,
         agent_name: str,
-        register_operations: Dict[str, Sequence[str]],
+        register_operations: Dict[str, list[str]],
         sys_paths: Optional[Sequence[str]] = None,
     ):
         """Initializes a module-based agent.
@@ -310,7 +332,7 @@ def __init__(
                 Required. The name of the module to import.
             agent_name (str):
                 Required. The name of the agent in the module to instantiate.
-            register_operations (Dict[str, Sequence[str]]):
+            register_operations (Dict[str, list[str]]):
                 Required. A dictionary of API modes to a list of method names.
             sys_paths (Sequence[str]):
                 Optional. The system paths to search for the module. It should
@@ -336,8 +358,11 @@ def clone(self) -> "ModuleAgent":
             sys_paths=self._tmpl_attrs.get("sys_paths"),
         )
 
-    def register_operations(self) -> Dict[str, Sequence[str]]:
-        self._tmpl_attrs.get("register_operations")
+    def register_operations(self, **kwargs: Any) -> dict[str, list[str]]:
+        reg_operations = self._tmpl_attrs.get("register_operations")
+        if reg_operations is None:
+            raise ValueError("Register operations is not set.")
+        return reg_operations
 
     def set_up(self) -> None:
         """Sets up the agent for execution of queries at runtime.
@@ -411,7 +436,7 @@ def __call__(
 class GetAsyncOperationFunction(Protocol):
     async def __call__(
         self, *, operation_name: str, **kwargs: Any
-    ) -> Awaitable[AgentEngineOperationUnion]:
+    ) -> AgentEngineOperationUnion:
         pass
 
 
@@ -507,7 +532,7 @@ def _await_operation(
 def _compare_requirements(
     *,
     requirements: Mapping[str, str],
-    constraints: Union[Sequence[str], Mapping[str, "_SpecifierSet"]],
+    constraints: Union[Sequence[str], Mapping[str, Optional["_SpecifierSet"]]],
     required_packages: Optional[Iterator[str]] = None,
 ) -> _RequirementsValidationResult:
     """Compares the requirements with the constraints.
@@ -536,7 +561,7 @@ def _compare_requirements(
     """
     packaging_version = _import_packaging_version_or_raise()
     if required_packages is None:
-        required_packages = _DEFAULT_REQUIRED_PACKAGES
+        required_packages = _DEFAULT_REQUIRED_PACKAGES  # type: ignore[assignment]
     result = _RequirementsValidationResult(
         warnings=_RequirementsValidationWarnings(missing=set(), incompatible=set()),
         actions=_RequirementsValidationActions(append=set()),
@@ -583,7 +608,7 @@ def _generate_class_methods_spec_or_raise(
     if isinstance(agent, ModuleAgent):
         # We do a dry-run of setting up the agent engine to have the operations
         # needed for registration.
-        agent: ModuleAgent = agent.clone()
+        agent: ModuleAgent = agent.clone()  # type: ignore[no-redef]
         try:
             agent.set_up()
         except Exception as e:
@@ -819,13 +844,13 @@ def _get_gcs_bucket(
         new_bucket = storage_client.bucket(staging_bucket)
         gcs_bucket = storage_client.create_bucket(new_bucket, location=location)
         logger.info(f"Creating bucket {staging_bucket} in {location=}")
-    return gcs_bucket  # type: ignore[no-any-return]
+    return gcs_bucket
 
 
 def _get_registered_operations(
     *,
     agent: _AgentEngineInterface,
-) -> Dict[str, List[str]]:
+) -> dict[str, list[str]]:
     """Retrieves registered operations for a AgentEngine."""
     if isinstance(agent, OperationRegistrable):
         return agent.register_operations()
@@ -859,13 +884,13 @@ def _import_cloudpickle_or_raise() -> types.ModuleType:
 def _import_cloud_storage_or_raise() -> types.ModuleType:
     """Tries to import the Cloud Storage module."""
     try:
-        from google.cloud import storage
+        from google.cloud import storage  # type: ignore[attr-defined]
     except ImportError as e:
         raise ImportError(
             "Cloud Storage is not installed. Please call "
             "'pip install google-cloud-aiplatform[agent_engines]'."
         ) from e
-    return storage
+    return storage  # type: ignore[no-any-return]
 
 
 def _import_packaging_requirements_or_raise() -> types.ModuleType:
@@ -1202,7 +1227,7 @@ def _upload_agent_engine(
 ) -> None:
     """Uploads the agent engine to GCS."""
     cloudpickle = _import_cloudpickle_or_raise()
-    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}")  # type: ignore[attr-defined]
+    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}")
     with blob.open("wb") as f:
         try:
             cloudpickle.dump(agent, f)
@@ -1216,7 +1241,7 @@ def _upload_agent_engine(
             _ = cloudpickle.load(f)
         except Exception as e:
             raise TypeError("Agent engine serialized to an invalid format") from e
-    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"  # type: ignore[attr-defined]
+    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"
     logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}")
 
 
@@ -1227,9 +1252,9 @@ def _upload_requirements(
     gcs_dir_name: str,
 ) -> None:
     """Uploads the requirements file to GCS."""
-    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}")  # type: ignore[attr-defined]
+    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}")
     blob.upload_from_string("\n".join(requirements))
-    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"  # type: ignore[attr-defined]
+    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"
     logger.info(f"Writing to {dir_name}/{_REQUIREMENTS_FILE}")
 
 
@@ -1246,9 +1271,9 @@ def _upload_extra_packages(
         for file in extra_packages:
             tar.add(file)
     tar_fileobj.seek(0)
-    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}")  # type: ignore[attr-defined]
+    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}")
     blob.upload_from_string(tar_fileobj.read())
-    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"  # type: ignore[attr-defined]
+    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"
     logger.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}")
 
 
@@ -1369,7 +1394,7 @@ def _validate_requirements_or_warn(
     *,
     obj: Any,
     requirements: List[str],
-) -> Mapping[str, str]:
+) -> List[str]:
     """Compiles the requirements into a list of requirements."""
     requirements = requirements.copy()
     try:
@@ -1380,16 +1405,14 @@ def _validate_requirements_or_warn(
             requirements=current_requirements,
             constraints=constraints,
         )
-        for warning_type, warnings in missing_requirements.get(
-            _WARNINGS_KEY, {}
-        ).items():
+        for warning_type, warnings in missing_requirements["warnings"].items():
             if warnings:
                 logger.warning(
                     f"The following requirements are {warning_type}: {warnings}"
                 )
-        for action_type, actions in missing_requirements.get(_ACTIONS_KEY, {}).items():
+        for action_type, actions in missing_requirements["actions"].items():
             if actions and action_type == _ACTION_APPEND:
-                for action in actions:
+                for action in actions:  # type: ignore[attr-defined]
                     requirements.append(action)
                 logger.info(f"The following requirements are appended: {actions}")
     except Exception as e:
@@ -1413,7 +1436,7 @@ def _validate_requirements_or_raise(
                 logger.info(f"Read the following lines: {requirements}")
         except IOError as err:
             raise IOError(f"Failed to read requirements from {requirements=}") from err
-    requirements = _validate_requirements_or_warn(  # type: ignore[assignment]
+    requirements = _validate_requirements_or_warn(
         obj=agent,
         requirements=requirements,
     )
@@ -1560,19 +1583,6 @@ def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
     return _method
 
 
-AgentEngineOperationUnion = Union[
-    genai_types.AgentEngineOperation,
-    genai_types.AgentEngineMemoryOperation,
-    genai_types.AgentEngineGenerateMemoriesOperation,
-]
-
-
-class GetOperationFunction(Protocol):
-    def __call__(  # noqa: E704
-        self, *, operation_name: str, **kwargs: Any
-    ) -> AgentEngineOperationUnion: ...
-
-
 def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]:
     """Wraps an Agent Engine method, creating a callable for `query` API.
 
@@ -1835,7 +1845,7 @@ async def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
 
         return response
 
-    return _method
+    return _method  # type: ignore[return-value]
 
 
 def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py
index 6ad21e886b..0b4d8ecb9a 100644
--- a/vertexai/_genai/agent_engines.py
+++ b/vertexai/_genai/agent_engines.py
@@ -19,6 +19,7 @@
 import importlib
 import json
 import logging
+import typing
 from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Tuple, Union
 from urllib.parse import urlencode
 import warnings
@@ -33,6 +34,13 @@
 from . import _agent_engines_utils
 from . import types
 
+if typing.TYPE_CHECKING:
+    from . import sessions as sessions_module
+    from . import memories as memories_module
+
+    _ = sessions_module
+    __ = memories_module
+
 
 logger = logging.getLogger("vertexai_genai.agentengines")
 
@@ -703,7 +711,7 @@ def _update(
     _sessions = None
 
     @property
-    def memories(self) -> Any:
+    def memories(self) -> "memories_module.Memories":
         if self._memories is None:
             try:
                 # We need to lazy load the memories module to handle the
@@ -715,7 +723,7 @@ def memories(self) -> Any:
                     "packages. Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._memories.Memories(self._api_client)
+        return self._memories.Memories(self._api_client)  # type: ignore[no-any-return]
 
     @property
     def sandboxes(self) -> Any:
@@ -733,7 +741,7 @@ def sandboxes(self) -> Any:
         return self._sandboxes.Sandboxes(self._api_client)
 
     @property
-    def sessions(self) -> Any:
+    def sessions(self) -> "sessions_module.Sessions":
         if self._sessions is None:
             try:
                 # We need to lazy load the sessions module to handle the
@@ -745,7 +753,7 @@ def sessions(self) -> Any:
                     "Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._sessions.Sessions(self._api_client)
+        return self._sessions.Sessions(self._api_client)  # type: ignore[no-any-return]
 
     def _list_pager(
         self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None
@@ -987,7 +995,7 @@ def create(
             # If the user did not provide an agent_engine (e.g. lightweight
             # provisioning), it will not have any API methods registered.
             agent_engine = self._register_api_methods(agent_engine=agent_engine)
-        return agent_engine
+        return agent_engine  # type: ignore[no-any-return]
 
     def _set_source_code_spec(
         self,
@@ -1004,16 +1012,16 @@ def _set_source_code_spec(
         requirements_file: Optional[str] = None,
         sys_version: str,
         build_options: Optional[dict[str, list[str]]] = None,
-    ):
+    ) -> None:
         """Sets source_code_spec for agent engine inside the `spec`."""
-        source_code_spec = {}
+        source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict()
         if source_packages:
             source_packages = _agent_engines_utils._validate_packages_or_raise(
                 packages=source_packages,
                 build_options=build_options,
             )
             update_masks.append("spec.source_code_spec.inline_source.source_archive")
-            source_code_spec["inline_source"] = {
+            source_code_spec["inline_source"] = {  # type: ignore[typeddict-item]
                 "source_archive": _agent_engines_utils._create_base64_encoded_tarball(
                     source_packages=source_packages
                 )
@@ -1027,10 +1035,9 @@ def _set_source_code_spec(
             raise ValueError(
                 "Please specify one of `source_packages` or `developer_connect_source`."
             )
-            return
 
         update_masks.append("spec.source_code_spec.python_spec.version")
-        python_spec = {
+        python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = {
             "version": sys_version,
         }
         if not entrypoint_module:
@@ -1079,7 +1086,7 @@ def _set_package_spec(
         class_methods: Optional[Sequence[dict[str, Any]]] = None,
         sys_version: str,
         build_options: Optional[dict[str, list[str]]] = None,
-    ):
+    ) -> None:
         """Sets package spec for agent engine."""
         project = self._api_client.project
         if project is None:
@@ -1114,7 +1121,7 @@ def _set_package_spec(
         )
         # Update the package spec.
         update_masks.append("spec.package_spec.pickle_object_gcs_uri")
-        package_spec = {
+        package_spec: types.ReasoningEngineSpecPackageSpecDict = {
             "python_version": sys_version,
             "pickle_object_gcs_uri": "{}/{}/{}".format(
                 staging_bucket,
@@ -1441,10 +1448,10 @@ def _register_api_methods(
             _agent_engines_utils._register_api_methods_or_raise(
                 agent_engine=agent_engine,
                 wrap_operation_fn={
-                    "": _agent_engines_utils._wrap_query_operation,
-                    "async": _agent_engines_utils._wrap_async_query_operation,
-                    "stream": _agent_engines_utils._wrap_stream_query_operation,
-                    "async_stream": _agent_engines_utils._wrap_async_stream_query_operation,
+                    "": _agent_engines_utils._wrap_query_operation,  # type: ignore[dict-item]
+                    "async": _agent_engines_utils._wrap_async_query_operation,  # type: ignore[dict-item]
+                    "stream": _agent_engines_utils._wrap_stream_query_operation,  # type: ignore[dict-item]
+                    "async_stream": _agent_engines_utils._wrap_async_stream_query_operation,  # type: ignore[dict-item]
                     "a2a_extension": _agent_engines_utils._wrap_a2a_operation,
                 },
             )
@@ -1610,7 +1617,7 @@ def update(
             raise RuntimeError(f"Failed to update Agent Engine: {operation.error}")
         if agent_engine.api_resource.spec:
             self._register_api_methods(agent_engine=agent_engine)
-        return agent_engine
+        return agent_engine  # type: ignore[no-any-return]
 
     def _stream_query(
         self, *, name: str, config: Optional[types.QueryAgentEngineConfigOrDict] = None
@@ -1910,7 +1917,7 @@ def append_session_event(
             DeprecationWarning,
             stacklevel=2,
         )
-        return self.sessions.events.append(
+        return self.sessions.events.append(  # type: ignore[no-any-return]
             name=name,
             author=author,
             invocation_id=invocation_id,
@@ -1933,7 +1940,7 @@ def list_session_events(
             DeprecationWarning,
             stacklevel=2,
         )
-        return self.sessions.events.list(name=name, config=config)
+        return self.sessions.events.list(name=name, config=config)  # type: ignore[no-any-return]
 
 
 class AsyncAgentEngines(_api_module.BaseModule):
@@ -2355,7 +2362,7 @@ async def delete(
         return operation
 
     @property
-    def memories(self):
+    def memories(self) -> "memories_module.AsyncMemories":
         if self._memories is None:
             try:
                 # We need to lazy load the memories module to handle the
@@ -2367,10 +2374,10 @@ def memories(self):
                     "packages. Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._memories.AsyncMemories(self._api_client)
+        return self._memories.AsyncMemories(self._api_client)  # type: ignore[no-any-return]
 
     @property
-    def sessions(self):
+    def sessions(self) -> "sessions_module.AsyncSessions":
         if self._sessions is None:
             try:
                 # We need to lazy load the sessions module to handle the
@@ -2382,7 +2389,7 @@ def sessions(self):
                     "Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._sessions.AsyncSessions(self._api_client)
+        return self._sessions.AsyncSessions(self._api_client)  # type: ignore[no-any-return]
 
     async def append_session_event(
         self,
@@ -2402,7 +2409,7 @@ async def append_session_event(
             DeprecationWarning,
             stacklevel=2,
         )
-        return await self.sessions.events.append(name=name, config=config)
+        return await self.sessions.events.append(name=name, config=config)  # type: ignore[no-any-return]
 
     async def delete_memory(
         self,
diff --git a/vertexai/_genai/memories.py b/vertexai/_genai/memories.py
index 9f4f1dcde2..ac8ddeb5a3 100644
--- a/vertexai/_genai/memories.py
+++ b/vertexai/_genai/memories.py
@@ -19,6 +19,7 @@
 import importlib
 import json
 import logging
+import typing
 from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
@@ -31,6 +32,11 @@
 from . import _agent_engines_utils
 from . import types
 
+if typing.TYPE_CHECKING:
+    from . import memory_revisions as memory_revisions_module
+
+    _ = memory_revisions_module
+
 
 logger = logging.getLogger("vertexai_genai.memories")
 
@@ -1124,7 +1130,7 @@ def _purge(
     _revisions = None
 
     @property
-    def revisions(self):
+    def revisions(self) -> "memory_revisions_module.MemoryRevisions":
         if self._revisions is None:
             try:
                 # We need to lazy load the revisions module to handle the
@@ -1138,7 +1144,7 @@ def revisions(self):
                     "additional packages. Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._revisions.MemoryRevisions(self._api_client)
+        return self._revisions.MemoryRevisions(self._api_client)  # type: ignore[no-any-return]
 
     def create(
         self,
@@ -2096,7 +2102,7 @@ async def _purge(
     _revisions = None
 
     @property
-    def revisions(self):
+    def revisions(self) -> "memory_revisions_module.AsyncMemoryRevisions":
         if self._revisions is None:
             try:
                 # We need to lazy load the revisions module to handle the
@@ -2110,7 +2116,7 @@ def revisions(self):
                     "additional packages. Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._revisions.AsyncMemoryRevisions(self._api_client)
+        return self._revisions.AsyncMemoryRevisions(self._api_client)  # type: ignore[no-any-return]
 
     async def create(
         self,
diff --git a/vertexai/_genai/sandboxes.py b/vertexai/_genai/sandboxes.py
index 5a8c3f8fea..c88be4e269 100644
--- a/vertexai/_genai/sandboxes.py
+++ b/vertexai/_genai/sandboxes.py
@@ -25,7 +25,7 @@
 from urllib.parse import urlencode
 
 from google import genai
-from google.cloud import iam_credentials_v1
+from google.cloud import iam_credentials_v1  # type: ignore[attr-defined]
 from google.genai import _api_module
 from google.genai import _common
 from google.genai import types as genai_types
@@ -666,20 +666,20 @@ def execute_code(
         )
 
         output_chunks = []
-        for output in response.outputs:
-            if output.mime_type is None:
-                # if mime_type is not available, try to guess the mime_type from the file_name.
-                if (
-                    output.metadata is not None
-                    and output.metadata.attributes is not None
-                ):
-                    file_name = output.metadata.attributes.get("file_name", b"").decode(
-                        "utf-8"
-                    )
-                    mime_type, _ = mimetypes.guess_type(file_name)
-                    output.mime_type = mime_type
-
-            output_chunks.append(output)
+        if response.outputs is not None:
+            for output in response.outputs:
+                if output.mime_type is None:
+                    # if mime_type is not available, try to guess the mime_type from the file_name.
+                    if (
+                        output.metadata is not None
+                        and output.metadata.attributes is not None
+                    ):
+                        file_name = output.metadata.attributes.get(
+                            "file_name", b""
+                        ).decode("utf-8")
+                        mime_type, _ = mimetypes.guess_type(file_name)
+                        output.mime_type = mime_type
+                output_chunks.append(output)
 
         response = types.ExecuteSandboxEnvironmentResponse(outputs=output_chunks)
 
@@ -758,7 +758,7 @@ def generate_access_token(
             payload=json.dumps(payload),
         )
         response = client.sign_jwt(request=request)
-        return response.signed_jwt
+        return response.signed_jwt  # type: ignore[no-any-return]
 
     def send_command(
         self,
@@ -766,7 +766,7 @@ def send_command(
         http_method: str,
         access_token: str,
         sandbox_environment: types.SandboxEnvironment,
-        path: str = None,
+        path: Optional[str] = None,
         query_params: Optional[dict[str, object]] = None,
         headers: Optional[dict[str, str]] = None,
         request_dict: Optional[dict[str, object]] = None,
diff --git a/vertexai/_genai/sessions.py b/vertexai/_genai/sessions.py
index 31b83ec6dc..2e5cfdbb80 100644
--- a/vertexai/_genai/sessions.py
+++ b/vertexai/_genai/sessions.py
@@ -19,6 +19,7 @@
 import importlib
 import json
 import logging
+import typing
 from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
@@ -31,6 +32,11 @@
 from . import _agent_engines_utils
 from . import types
 
+if typing.TYPE_CHECKING:
+    from . import session_events as session_events_module
+
+    _ = session_events_module
+
 
 logger = logging.getLogger("vertexai_genai.sessions")
 
@@ -606,7 +612,7 @@ def _update(
         "The Vertex SDK GenAI agent_engines.sessions.events module is "
         "experimental, and may change in future versions."
     )
-    def events(self):
+    def events(self) -> "session_events_module.SessionEvents":
         if self._events is None:
             try:
                 # We need to lazy load the sessions.events module to handle the
@@ -618,7 +624,7 @@ def events(self):
                     "additional packages. Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._events.SessionEvents(self._api_client)
+        return self._events.SessionEvents(self._api_client)  # type: ignore[no-any-return]
 
     def create(
         self,
@@ -1092,7 +1098,7 @@ async def _update(
         "The Vertex SDK GenAI agent_engines.sessions.events module is "
         "experimental, and may change in future versions."
     )
-    def events(self):
+    def events(self) -> "session_events_module.AsyncSessionEvents":
         if self._events is None:
             try:
                 # We need to lazy load the sessions.events module to handle the
@@ -1104,7 +1110,7 @@ def events(self):
                     "additional packages. Please install them using pip install "
                     "google-cloud-aiplatform[agent_engines]"
                 ) from e
-        return self._events.AsyncSessionEvents(self._api_client)
+        return self._events.AsyncSessionEvents(self._api_client)  # type: ignore[no-any-return]
 
     async def create(
         self,