PLTF-2686: Add sdk to mypy checking in the enterprise folder and fix the resulting errors (#14463)

Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: aivong-openhands <ai.vong@openhands.dev>

enterprise/dev_config/python/.pre-commit-config.yaml

@@ -50,6 +50,8 @@ repos:
           - ./
           - stripe==11.5.0
           - pygithub==2.6.1
+          - openhands-sdk==1.22.1
+          - openhands-tools==1.22.1
           - sqlalchemy>=2.0
         # Use -p (package) to avoid dual module name conflict when using MYPYPATH
         # MYPYPATH=enterprise allows resolving bare imports like "from integrations.xxx"

enterprise/dev_config/python/mypy.ini

@@ -14,3 +14,20 @@ disable_error_code = union-attr
 
 [mypy-enterprise.sync.install_gitlab_webhooks.*]
 disable_error_code = redundant-cast
+
+[mypy-openhands.memory.condenser.impl.*]
+disable_error_code = override
+
+[mypy-openhands.controller.agent]
+disable_error_code = typeddict-item
+
+[mypy-openhands.agenthub.*]
+disable_error_code = unreachable
+
+[mypy-openhands.llm.*]
+# Deprecated V0 folder
+disable_error_code = unreachable,assignment,arg-type
+
+[mypy-openai.*]
+follow_imports = skip
+ignore_missing_imports = True

enterprise/integrations/bitbucket/bitbucket_service.py

@@ -50,20 +50,29 @@ class SaaSBitBucketService(BitBucketService):
             offline_token = await self.token_manager.load_offline_token(
                 self.external_auth_id
             )
-            bitbucket_token = SecretStr(
-                await self.token_manager.get_idp_token_from_offline_token(
+            if offline_token:
+                bitbucket_token_str: (
+                    str | None
+                ) = await self.token_manager.get_idp_token_from_offline_token(
                     offline_token, ProviderType.BITBUCKET
                 )
-            )
+                bitbucket_token = (
+                    SecretStr(bitbucket_token_str) if bitbucket_token_str else None
+                )
+            else:
+                bitbucket_token = None
             logger.info(
-                f'Got BitBucket token {bitbucket_token.get_secret_value()} from external auth user ID: {self.external_auth_id}'
+                f'Got BitBucket token {bitbucket_token} from external auth user ID: {self.external_auth_id}'
             )
         elif self.user_id:
-            bitbucket_token = SecretStr(
+            bitbucket_token_str = (
                 await self.token_manager.get_idp_token_from_idp_user_id(
                     self.user_id, ProviderType.BITBUCKET
                 )
             )
+            bitbucket_token = (
+                SecretStr(bitbucket_token_str) if bitbucket_token_str else None
+            )
             logger.debug(
                 f'Got BitBucket token {bitbucket_token} from user ID: {self.user_id}'
             )

enterprise/integrations/bitbucket_data_center/bitbucket_dc_service.py

@@ -47,18 +47,26 @@ class SaaSBitbucketDCService(BitbucketDCService):
             offline_token = await self.token_manager.load_offline_token(
                 self.external_auth_id
             )
-            bitbucket_dc_token = SecretStr(
+            bitbucket_dc_token_str: str | None = (
                 await self.token_manager.get_idp_token_from_offline_token(
                     offline_token, ProviderType.BITBUCKET_DATA_CENTER
                 )
+                if offline_token
+                else None
+            )
+            bitbucket_dc_token = (
+                SecretStr(bitbucket_dc_token_str) if bitbucket_dc_token_str else None
             )
             logger.debug('Got Bitbucket DC token via external_auth_id')
         elif self.user_id:
-            bitbucket_dc_token = SecretStr(
+            bitbucket_dc_token_str = (
                 await self.token_manager.get_idp_token_from_idp_user_id(
                     self.user_id, ProviderType.BITBUCKET_DATA_CENTER
                 )
             )
+            bitbucket_dc_token = (
+                SecretStr(bitbucket_dc_token_str) if bitbucket_dc_token_str else None
+            )
             logger.debug('Got Bitbucket DC token via user_id')
         else:
             logger.warning('external_auth_token and user_id not set!')

enterprise/integrations/github/github_service.py

@@ -51,20 +51,22 @@ class SaaSGitHubService(GitHubService):
             offline_token = await self.token_manager.load_offline_token(
                 self.external_auth_id
             )
-            github_token = SecretStr(
+            github_token_str: str | None = (
                 await self.token_manager.get_idp_token_from_offline_token(
                     offline_token, ProviderType.GITHUB
                 )
+                if offline_token
+                else None
             )
+            github_token = SecretStr(github_token_str) if github_token_str else None
             logger.debug(
                 f'Got GitHub token {github_token} from external auth user ID: {self.external_auth_id}'
             )
         elif self.user_id:
-            github_token = SecretStr(
-                await self.token_manager.get_idp_token_from_idp_user_id(
-                    self.user_id, ProviderType.GITHUB
-                )
+            github_token_str = await self.token_manager.get_idp_token_from_idp_user_id(
+                self.user_id, ProviderType.GITHUB
             )
+            github_token = SecretStr(github_token_str) if github_token_str else None
             logger.debug(
                 f'Got GitHub token {github_token} from user ID: {self.user_id}'
             )

enterprise/integrations/gitlab/gitlab_service.py

@@ -59,20 +59,22 @@ class SaaSGitLabService(GitLabService):
             offline_token = await self.token_manager.load_offline_token(
                 self.external_auth_id
             )
-            gitlab_token = SecretStr(
+            gitlab_token_str: str | None = (
                 await self.token_manager.get_idp_token_from_offline_token(
                     offline_token, ProviderType.GITLAB
                 )
+                if offline_token
+                else None
             )
+            gitlab_token = SecretStr(gitlab_token_str) if gitlab_token_str else None
             logger.info(
-                f'Got GitLab token {gitlab_token.get_secret_value()} from external auth user ID: {self.external_auth_id}'
+                f'Got GitLab token {gitlab_token} from external auth user ID: {self.external_auth_id}'
             )
         elif self.user_id:
-            gitlab_token = SecretStr(
-                await self.token_manager.get_idp_token_from_idp_user_id(
-                    self.user_id, ProviderType.GITLAB
-                )
+            gitlab_token_str = await self.token_manager.get_idp_token_from_idp_user_id(
+                self.user_id, ProviderType.GITLAB
             )
+            gitlab_token = SecretStr(gitlab_token_str) if gitlab_token_str else None
             logger.debug(
                 f'Got Gitlab token {gitlab_token} from user ID: {self.user_id}'
             )

enterprise/integrations/gitlab/gitlab_view.py

@@ -400,9 +400,11 @@ class GitlabFactory:
         keycloak_user_id = await token_manager.get_user_id_from_idp_user_id(
             user_id, ProviderType.GITLAB
         )
-
+        # TODO: When keycloak_user_id is None, perhaps this should raise unauthorized.
         user_info = UserData(
-            user_id=user_id, username=username, keycloak_user_id=keycloak_user_id
+            user_id=user_id,
+            username=username,
+            keycloak_user_id=keycloak_user_id,  # type: ignore[arg-type]
         )
 
         if GitlabFactory.is_labeled_issue(message):

enterprise/integrations/resolver_context.py

@@ -81,8 +81,8 @@ class ResolverUserContext(UserContext):
             converted_secrets = {}
             for key, custom_secret in secrets.custom_secrets.items():
                 # Extract the secret value from CustomSecret and convert to StaticSecret
-                secret_value = custom_secret.secret.get_secret_value()
-                converted_secrets[key] = StaticSecret(value=secret_value)
+                secret_source: SecretSource = StaticSecret(value=custom_secret.secret)
+                converted_secrets[key] = secret_source
             return converted_secrets
         return {}
 

enterprise/integrations/v1_utils.py

@@ -92,6 +92,6 @@ async def get_saas_user_auth(
 
     user_auth = SaasUserAuth(
         user_id=keycloak_user_id,
-        refresh_token=SecretStr(offline_token),
+        refresh_token=SecretStr(offline_token or ''),
     )
     return user_auth

enterprise/server/auth/saas_user_auth.py

@@ -564,7 +564,7 @@ async def saas_user_auth_from_bearer(request: Request) -> SaasUserAuth | None:
         )
         saas_user_auth = SaasUserAuth(
             user_id=validation_result.user_id,
-            refresh_token=SecretStr(offline_token),
+            refresh_token=SecretStr(offline_token or ''),
             auth_type=AuthType.BEARER,
             api_key_org_id=validation_result.org_id,
             api_key_id=validation_result.key_id,
@@ -641,6 +641,6 @@ async def get_user_auth_from_keycloak_id(keycloak_user_id: str) -> UserAuth:
 
     user_auth = SaasUserAuth(
         user_id=keycloak_user_id,
-        refresh_token=SecretStr(offline_token),
+        refresh_token=SecretStr(offline_token or ''),
     )
     return user_auth

enterprise/server/routes/api_keys.py

@@ -3,7 +3,7 @@ from typing import cast
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, SecretStr, field_validator
 from server.auth.org_context import EFFECTIVE_ORG_ID
 from server.auth.saas_user_auth import SaasUserAuth
 from storage.api_key import ApiKey
@@ -51,7 +51,7 @@ async def store_byor_key_in_db(user_id: str, org_id: UUID, key: str) -> None:
             break
     if not org_member:
         return None
-    org_member.llm_api_key_for_byor = key
+    org_member.llm_api_key_for_byor = SecretStr(key)
     await OrgMemberStore.update_org_member(org_member)
 
 
@@ -312,7 +312,11 @@ async def get_current_api_key(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail='This API key was created before organization support. Please regenerate your API key to use this endpoint.',
         )
-
+    if saas_user_auth.api_key_id is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail='This endpoint requires API key authentication.',
+        )
     return CurrentApiKeyResponse(
         id=saas_user_auth.api_key_id,
         name=saas_user_auth.api_key_name,

enterprise/server/routes/billing.py

@@ -169,7 +169,7 @@ async def create_customer_setup_session(
         success_url=f'{base_url}?setup=success',
         cancel_url=f'{base_url}',
     )
-    return CreateBillingSessionResponse(redirect_url=checkout_session.url)
+    return CreateBillingSessionResponse(redirect_url=checkout_session.url)  # type: ignore[arg-type]
 
 
 # Endpoint to create a new Stripe checkout session for credit purchase
@@ -235,7 +235,7 @@ async def create_checkout_session(
         session.add(billing_session)
         await session.commit()
 
-    return CreateBillingSessionResponse(redirect_url=checkout_session.url)
+    return CreateBillingSessionResponse(redirect_url=checkout_session.url)  # type: ignore[arg-type]
 
 
 # Callback endpoint for successful Stripe payments - updates user credits and billing session status

enterprise/server/routes/org_models.py

@@ -199,8 +199,8 @@ class OrgResponse(BaseModel):
         return cls(
             id=str(org.id),
             name=org.name,
-            contact_name=org.contact_name,
-            contact_email=org.contact_email,
+            contact_name=org.contact_name,  # type: ignore[arg-type]
+            contact_email=org.contact_email,  # type: ignore[arg-type]
             conversation_expiration=org.conversation_expiration,
             remote_runtime_resource_factor=org.remote_runtime_resource_factor,
             billing_margin=org.billing_margin,
@@ -386,7 +386,7 @@ class OrgUpdate(BaseModel):
         member_settings = OrgMemberSettingsUpdate(
             agent_settings_diff=self.agent_settings_diff,
             conversation_settings_diff=self.conversation_settings_diff,
-            llm_api_key=self.llm_api_key or None,
+            llm_api_key=SecretStr(self.llm_api_key) if self.llm_api_key else None,
         )
         return member_settings if member_settings.has_updates() else None
 

enterprise/server/services/org_invitation_service.py

@@ -386,7 +386,7 @@ class OrgInvitationService:
         # continue to flow through automatically.
         llm_api_key_secret = settings.agent_settings.llm.api_key
         llm_api_key = (
-            llm_api_key_secret.get_secret_value() if llm_api_key_secret else ''
+            llm_api_key_secret.get_secret_value() if llm_api_key_secret else ''  # type: ignore[union-attr]
         )
 
         await OrgMemberStore.add_user_to_org(

enterprise/server/sharing/sql_shared_conversation_info_service.py

@@ -98,17 +98,17 @@ class SQLSharedConversationInfoService(SharedConversationInfoService):
 
         # Rebuild token usage
         token_usage = TokenUsage(
-            prompt_tokens=stored.prompt_tokens,
-            completion_tokens=stored.completion_tokens,
-            cache_read_tokens=stored.cache_read_tokens,
-            cache_write_tokens=stored.cache_write_tokens,
-            context_window=stored.context_window,
-            per_turn_token=stored.per_turn_token,
+            prompt_tokens=stored.prompt_tokens,  # type: ignore[arg-type]
+            completion_tokens=stored.completion_tokens,  # type: ignore[arg-type]
+            cache_read_tokens=stored.cache_read_tokens,  # type: ignore[arg-type]
+            cache_write_tokens=stored.cache_write_tokens,  # type: ignore[arg-type]
+            context_window=stored.context_window,  # type: ignore[arg-type]
+            per_turn_token=stored.per_turn_token,  # type: ignore[arg-type]
         )
 
         # Rebuild metrics object
         metrics = MetricsSnapshot(
-            accumulated_cost=stored.accumulated_cost,
+            accumulated_cost=stored.accumulated_cost,  # type: ignore[arg-type]
             max_budget_per_task=stored.max_budget_per_task,
             accumulated_token_usage=token_usage,
         )
@@ -127,14 +127,14 @@ class SQLSharedConversationInfoService(SharedConversationInfoService):
         return SharedConversation(
             id=UUID(stored.conversation_id),
             created_by_user_id=created_by_user_id,
-            sandbox_id=stored.sandbox_id,
+            sandbox_id=stored.sandbox_id,  # type: ignore[arg-type]
             selected_repository=stored.selected_repository,
             selected_branch=stored.selected_branch,
             git_provider=(
                 ProviderType(stored.git_provider) if stored.git_provider else None
             ),
             title=stored.title,
-            pr_number=stored.pr_number,
+            pr_number=stored.pr_number,  # type: ignore[arg-type]
             llm_model=stored.llm_model,
             metrics=metrics,
             parent_conversation_id=(

enterprise/storage/lite_llm_manager.py

@@ -224,7 +224,7 @@ class LiteLlmManager:
                 try:
                     await LiteLlmManager._delete_key_by_alias(client, key_alias)
                 except httpx.HTTPStatusError as ex:
-                    if ex.status_code == 404:
+                    if ex.response and ex.response.status_code == 404:
                         logger.debug(f'Key "{key_alias}" did not exist - continuing')
                     else:
                         raise
@@ -1635,8 +1635,9 @@ class LiteLlmManager:
     ) -> Callable[..., Awaitable[Any]]:
         @functools.wraps(internal_fn)
         async def wrapper(*args, **kwargs):
+            headers = {'x-goog-api-key': LITE_LLM_API_KEY} if LITE_LLM_API_KEY else {}
             async with httpx.AsyncClient(
-                headers={'x-goog-api-key': LITE_LLM_API_KEY},
+                headers=headers,
                 timeout=httpx.Timeout(30.0),
             ) as client:
                 return await internal_fn(client, *args, **kwargs)

enterprise/storage/org_store.py

@@ -274,7 +274,7 @@ class OrgStore:
             merged_settings['agent_kind'] = 'openhands'
             return OpenHandsAgentSettings.model_validate(merged_settings)
 
-        base_settings = _load_persisted_conversation_settings(current_settings)
+        base_settings = _load_persisted_conversation_settings(current_settings)  # type: ignore[assignment]
         merged_settings = deep_merge(
             base_settings.model_dump(mode='json'), settings_diff
         )

enterprise/storage/saas_settings_store.py

@@ -277,7 +277,7 @@ class SaasSettingsStore(SettingsStore):
                 else None
             )
             current_member_llm_api_key_raw = (
-                current_member_llm_api_key.get_secret_value()
+                current_member_llm_api_key.get_secret_value()  # type: ignore[union-attr]
                 if current_member_llm_api_key
                 else None
             )
@@ -289,7 +289,7 @@ class SaasSettingsStore(SettingsStore):
                     agent_settings_diff=effective_agent_settings_diff,
                     conversation_settings_diff=effective_conversation_diff,
                     llm_api_key=(
-                        current_member_llm_api_key_raw
+                        current_member_llm_api_key_raw  # type: ignore[arg-type]
                         if not uses_managed_llm_key
                         else None
                     ),
@@ -298,11 +298,11 @@ class SaasSettingsStore(SettingsStore):
 
             if uses_managed_llm_key and current_member_llm_api_key is not None:
                 # Managed/proxy key — store on this member but mark as org-managed
-                org_member.llm_api_key = current_member_llm_api_key
+                org_member.llm_api_key = current_member_llm_api_key  # type: ignore[assignment]
                 org_member.has_custom_llm_api_key = False
             elif current_member_llm_api_key_raw is not None:
                 # BYOR: member supplied their own (non-managed) API key
-                org_member.llm_api_key = current_member_llm_api_key
+                org_member.llm_api_key = current_member_llm_api_key  # type: ignore[assignment]
                 org_member.has_custom_llm_api_key = True
             elif org_default_llm_api_key_raw is not None:
                 # No member key, falling back to org default
@@ -343,7 +343,7 @@ class SaasSettingsStore(SettingsStore):
 
         # First, check if our current key is valid
         if llm_api_key and not await LiteLlmManager.verify_existing_key(
-            llm_api_key.get_secret_value(),
+            llm_api_key.get_secret_value(),  # type: ignore[union-attr]
             self.user_id,
             org_id,
             openhands_type=openhands_type,