Fix V1 MCP services (Fix tavily search) (#11840)

Co-authored-by: openhands <openhands@all-hands.dev>
This commit is contained in:
Tim O'Farrell 2025-12-01 14:19:19 -07:00 committed by GitHub
parent 33eec7cb09
commit fd6e0cab3f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 1031 additions and 82 deletions

View File

@ -97,6 +97,9 @@ class GithubUserContext(UserContext):
user_secrets = await self.secrets_store.load()
return dict(user_secrets.custom_secrets) if user_secrets else {}
async def get_mcp_api_key(self) -> str | None:
raise NotImplementedError()
async def get_user_proactive_conversation_setting(user_id: str | None) -> bool:
"""Get the user's proactive conversation setting.

View File

@ -203,6 +203,15 @@ class SaasUserAuth(UserAuth):
self.settings_store = settings_store
return settings_store
async def get_mcp_api_key(self) -> str:
api_key_store = ApiKeyStore.get_instance()
mcp_api_key = api_key_store.retrieve_mcp_api_key(self.user_id)
if not mcp_api_key:
mcp_api_key = api_key_store.create_api_key(
self.user_id, 'MCP_API_KEY', None
)
return mcp_api_key
@classmethod
async def get_instance(cls, request: Request) -> UserAuth:
logger.debug('saas_user_auth_get_instance')

View File

@ -4,12 +4,12 @@ from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timedelta
from time import time
from typing import AsyncGenerator, Sequence
from typing import Any, AsyncGenerator, Sequence
from uuid import UUID, uuid4
import httpx
from fastapi import Request
from pydantic import Field, TypeAdapter
from pydantic import Field, SecretStr, TypeAdapter
from openhands.agent_server.models import (
ConversationInfo,
@ -63,19 +63,27 @@ from openhands.app_server.sandbox.sandbox_spec_service import SandboxSpecService
from openhands.app_server.services.injector import InjectorState
from openhands.app_server.services.jwt_service import JwtService
from openhands.app_server.user.user_context import UserContext
from openhands.app_server.user.user_models import UserInfo
from openhands.app_server.utils.docker_utils import (
replace_localhost_hostname_for_docker,
)
from openhands.experiments.experiment_manager import ExperimentManagerImpl
from openhands.integrations.provider import ProviderType
from openhands.sdk import AgentContext, LocalWorkspace
from openhands.sdk import Agent, AgentContext, LocalWorkspace
from openhands.sdk.conversation.secret_source import LookupSecret, StaticSecret
from openhands.sdk.llm import LLM
from openhands.sdk.security.confirmation_policy import AlwaysConfirm
from openhands.sdk.workspace.remote.async_remote_workspace import AsyncRemoteWorkspace
from openhands.server.types import AppMode
from openhands.tools.preset.default import get_default_agent
from openhands.tools.preset.planning import get_planning_agent
from openhands.tools.preset.default import (
get_default_condenser,
get_default_tools,
)
from openhands.tools.preset.planning import (
format_plan_structure,
get_planning_condenser,
get_planning_tools,
)
_conversation_info_type_adapter = TypeAdapter(list[ConversationInfo | None])
_logger = logging.getLogger(__name__)
@ -99,6 +107,7 @@ class LiveStatusAppConversationService(AppConversationServiceBase):
access_token_hard_timeout: timedelta | None
app_mode: str | None = None
keycloak_auth_cookie: str | None = None
tavily_api_key: str | None = None
async def search_app_conversations(
self,
@ -519,6 +528,223 @@ class LiveStatusAppConversationService(AppConversationServiceBase):
if not request.llm_model and parent_info.llm_model:
request.llm_model = parent_info.llm_model
async def _setup_secrets_for_git_provider(
self, git_provider: ProviderType | None, user: UserInfo
) -> dict:
"""Set up secrets for git provider authentication.
Args:
git_provider: The git provider type (GitHub, GitLab, etc.)
user: User information containing authentication details
Returns:
Dictionary of secrets for the conversation
"""
secrets = await self.user_context.get_secrets()
if not git_provider:
return secrets
secret_name = f'{git_provider.name}_TOKEN'
if self.web_url:
# Create an access token for web-based authentication
access_token = self.jwt_service.create_jws_token(
payload={
'user_id': user.id,
'provider_type': git_provider.value,
},
expires_in=self.access_token_hard_timeout,
)
headers = {'X-Access-Token': access_token}
# Include keycloak_auth cookie in headers if app_mode is SaaS
if self.app_mode == 'saas' and self.keycloak_auth_cookie:
headers['Cookie'] = f'keycloak_auth={self.keycloak_auth_cookie}'
secrets[secret_name] = LookupSecret(
url=self.web_url + '/api/v1/webhooks/secrets',
headers=headers,
)
else:
# Use static token for environments without web URL access
static_token = await self.user_context.get_latest_token(git_provider)
if static_token:
secrets[secret_name] = StaticSecret(value=static_token)
return secrets
async def _configure_llm_and_mcp(
self, user: UserInfo, llm_model: str | None
) -> tuple[LLM, dict]:
"""Configure LLM and MCP (Model Context Protocol) settings.
Args:
user: User information containing LLM preferences
llm_model: Optional specific model to use, falls back to user default
Returns:
Tuple of (configured LLM instance, MCP config dictionary)
"""
# Configure LLM
model = llm_model or user.llm_model
llm = LLM(
model=model,
base_url=user.llm_base_url,
api_key=user.llm_api_key,
usage_id='agent',
)
# Configure MCP
mcp_config: dict[str, Any] = {}
if self.web_url:
mcp_url = f'{self.web_url}/mcp/mcp'
mcp_config = {
'default': {
'url': mcp_url,
}
}
# Add API key if available
mcp_api_key = await self.user_context.get_mcp_api_key()
if mcp_api_key:
mcp_config['default']['headers'] = {
'X-Session-API-Key': mcp_api_key,
}
# Get the actual API key values, prioritizing user's key over service key
user_search_key = None
if user.search_api_key:
key_value = user.search_api_key.get_secret_value()
if key_value and key_value.strip():
user_search_key = key_value
service_tavily_key = None
if self.tavily_api_key:
# tavily_api_key is already a string (extracted in the factory method)
if self.tavily_api_key.strip():
service_tavily_key = self.tavily_api_key
tavily_api_key = user_search_key or service_tavily_key
if tavily_api_key:
_logger.info('Adding search engine to MCP config')
mcp_config['tavily'] = {
'url': f'https://mcp.tavily.com/mcp/?tavilyApiKey={tavily_api_key}'
}
else:
_logger.info('No search engine API key found, skipping search engine')
return llm, mcp_config
def _create_agent_with_context(
self,
llm: LLM,
agent_type: AgentType,
system_message_suffix: str | None,
mcp_config: dict,
) -> Agent:
"""Create an agent with appropriate tools and context based on agent type.
Args:
llm: Configured LLM instance
agent_type: Type of agent to create (PLAN or DEFAULT)
system_message_suffix: Optional suffix for system messages
mcp_config: MCP configuration dictionary
Returns:
Configured Agent instance with context
"""
# Create agent based on type
if agent_type == AgentType.PLAN:
agent = Agent(
llm=llm,
tools=get_planning_tools(),
system_prompt_filename='system_prompt_planning.j2',
system_prompt_kwargs={'plan_structure': format_plan_structure()},
condenser=get_planning_condenser(
llm=llm.model_copy(update={'usage_id': 'planning_condenser'})
),
security_analyzer=None,
mcp_config=mcp_config,
)
else:
agent = Agent(
llm=llm,
tools=get_default_tools(enable_browser=True),
system_prompt_kwargs={'cli_mode': False},
condenser=get_default_condenser(
llm=llm.model_copy(update={'usage_id': 'condenser'})
),
mcp_config=mcp_config,
)
# Add agent context
agent_context = AgentContext(system_message_suffix=system_message_suffix)
agent = agent.model_copy(update={'agent_context': agent_context})
return agent
async def _finalize_conversation_request(
self,
agent: Agent,
conversation_id: UUID | None,
user: UserInfo,
workspace: LocalWorkspace,
initial_message: SendMessageRequest | None,
secrets: dict,
sandbox: SandboxInfo,
remote_workspace: AsyncRemoteWorkspace | None,
selected_repository: str | None,
working_dir: str,
) -> StartConversationRequest:
"""Finalize the conversation request with experiment variants and skills.
Args:
agent: The configured agent
conversation_id: Optional conversation ID, generates new one if None
user: User information
workspace: Local workspace instance
initial_message: Optional initial message for the conversation
secrets: Dictionary of secrets for authentication
sandbox: Sandbox information
remote_workspace: Optional remote workspace for skills loading
selected_repository: Optional repository name
working_dir: Working directory path
Returns:
Complete StartConversationRequest ready for use
"""
# Generate conversation ID if not provided
conversation_id = conversation_id or uuid4()
# Apply experiment variants
agent = ExperimentManagerImpl.run_agent_variant_tests__v1(
user.id, conversation_id, agent
)
# Load and merge skills if remote workspace is available
if remote_workspace:
try:
agent = await self._load_skills_and_update_agent(
sandbox, agent, remote_workspace, selected_repository, working_dir
)
except Exception as e:
_logger.warning(f'Failed to load skills: {e}', exc_info=True)
# Continue without skills - don't fail conversation startup
# Create and return the final request
return StartConversationRequest(
conversation_id=conversation_id,
agent=agent,
workspace=workspace,
confirmation_policy=(
AlwaysConfirm() if user.confirmation_mode else NeverConfirm()
),
initial_message=initial_message,
secrets=secrets,
)
async def _build_start_conversation_request_for_user(
self,
sandbox: SandboxInfo,
@ -532,87 +758,41 @@ class LiveStatusAppConversationService(AppConversationServiceBase):
remote_workspace: AsyncRemoteWorkspace | None = None,
selected_repository: str | None = None,
) -> StartConversationRequest:
"""Build a complete conversation request for a user.
This method orchestrates the creation of a conversation request by:
1. Setting up git provider secrets
2. Configuring LLM and MCP settings
3. Creating an agent with appropriate context
4. Finalizing the request with skills and experiment variants
"""
user = await self.user_context.get_user_info()
# Set up a secret for the git token
secrets = await self.user_context.get_secrets()
if git_provider:
secret_name = f'{git_provider.name}_TOKEN'
if self.web_url:
# If there is a web url, then we create an access token to access it.
# For security reasons, we are explicit here - only this user, and
# only this provider, with a timeout
access_token = self.jwt_service.create_jws_token(
payload={
'user_id': user.id,
'provider_type': git_provider.value,
},
expires_in=self.access_token_hard_timeout,
)
headers = {'X-Access-Token': access_token}
# Include keycloak_auth cookie in headers if app_mode is SaaS
if self.app_mode == 'saas' and self.keycloak_auth_cookie:
headers['Cookie'] = f'keycloak_auth={self.keycloak_auth_cookie}'
secrets[secret_name] = LookupSecret(
url=self.web_url + '/api/v1/webhooks/secrets',
headers=headers,
)
else:
# If there is no URL specified where the sandbox can access the app server
# then we supply a static secret with the most recent value. Depending
# on the type, this may eventually expire.
static_token = await self.user_context.get_latest_token(git_provider)
if static_token:
secrets[secret_name] = StaticSecret(value=static_token)
workspace = LocalWorkspace(working_dir=working_dir)
# Use provided llm_model if available, otherwise fall back to user's default
model = llm_model or user.llm_model
llm = LLM(
model=model,
base_url=user.llm_base_url,
api_key=user.llm_api_key,
usage_id='agent',
)
# The agent gets passed initial instructions
# Select agent based on agent_type
if agent_type == AgentType.PLAN:
agent = get_planning_agent(llm=llm)
else:
agent = get_default_agent(llm=llm)
# Set up secrets for git provider
secrets = await self._setup_secrets_for_git_provider(git_provider, user)
agent_context = AgentContext(system_message_suffix=system_message_suffix)
agent = agent.model_copy(update={'agent_context': agent_context})
# Configure LLM and MCP
llm, mcp_config = await self._configure_llm_and_mcp(user, llm_model)
conversation_id = conversation_id or uuid4()
agent = ExperimentManagerImpl.run_agent_variant_tests__v1(
user.id, conversation_id, agent
# Create agent with context
agent = self._create_agent_with_context(
llm, agent_type, system_message_suffix, mcp_config
)
# Load and merge all skills if remote_workspace is available
if remote_workspace:
try:
agent = await self._load_skills_and_update_agent(
sandbox, agent, remote_workspace, selected_repository, working_dir
)
except Exception as e:
_logger.warning(f'Failed to load skills: {e}', exc_info=True)
# Continue without skills - don't fail conversation startup
start_conversation_request = StartConversationRequest(
conversation_id=conversation_id,
agent=agent,
workspace=workspace,
confirmation_policy=(
AlwaysConfirm() if user.confirmation_mode else NeverConfirm()
),
initial_message=initial_message,
secrets=secrets,
# Finalize and return the conversation request
return await self._finalize_conversation_request(
agent,
conversation_id,
user,
workspace,
initial_message,
secrets,
sandbox,
remote_workspace,
selected_repository,
working_dir,
)
return start_conversation_request
async def update_agent_server_conversation_title(
self,
@ -817,6 +997,10 @@ class LiveStatusAppConversationServiceInjector(AppConversationServiceInjector):
'be retrieved by a sandboxed conversation.'
),
)
tavily_api_key: SecretStr | None = Field(
default=None,
description='The Tavily Search API key to add to MCP integration',
)
async def inject(
self, state: InjectorState, request: Request | None = None
@ -874,6 +1058,14 @@ class LiveStatusAppConversationServiceInjector(AppConversationServiceInjector):
# If server_config is not available (e.g., in tests), continue without it
pass
# We supply the global tavily key only if the app mode is not SAAS, where
# currently the search endpoints are patched into the app server instead
# so the tavily key does not need to be shared
if self.tavily_api_key and app_mode != AppMode.SAAS:
tavily_api_key = self.tavily_api_key.get_secret_value()
else:
tavily_api_key = None
yield LiveStatusAppConversationService(
init_git_in_empty_workspace=self.init_git_in_empty_workspace,
user_context=user_context,
@ -890,4 +1082,5 @@ class LiveStatusAppConversationServiceInjector(AppConversationServiceInjector):
access_token_hard_timeout=access_token_hard_timeout,
app_mode=app_mode,
keycloak_auth_cookie=keycloak_auth_cookie,
tavily_api_key=tavily_api_key,
)

View File

@ -6,7 +6,7 @@ from typing import AsyncContextManager
import httpx
from fastapi import Depends, Request
from pydantic import Field
from pydantic import Field, SecretStr
from sqlalchemy.ext.asyncio import AsyncSession
# Import the event_callback module to ensure all processors are registered
@ -185,7 +185,13 @@ def config_from_env() -> AppServerConfig:
)
if config.app_conversation is None:
config.app_conversation = LiveStatusAppConversationServiceInjector()
tavily_api_key = None
tavily_api_key_str = os.getenv('TAVILY_API_KEY') or os.getenv('SEARCH_API_KEY')
if tavily_api_key_str:
tavily_api_key = SecretStr(tavily_api_key_str)
config.app_conversation = LiveStatusAppConversationServiceInjector(
tavily_api_key=tavily_api_key
)
if config.user is None:
config.user = AuthUserContextInjector()

View File

@ -78,6 +78,10 @@ class AuthUserContext(UserContext):
return results
async def get_mcp_api_key(self) -> str | None:
mcp_api_key = await self.user_auth.get_mcp_api_key()
return mcp_api_key
USER_ID_ATTR = 'user_id'

View File

@ -30,6 +30,9 @@ class SpecifyUserContext(UserContext):
async def get_secrets(self) -> dict[str, SecretSource]:
raise NotImplementedError()
async def get_mcp_api_key(self) -> str | None:
raise NotImplementedError()
USER_CONTEXT_ATTR = 'user_context'
ADMIN = SpecifyUserContext(user_id=None)

View File

@ -34,6 +34,10 @@ class UserContext(ABC):
async def get_secrets(self) -> dict[str, SecretSource]:
"""Get custom secrets and github provider secrets for the conversation."""
@abstractmethod
async def get_mcp_api_key(self) -> str | None:
"""Get an MCP API Key."""
class UserContextInjector(DiscriminatedUnionMixin, Injector[UserContext], ABC):
"""Injector for user contexts."""

View File

@ -88,6 +88,9 @@ class DefaultUserAuth(UserAuth):
return None
return user_secrets.provider_tokens
async def get_mcp_api_key(self) -> str | None:
return None
@classmethod
async def get_instance(cls, request: Request) -> UserAuth:
user_auth = DefaultUserAuth()

View File

@ -75,6 +75,10 @@ class UserAuth(ABC):
def get_auth_type(self) -> AuthType | None:
return None
@abstractmethod
async def get_mcp_api_key(self) -> str | None:
"""Get an mcp api key for the user"""
@classmethod
@abstractmethod
async def get_instance(cls, request: Request) -> UserAuth:

View File

@ -0,0 +1,698 @@
"""Unit tests for the methods in LiveStatusAppConversationService."""
from unittest.mock import AsyncMock, Mock, patch
from uuid import UUID, uuid4
import pytest
from openhands.agent_server.models import SendMessageRequest, StartConversationRequest
from openhands.app_server.app_conversation.app_conversation_models import AgentType
from openhands.app_server.app_conversation.live_status_app_conversation_service import (
LiveStatusAppConversationService,
)
from openhands.app_server.sandbox.sandbox_models import SandboxInfo, SandboxStatus
from openhands.app_server.user.user_context import UserContext
from openhands.integrations.provider import ProviderType
from openhands.sdk import Agent
from openhands.sdk.conversation.secret_source import LookupSecret, StaticSecret
from openhands.sdk.llm import LLM
from openhands.sdk.workspace import LocalWorkspace
from openhands.sdk.workspace.remote.async_remote_workspace import AsyncRemoteWorkspace
from openhands.server.types import AppMode
class TestLiveStatusAppConversationService:
"""Test cases for the methods in LiveStatusAppConversationService."""
def setup_method(self):
"""Set up test fixtures."""
# Create mock dependencies
self.mock_user_context = Mock(spec=UserContext)
self.mock_jwt_service = Mock()
self.mock_sandbox_service = Mock()
self.mock_sandbox_spec_service = Mock()
self.mock_app_conversation_info_service = Mock()
self.mock_app_conversation_start_task_service = Mock()
self.mock_event_callback_service = Mock()
self.mock_httpx_client = Mock()
# Create service instance
self.service = LiveStatusAppConversationService(
init_git_in_empty_workspace=True,
user_context=self.mock_user_context,
app_conversation_info_service=self.mock_app_conversation_info_service,
app_conversation_start_task_service=self.mock_app_conversation_start_task_service,
event_callback_service=self.mock_event_callback_service,
sandbox_service=self.mock_sandbox_service,
sandbox_spec_service=self.mock_sandbox_spec_service,
jwt_service=self.mock_jwt_service,
sandbox_startup_timeout=30,
sandbox_startup_poll_frequency=1,
httpx_client=self.mock_httpx_client,
web_url='https://test.example.com',
access_token_hard_timeout=None,
app_mode='test',
keycloak_auth_cookie=None,
)
# Mock user info
self.mock_user = Mock()
self.mock_user.id = 'test_user_123'
self.mock_user.llm_model = 'gpt-4'
self.mock_user.llm_base_url = 'https://api.openai.com/v1'
self.mock_user.llm_api_key = 'test_api_key'
self.mock_user.confirmation_mode = False
self.mock_user.search_api_key = None # Default to None
# Mock sandbox
self.mock_sandbox = Mock(spec=SandboxInfo)
self.mock_sandbox.id = uuid4()
self.mock_sandbox.status = SandboxStatus.RUNNING
@pytest.mark.asyncio
async def test_setup_secrets_for_git_provider_no_provider(self):
"""Test _setup_secrets_for_git_provider with no git provider."""
# Arrange
base_secrets = {'existing': 'secret'}
self.mock_user_context.get_secrets.return_value = base_secrets
# Act
result = await self.service._setup_secrets_for_git_provider(
None, self.mock_user
)
# Assert
assert result == base_secrets
self.mock_user_context.get_secrets.assert_called_once()
@pytest.mark.asyncio
async def test_setup_secrets_for_git_provider_with_web_url(self):
"""Test _setup_secrets_for_git_provider with web URL (creates access token)."""
# Arrange
base_secrets = {}
self.mock_user_context.get_secrets.return_value = base_secrets
self.mock_jwt_service.create_jws_token.return_value = 'test_access_token'
git_provider = ProviderType.GITHUB
# Act
result = await self.service._setup_secrets_for_git_provider(
git_provider, self.mock_user
)
# Assert
assert 'GITHUB_TOKEN' in result
assert isinstance(result['GITHUB_TOKEN'], LookupSecret)
assert (
result['GITHUB_TOKEN'].url
== 'https://test.example.com/api/v1/webhooks/secrets'
)
assert result['GITHUB_TOKEN'].headers['X-Access-Token'] == 'test_access_token'
self.mock_jwt_service.create_jws_token.assert_called_once_with(
payload={
'user_id': self.mock_user.id,
'provider_type': git_provider.value,
},
expires_in=None,
)
@pytest.mark.asyncio
async def test_setup_secrets_for_git_provider_with_saas_mode(self):
"""Test _setup_secrets_for_git_provider with SaaS mode (includes keycloak cookie)."""
# Arrange
self.service.app_mode = 'saas'
self.service.keycloak_auth_cookie = 'test_cookie'
base_secrets = {}
self.mock_user_context.get_secrets.return_value = base_secrets
self.mock_jwt_service.create_jws_token.return_value = 'test_access_token'
git_provider = ProviderType.GITLAB
# Act
result = await self.service._setup_secrets_for_git_provider(
git_provider, self.mock_user
)
# Assert
assert 'GITLAB_TOKEN' in result
lookup_secret = result['GITLAB_TOKEN']
assert isinstance(lookup_secret, LookupSecret)
assert 'Cookie' in lookup_secret.headers
assert lookup_secret.headers['Cookie'] == 'keycloak_auth=test_cookie'
@pytest.mark.asyncio
async def test_setup_secrets_for_git_provider_without_web_url(self):
"""Test _setup_secrets_for_git_provider without web URL (uses static token)."""
# Arrange
self.service.web_url = None
base_secrets = {}
self.mock_user_context.get_secrets.return_value = base_secrets
self.mock_user_context.get_latest_token.return_value = 'static_token_value'
git_provider = ProviderType.GITHUB
# Act
result = await self.service._setup_secrets_for_git_provider(
git_provider, self.mock_user
)
# Assert
assert 'GITHUB_TOKEN' in result
assert isinstance(result['GITHUB_TOKEN'], StaticSecret)
assert result['GITHUB_TOKEN'].value.get_secret_value() == 'static_token_value'
self.mock_user_context.get_latest_token.assert_called_once_with(git_provider)
@pytest.mark.asyncio
async def test_setup_secrets_for_git_provider_no_static_token(self):
"""Test _setup_secrets_for_git_provider when no static token is available."""
# Arrange
self.service.web_url = None
base_secrets = {}
self.mock_user_context.get_secrets.return_value = base_secrets
self.mock_user_context.get_latest_token.return_value = None
git_provider = ProviderType.GITHUB
# Act
result = await self.service._setup_secrets_for_git_provider(
git_provider, self.mock_user
)
# Assert
assert 'GITHUB_TOKEN' not in result
assert result == base_secrets
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_with_custom_model(self):
"""Test _configure_llm_and_mcp with custom LLM model."""
# Arrange
custom_model = 'gpt-3.5-turbo'
self.mock_user_context.get_mcp_api_key.return_value = 'mcp_api_key'
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, custom_model
)
# Assert
assert isinstance(llm, LLM)
assert llm.model == custom_model
assert llm.base_url == self.mock_user.llm_base_url
assert llm.api_key.get_secret_value() == self.mock_user.llm_api_key
assert llm.usage_id == 'agent'
assert 'default' in mcp_config
assert mcp_config['default']['url'] == 'https://test.example.com/mcp/mcp'
assert mcp_config['default']['headers']['X-Session-API-Key'] == 'mcp_api_key'
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_with_user_default_model(self):
"""Test _configure_llm_and_mcp using user's default model."""
# Arrange
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert llm.model == self.mock_user.llm_model
assert 'default' in mcp_config
assert 'headers' not in mcp_config['default']
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_without_web_url(self):
"""Test _configure_llm_and_mcp without web URL (no MCP config)."""
# Arrange
self.service.web_url = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert mcp_config == {}
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_tavily_with_user_search_api_key(self):
"""Test _configure_llm_and_mcp adds tavily when user has search_api_key."""
# Arrange
from pydantic import SecretStr
self.mock_user.search_api_key = SecretStr('user_search_key')
self.mock_user_context.get_mcp_api_key.return_value = 'mcp_api_key'
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'default' in mcp_config
assert 'tavily' in mcp_config
assert (
mcp_config['tavily']['url']
== 'https://mcp.tavily.com/mcp/?tavilyApiKey=user_search_key'
)
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_tavily_with_env_tavily_key(self):
"""Test _configure_llm_and_mcp adds tavily when service has tavily_api_key."""
# Arrange
self.service.tavily_api_key = 'env_tavily_key'
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'default' in mcp_config
assert 'tavily' in mcp_config
assert (
mcp_config['tavily']['url']
== 'https://mcp.tavily.com/mcp/?tavilyApiKey=env_tavily_key'
)
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_tavily_user_key_takes_precedence(self):
"""Test _configure_llm_and_mcp user search_api_key takes precedence over env key."""
# Arrange
from pydantic import SecretStr
self.mock_user.search_api_key = SecretStr('user_search_key')
self.service.tavily_api_key = 'env_tavily_key'
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'tavily' in mcp_config
assert (
mcp_config['tavily']['url']
== 'https://mcp.tavily.com/mcp/?tavilyApiKey=user_search_key'
)
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_no_tavily_without_keys(self):
"""Test _configure_llm_and_mcp does not add tavily when no keys are available."""
# Arrange
self.mock_user.search_api_key = None
self.service.tavily_api_key = None
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'default' in mcp_config
assert 'tavily' not in mcp_config
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_saas_mode_no_tavily_without_user_key(self):
"""Test _configure_llm_and_mcp does not add tavily in SAAS mode without user search_api_key.
In SAAS mode, the global tavily_api_key should not be passed to the service instance,
so tavily should only be added if the user has their own search_api_key.
"""
# Arrange - simulate SAAS mode where no global tavily key is available
self.service.app_mode = AppMode.SAAS.value
self.service.tavily_api_key = None # In SAAS mode, this should be None
self.mock_user.search_api_key = None
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'default' in mcp_config
assert 'tavily' not in mcp_config
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_saas_mode_with_user_search_key(self):
"""Test _configure_llm_and_mcp adds tavily in SAAS mode when user has search_api_key.
Even in SAAS mode, if the user has their own search_api_key, tavily should be added.
"""
# Arrange - simulate SAAS mode with user having their own search key
from pydantic import SecretStr
self.service.app_mode = AppMode.SAAS.value
self.service.tavily_api_key = None # In SAAS mode, this should be None
self.mock_user.search_api_key = SecretStr('user_search_key')
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'default' in mcp_config
assert 'tavily' in mcp_config
assert (
mcp_config['tavily']['url']
== 'https://mcp.tavily.com/mcp/?tavilyApiKey=user_search_key'
)
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_tavily_with_empty_user_search_key(self):
"""Test _configure_llm_and_mcp handles empty user search_api_key correctly."""
# Arrange
from pydantic import SecretStr
self.mock_user.search_api_key = SecretStr('') # Empty string
self.service.tavily_api_key = 'env_tavily_key'
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'tavily' in mcp_config
# Should fall back to env key since user key is empty
assert (
mcp_config['tavily']['url']
== 'https://mcp.tavily.com/mcp/?tavilyApiKey=env_tavily_key'
)
@pytest.mark.asyncio
async def test_configure_llm_and_mcp_tavily_with_whitespace_user_search_key(self):
"""Test _configure_llm_and_mcp handles whitespace-only user search_api_key correctly."""
# Arrange
from pydantic import SecretStr
self.mock_user.search_api_key = SecretStr(' ') # Whitespace only
self.service.tavily_api_key = 'env_tavily_key'
self.mock_user_context.get_mcp_api_key.return_value = None
# Act
llm, mcp_config = await self.service._configure_llm_and_mcp(
self.mock_user, None
)
# Assert
assert isinstance(llm, LLM)
assert 'tavily' in mcp_config
# Should fall back to env key since user key is whitespace only
assert (
mcp_config['tavily']['url']
== 'https://mcp.tavily.com/mcp/?tavilyApiKey=env_tavily_key'
)
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.get_planning_tools'
)
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.get_planning_condenser'
)
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.format_plan_structure'
)
def test_create_agent_with_context_planning_agent(
self, mock_format_plan, mock_get_condenser, mock_get_tools
):
"""Test _create_agent_with_context for planning agent type."""
# Arrange
mock_llm = Mock(spec=LLM)
mock_llm.model_copy.return_value = mock_llm
mock_get_tools.return_value = []
mock_get_condenser.return_value = Mock()
mock_format_plan.return_value = 'test_plan_structure'
mcp_config = {'default': {'url': 'test'}}
system_message_suffix = 'Test suffix'
# Act
with patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.Agent'
) as mock_agent_class:
mock_agent_instance = Mock()
mock_agent_instance.model_copy.return_value = mock_agent_instance
mock_agent_class.return_value = mock_agent_instance
self.service._create_agent_with_context(
mock_llm, AgentType.PLAN, system_message_suffix, mcp_config
)
# Assert
mock_agent_class.assert_called_once()
call_kwargs = mock_agent_class.call_args[1]
assert call_kwargs['llm'] == mock_llm
assert call_kwargs['system_prompt_filename'] == 'system_prompt_planning.j2'
assert (
call_kwargs['system_prompt_kwargs']['plan_structure']
== 'test_plan_structure'
)
assert call_kwargs['mcp_config'] == mcp_config
assert call_kwargs['security_analyzer'] is None
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.get_default_tools'
)
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.get_default_condenser'
)
def test_create_agent_with_context_default_agent(
self, mock_get_condenser, mock_get_tools
):
"""Test _create_agent_with_context for default agent type."""
# Arrange
mock_llm = Mock(spec=LLM)
mock_llm.model_copy.return_value = mock_llm
mock_get_tools.return_value = []
mock_get_condenser.return_value = Mock()
mcp_config = {'default': {'url': 'test'}}
# Act
with patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.Agent'
) as mock_agent_class:
mock_agent_instance = Mock()
mock_agent_instance.model_copy.return_value = mock_agent_instance
mock_agent_class.return_value = mock_agent_instance
self.service._create_agent_with_context(
mock_llm, AgentType.DEFAULT, None, mcp_config
)
# Assert
mock_agent_class.assert_called_once()
call_kwargs = mock_agent_class.call_args[1]
assert call_kwargs['llm'] == mock_llm
assert call_kwargs['system_prompt_kwargs']['cli_mode'] is False
assert call_kwargs['mcp_config'] == mcp_config
mock_get_tools.assert_called_once_with(enable_browser=True)
@pytest.mark.asyncio
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.ExperimentManagerImpl'
)
async def test_finalize_conversation_request_with_skills(
self, mock_experiment_manager
):
"""Test _finalize_conversation_request with skills loading."""
# Arrange
mock_agent = Mock(spec=Agent)
mock_updated_agent = Mock(spec=Agent)
mock_experiment_manager.run_agent_variant_tests__v1.return_value = (
mock_updated_agent
)
conversation_id = uuid4()
workspace = LocalWorkspace(working_dir='/test')
initial_message = Mock(spec=SendMessageRequest)
secrets = {'test': StaticSecret(value='secret')}
remote_workspace = Mock(spec=AsyncRemoteWorkspace)
# Mock the skills loading method
self.service._load_skills_and_update_agent = AsyncMock(
return_value=mock_updated_agent
)
# Act
result = await self.service._finalize_conversation_request(
mock_agent,
conversation_id,
self.mock_user,
workspace,
initial_message,
secrets,
self.mock_sandbox,
remote_workspace,
'test_repo',
'/test/dir',
)
# Assert
assert isinstance(result, StartConversationRequest)
assert result.conversation_id == conversation_id
assert result.agent == mock_updated_agent
assert result.workspace == workspace
assert result.initial_message == initial_message
assert result.secrets == secrets
mock_experiment_manager.run_agent_variant_tests__v1.assert_called_once_with(
self.mock_user.id, conversation_id, mock_agent
)
self.service._load_skills_and_update_agent.assert_called_once_with(
self.mock_sandbox,
mock_updated_agent,
remote_workspace,
'test_repo',
'/test/dir',
)
@pytest.mark.asyncio
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.ExperimentManagerImpl'
)
async def test_finalize_conversation_request_without_skills(
self, mock_experiment_manager
):
"""Test _finalize_conversation_request without remote workspace (no skills)."""
# Arrange
mock_agent = Mock(spec=Agent)
mock_updated_agent = Mock(spec=Agent)
mock_experiment_manager.run_agent_variant_tests__v1.return_value = (
mock_updated_agent
)
workspace = LocalWorkspace(working_dir='/test')
secrets = {'test': StaticSecret(value='secret')}
# Act
result = await self.service._finalize_conversation_request(
mock_agent,
None,
self.mock_user,
workspace,
None,
secrets,
self.mock_sandbox,
None,
None,
'/test/dir',
)
# Assert
assert isinstance(result, StartConversationRequest)
assert isinstance(result.conversation_id, UUID)
assert result.agent == mock_updated_agent
mock_experiment_manager.run_agent_variant_tests__v1.assert_called_once()
@pytest.mark.asyncio
@patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.ExperimentManagerImpl'
)
async def test_finalize_conversation_request_skills_loading_fails(
self, mock_experiment_manager
):
"""Test _finalize_conversation_request when skills loading fails."""
# Arrange
mock_agent = Mock(spec=Agent)
mock_updated_agent = Mock(spec=Agent)
mock_experiment_manager.run_agent_variant_tests__v1.return_value = (
mock_updated_agent
)
workspace = LocalWorkspace(working_dir='/test')
secrets = {'test': StaticSecret(value='secret')}
remote_workspace = Mock(spec=AsyncRemoteWorkspace)
# Mock skills loading to raise an exception
self.service._load_skills_and_update_agent = AsyncMock(
side_effect=Exception('Skills loading failed')
)
# Act
with patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service._logger'
) as mock_logger:
result = await self.service._finalize_conversation_request(
mock_agent,
None,
self.mock_user,
workspace,
None,
secrets,
self.mock_sandbox,
remote_workspace,
'test_repo',
'/test/dir',
)
# Assert
assert isinstance(result, StartConversationRequest)
assert (
result.agent == mock_updated_agent
) # Should still use the experiment-modified agent
mock_logger.warning.assert_called_once()
@pytest.mark.asyncio
async def test_build_start_conversation_request_for_user_integration(self):
"""Test the main _build_start_conversation_request_for_user method integration."""
# Arrange
self.mock_user_context.get_user_info.return_value = self.mock_user
# Mock all the helper methods
mock_secrets = {'GITHUB_TOKEN': Mock()}
mock_llm = Mock(spec=LLM)
mock_mcp_config = {'default': {'url': 'test'}}
mock_agent = Mock(spec=Agent)
mock_final_request = Mock(spec=StartConversationRequest)
self.service._setup_secrets_for_git_provider = AsyncMock(
return_value=mock_secrets
)
self.service._configure_llm_and_mcp = AsyncMock(
return_value=(mock_llm, mock_mcp_config)
)
self.service._create_agent_with_context = Mock(return_value=mock_agent)
self.service._finalize_conversation_request = AsyncMock(
return_value=mock_final_request
)
# Act
result = await self.service._build_start_conversation_request_for_user(
sandbox=self.mock_sandbox,
initial_message=None,
system_message_suffix='Test suffix',
git_provider=ProviderType.GITHUB,
working_dir='/test/dir',
agent_type=AgentType.DEFAULT,
llm_model='gpt-4',
conversation_id=None,
remote_workspace=None,
selected_repository='test/repo',
)
# Assert
assert result == mock_final_request
self.service._setup_secrets_for_git_provider.assert_called_once_with(
ProviderType.GITHUB, self.mock_user
)
self.service._configure_llm_and_mcp.assert_called_once_with(
self.mock_user, 'gpt-4'
)
self.service._create_agent_with_context.assert_called_once_with(
mock_llm, AgentType.DEFAULT, 'Test suffix', mock_mcp_config
)
self.service._finalize_conversation_request.assert_called_once()

View File

@ -200,8 +200,24 @@ class TestExperimentManagerIntegration:
# Patch the pieces invoked by the service
with (
patch(
'openhands.app_server.app_conversation.live_status_app_conversation_service.get_default_agent',
patch.object(
service,
'_setup_secrets_for_git_provider',
return_value={},
),
patch.object(
service,
'_configure_llm_and_mcp',
return_value=(mock_llm, {}),
),
patch.object(
service,
'_create_agent_with_context',
return_value=mock_agent,
),
patch.object(
service,
'_load_skills_and_update_agent',
return_value=mock_agent,
),
patch(

View File

@ -46,6 +46,9 @@ class MockUserAuth(UserAuth):
async def get_secrets(self) -> Secrets | None:
return None
async def get_mcp_api_key(self) -> str | None:
return None
@classmethod
async def get_instance(cls, request: Request) -> UserAuth:
return MockUserAuth()

View File

@ -46,6 +46,9 @@ class MockUserAuth(UserAuth):
async def get_secrets(self) -> Secrets | None:
return None
async def get_mcp_api_key(self) -> str | None:
return None
@classmethod
async def get_instance(cls, request: Request) -> UserAuth:
return MockUserAuth()