From 7dd3504add86196887bf6ac482f6dcfe9b4fd0b3 Mon Sep 17 00:00:00 2001 From: Anxhela Coba Date: Thu, 26 Feb 2026 16:36:35 -0500 Subject: [PATCH 1/6] updated new mcp header Signed-off-by: Anxhela Coba --- src/lightspeed_evaluation/core/api/client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/lightspeed_evaluation/core/api/client.py b/src/lightspeed_evaluation/core/api/client.py index a3acb821..8249a0ce 100644 --- a/src/lightspeed_evaluation/core/api/client.py +++ b/src/lightspeed_evaluation/core/api/client.py @@ -94,7 +94,10 @@ def _setup_client(self) -> None: # Use API_KEY environment variable for authentication api_key = os.getenv("API_KEY") if api_key and self.client: - self.client.headers.update({"Authorization": f"Bearer {api_key}"}) + mcp_headers = { + "filesystem-tools": {"Authorization": f"Bearer {api_key}"} + } + self.client.headers.update({"MCP-HEADERS": json.dumps(mcp_headers)}) except Exception as e: raise APIError(f"Failed to setup API client: {e}") from e From 2df22a8a6886b29a9b916cd52304ecda1d079abf Mon Sep 17 00:00:00 2001 From: Anxhela Coba Date: Mon, 2 Mar 2026 10:31:38 -0500 Subject: [PATCH 2/6] update system config Signed-off-by: Anxhela Coba --- config/system.yaml | 10 +++ src/lightspeed_evaluation/core/api/client.py | 75 +++++++++++++++++-- .../core/models/system.py | 47 ++++++++++++ 3 files changed, 125 insertions(+), 7 deletions(-) diff --git a/config/system.yaml b/config/system.yaml index 5b51cf1f..98f9c5f8 100644 --- a/config/system.yaml +++ b/config/system.yaml @@ -46,6 +46,16 @@ api: cache_dir: ".caches/api_cache" # Directory with lightspeed-stack cache cache_enabled: true # Is lightspeed-stack cache enabled? + + # MCP Server Authentication Configuration + mcp_headers: + enabled: true # Enable MCP headers functionality + servers: # MCP server configurations + filesystem-tools: + auth_type: bearer # Authentication type: bearer, api_key, custom + env_var: API_KEY # Environment variable containing the token/key + + # Legacy authentication (fallback when mcp_headers.enabled is false) # Authentication via API_KEY environment variable only for MCP server # Retry configuration for 429 Too Many Requests API errors diff --git a/src/lightspeed_evaluation/core/api/client.py b/src/lightspeed_evaluation/core/api/client.py index 8249a0ce..1c4050fc 100644 --- a/src/lightspeed_evaluation/core/api/client.py +++ b/src/lightspeed_evaluation/core/api/client.py @@ -91,17 +91,78 @@ def _setup_client(self) -> None: ) self.client.headers.update({"Content-Type": "application/json"}) - # Use API_KEY environment variable for authentication - api_key = os.getenv("API_KEY") - if api_key and self.client: - mcp_headers = { - "filesystem-tools": {"Authorization": f"Bearer {api_key}"} - } - self.client.headers.update({"MCP-HEADERS": json.dumps(mcp_headers)}) + # Set up MCP headers based on configuration + mcp_headers_dict = self._build_mcp_headers() + if mcp_headers_dict: + self.client.headers.update( + {"MCP-HEADERS": json.dumps(mcp_headers_dict)} + ) except Exception as e: raise APIError(f"Failed to setup API client: {e}") from e + def _build_mcp_headers(self) -> dict[str, dict[str, str]]: + """Build MCP headers based on configuration. + + Returns: + Dictionary of MCP server headers, or empty dict if none configured. + """ + # Use new MCP headers configuration if available and enabled + if ( + self.config.mcp_headers + and self.config.mcp_headers.enabled + and self.config.mcp_headers.servers + ): + + mcp_headers = {} + for server_name, server_config in self.config.mcp_headers.servers.items(): + # Get token from environment variable + token = os.getenv(server_config.env_var) + if not token: + logger.warning( + "Environment variable '%s' not found " + "for MCP server '%s'. Skipping authentication.", + server_config.env_var, + server_name, + ) + continue + + # Determine header name and value based on auth_type + if server_config.auth_type == "bearer": + header_name = server_config.header_name or "Authorization" + header_value = f"Bearer {token}" + elif server_config.auth_type == "api_key": + header_name = server_config.header_name or "X-API-Key" + header_value = token + elif server_config.auth_type == "custom": + if not server_config.header_name: + logger.warning( + "Custom auth_type for server '%s' requires " + "header_name to be specified. Skipping.", + server_name, + ) + continue + header_name = server_config.header_name + header_value = token + else: + logger.warning( + "Unknown auth_type '%s' for server '%s'. Skipping.", + server_config.auth_type, + server_name, + ) + continue + + mcp_headers[server_name] = {header_name: header_value} + + return mcp_headers + + # Fallback to legacy API_KEY behavior for backward compatibility + api_key = os.getenv("API_KEY") + if api_key: + return {"filesystem-tools": {"Authorization": f"Bearer {api_key}"}} + + return {} + def query( self, query: str, diff --git a/src/lightspeed_evaluation/core/models/system.py b/src/lightspeed_evaluation/core/models/system.py index 33952702..615eb5bc 100644 --- a/src/lightspeed_evaluation/core/models/system.py +++ b/src/lightspeed_evaluation/core/models/system.py @@ -160,6 +160,50 @@ def _validate_provider(cls, v: str) -> str: return v +class MCPServerConfig(BaseModel): + """Configuration for a single MCP server authentication.""" + + model_config = ConfigDict(extra="forbid") + + auth_type: str = Field( + ..., + description="Authentication type: bearer, api_key", + ) + env_var: str = Field( + ..., + min_length=1, + description="Environment variable containing the token/key", + ) + header_name: Optional[str] = Field( + default=None, + description="Custom header name (optional, defaults based on auth_type)", + ) + + @field_validator("auth_type") + @classmethod + def validate_auth_type(cls, v: str) -> str: + """Validate auth_type is supported.""" + allowed = {"bearer", "api_key", "custom"} + if v not in allowed: + raise ValueError(f"auth_type must be one of {allowed}") + return v + + +class MCPHeadersConfig(BaseModel): + """Configuration for MCP headers functionality.""" + + model_config = ConfigDict(extra="forbid") + + enabled: bool = Field( + default=True, + description="Enable MCP headers functionality", + ) + servers: dict[str, MCPServerConfig] = Field( + default_factory=dict, + description="MCP server configurations", + ) + + class APIConfig(BaseModel): """API configuration for dynamic data generation.""" @@ -196,6 +240,9 @@ class APIConfig(BaseModel): cache_enabled: bool = Field( default=True, description="Is caching of lightspeed-stack queries enabled?" ) + mcp_headers: Optional[MCPHeadersConfig] = Field( + default=None, description="MCP headers configuration for authentication" + ) num_retries: int = Field( default=DEFAULT_API_NUM_RETRIES, ge=0, From f9ba4f29b927f17d1a421e572943242e62f8d1e0 Mon Sep 17 00:00:00 2001 From: Anxhela Coba Date: Mon, 2 Mar 2026 13:44:17 -0500 Subject: [PATCH 3/6] PR coderabbit comments Signed-off-by: Anxhela Coba --- src/lightspeed_evaluation/core/api/client.py | 5 +++++ src/lightspeed_evaluation/core/models/system.py | 9 +++++++++ 2 files changed, 14 insertions(+) diff --git a/src/lightspeed_evaluation/core/api/client.py b/src/lightspeed_evaluation/core/api/client.py index 1c4050fc..f1deb720 100644 --- a/src/lightspeed_evaluation/core/api/client.py +++ b/src/lightspeed_evaluation/core/api/client.py @@ -154,6 +154,11 @@ def _build_mcp_headers(self) -> dict[str, dict[str, str]]: mcp_headers[server_name] = {header_name: header_value} + if not mcp_headers: + raise APIError( + "MCP headers are enabled, but no valid server credentials were resolved. " + "Check mcp_headers.servers and required environment variables." + ) return mcp_headers # Fallback to legacy API_KEY behavior for backward compatibility diff --git a/src/lightspeed_evaluation/core/models/system.py b/src/lightspeed_evaluation/core/models/system.py index 615eb5bc..71ec68a6 100644 --- a/src/lightspeed_evaluation/core/models/system.py +++ b/src/lightspeed_evaluation/core/models/system.py @@ -188,6 +188,15 @@ def validate_auth_type(cls, v: str) -> str: raise ValueError(f"auth_type must be one of {allowed}") return v + @model_validator(mode="after") + def validate_custom_header_name(self) -> "MCPServerConfig": + """Validate custom auth has an explicit header name.""" + if self.auth_type == "custom" and not self.header_name: + raise ConfigurationError( + "For auth_type='custom', 'header_name' must be provided." + ) + return self + class MCPHeadersConfig(BaseModel): """Configuration for MCP headers functionality.""" From 99f8a5b9cfd8c9dc44e3d362800ab3752c25fa68 Mon Sep 17 00:00:00 2001 From: Anxhela Coba Date: Wed, 4 Mar 2026 12:11:09 -0500 Subject: [PATCH 4/6] format and simplify auth type Signed-off-by: Anxhela Coba --- config/system.yaml | 4 +-- src/lightspeed_evaluation/core/api/client.py | 27 +++---------------- .../core/models/system.py | 18 +++---------- 3 files changed, 9 insertions(+), 40 deletions(-) diff --git a/config/system.yaml b/config/system.yaml index 98f9c5f8..9f7b60a9 100644 --- a/config/system.yaml +++ b/config/system.yaml @@ -41,7 +41,7 @@ api: # API input configuration provider: "openai" # LLM provider for queries model: "gpt-4o-mini" # Model to use for queries - no_tools: null # Whether to bypass tools and MCP servers (optional) + no_tools: false # Whether to bypass tools and MCP servers (optional) system_prompt: null # System prompt (default None) cache_dir: ".caches/api_cache" # Directory with lightspeed-stack cache @@ -52,7 +52,7 @@ api: enabled: true # Enable MCP headers functionality servers: # MCP server configurations filesystem-tools: - auth_type: bearer # Authentication type: bearer, api_key, custom + auth_type: bearer # Authentication type: only bearer is supported env_var: API_KEY # Environment variable containing the token/key # Legacy authentication (fallback when mcp_headers.enabled is false) diff --git a/src/lightspeed_evaluation/core/api/client.py b/src/lightspeed_evaluation/core/api/client.py index f1deb720..cad3714b 100644 --- a/src/lightspeed_evaluation/core/api/client.py +++ b/src/lightspeed_evaluation/core/api/client.py @@ -127,30 +127,9 @@ def _build_mcp_headers(self) -> dict[str, dict[str, str]]: ) continue - # Determine header name and value based on auth_type - if server_config.auth_type == "bearer": - header_name = server_config.header_name or "Authorization" - header_value = f"Bearer {token}" - elif server_config.auth_type == "api_key": - header_name = server_config.header_name or "X-API-Key" - header_value = token - elif server_config.auth_type == "custom": - if not server_config.header_name: - logger.warning( - "Custom auth_type for server '%s' requires " - "header_name to be specified. Skipping.", - server_name, - ) - continue - header_name = server_config.header_name - header_value = token - else: - logger.warning( - "Unknown auth_type '%s' for server '%s'. Skipping.", - server_config.auth_type, - server_name, - ) - continue + # Set up bearer auth headers + header_name = server_config.header_name or "Authorization" + header_value = f"Bearer {token}" mcp_headers[server_name] = {header_name: header_value} diff --git a/src/lightspeed_evaluation/core/models/system.py b/src/lightspeed_evaluation/core/models/system.py index 71ec68a6..3d0dfac5 100644 --- a/src/lightspeed_evaluation/core/models/system.py +++ b/src/lightspeed_evaluation/core/models/system.py @@ -166,8 +166,8 @@ class MCPServerConfig(BaseModel): model_config = ConfigDict(extra="forbid") auth_type: str = Field( - ..., - description="Authentication type: bearer, api_key", + default="bearer", + description="Authentication type: only bearer is supported", ) env_var: str = Field( ..., @@ -183,20 +183,10 @@ class MCPServerConfig(BaseModel): @classmethod def validate_auth_type(cls, v: str) -> str: """Validate auth_type is supported.""" - allowed = {"bearer", "api_key", "custom"} - if v not in allowed: - raise ValueError(f"auth_type must be one of {allowed}") + if v != "bearer": + raise ValueError("auth_type must be 'bearer'") return v - @model_validator(mode="after") - def validate_custom_header_name(self) -> "MCPServerConfig": - """Validate custom auth has an explicit header name.""" - if self.auth_type == "custom" and not self.header_name: - raise ConfigurationError( - "For auth_type='custom', 'header_name' must be provided." - ) - return self - class MCPHeadersConfig(BaseModel): """Configuration for MCP headers functionality.""" From e3d90a96f3df69f764b035dcbfda5a2e136d0ea7 Mon Sep 17 00:00:00 2001 From: Anxhela Coba Date: Wed, 4 Mar 2026 13:43:50 -0500 Subject: [PATCH 5/6] update config.md Signed-off-by: Anxhela Coba --- docs/configuration.md | 57 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index fdcb7e80..a850fad3 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -107,8 +107,6 @@ embedding: This section configures the inference API for generating the responses. It can be any Lightspeed-Core compatible API. Note that it can be easily integrated with other APIs with a minimal change. -Authentication via `API_KEY` environment variable only for MCP server. - | Setting (api.) | Default | Description | |----------------|---------|-------------| | enabled | `"true"` | Enable/disable API calls | @@ -121,6 +119,30 @@ Authentication via `API_KEY` environment variable only for MCP server. | system_prompt | `null` | Custom system prompt (optional) | | cache_dir | `".caches/api_cache"` | Directory with cached API responses | | cache_enabled | `true` | Is API cache enabled? | +| mcp_headers | `null` | MCP headers configuration for authentication (see below) | +| num_retries | `3` | Maximum number of retry attempts for API calls on 429 errors | + +### MCP Server Authentication + +The framework supports two methods for MCP server authentication: + +#### 1. New MCP Headers Configuration (Recommended) +The `mcp_headers` configuration provides a flexible way to configure authentication for individual MCP servers: + +| Setting (api.mcp_headers.) | Default | Description | +|----------------------------|---------|-------------| +| enabled | `true` | Enable/disable MCP headers functionality | +| servers | `{}` | Dictionary of MCP server configurations | + +For each server in `servers`, you can configure: + +| Setting (api.mcp_headers.servers..) | Default | Description | +|---------------------------------------------------|---------|-------------| +| auth_type | `"bearer"` | Authentication type (currently only "bearer" is supported) | +| env_var | required | Environment variable containing the authentication token | + +#### 2. Legacy Authentication (Fallback) +When `mcp_headers.enabled` is `false`, the system falls back to using the `API_KEY` environment variable for all MCP server authentication. ### API Modes @@ -137,6 +159,8 @@ Authentication via `API_KEY` environment variable only for MCP server. - **Reproducible results**: Same response data used across runs ### Example +#### Example Configuration + ```yaml api: enabled: true @@ -150,8 +174,37 @@ api: system_prompt: null cache_dir: ".caches/api_cache" cache_enabled: true + num_retries: 3 + + # MCP Server Authentication Configuration + mcp_headers: + enabled: true # Enable MCP headers functionality + servers: # MCP server configurations + filesystem-tools: + auth_type: bearer # Authentication type: only bearer is supported + env_var: API_KEY # Environment variable containing the token/key + another-mcp-server: + auth_type: bearer + env_var: ANOTHER_API_KEY # Use a different environment variable +``` + +#### Lightspeed Stack API Compatibility + +**Important Note for lightspeed-stack API users**: To use the MCP headers functionality with the lightspeed-stack API, you need to modify the `llama_stack_api/openai_responses.py` file in your lightspeed-stack installation: + +In the `OpenAIResponsesToolMCP` class, change the `authorization` parameter's `exclude` field from `True` to `False`: + +```python +# In llama_stack_api/openai_responses.py +class OpenAIResponsesToolMCP: + authorization: Optional[str] = Field( + default=None, + exclude=False # Change this from True to False + ) ``` +This change allows the authorization headers to be properly passed through to MCP servers. + ## Metrics Metrics are enabled globally (as described below) or within the input data for each individual conversation or individual turn (question/answer pair). To enable a metrics globally you need to set `default` meta data attribute to `true` From d9a49336b84a663947707a0688e75650bd1661d1 Mon Sep 17 00:00:00 2001 From: Anxhela Coba Date: Sun, 15 Mar 2026 18:55:20 -0400 Subject: [PATCH 6/6] address PR comments Signed-off-by: Anxhela Coba --- config/system.yaml | 1 - docs/configuration.md | 4 +--- src/lightspeed_evaluation/core/api/client.py | 11 ++++++----- src/lightspeed_evaluation/core/models/system.py | 14 +------------- 4 files changed, 8 insertions(+), 22 deletions(-) diff --git a/config/system.yaml b/config/system.yaml index 9f7b60a9..dc0b4dfe 100644 --- a/config/system.yaml +++ b/config/system.yaml @@ -52,7 +52,6 @@ api: enabled: true # Enable MCP headers functionality servers: # MCP server configurations filesystem-tools: - auth_type: bearer # Authentication type: only bearer is supported env_var: API_KEY # Environment variable containing the token/key # Legacy authentication (fallback when mcp_headers.enabled is false) diff --git a/docs/configuration.md b/docs/configuration.md index a850fad3..5e4930bf 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -138,8 +138,8 @@ For each server in `servers`, you can configure: | Setting (api.mcp_headers.servers..) | Default | Description | |---------------------------------------------------|---------|-------------| -| auth_type | `"bearer"` | Authentication type (currently only "bearer" is supported) | | env_var | required | Environment variable containing the authentication token | +| header_name | `"Authorization"` | Custom header name (optional) | #### 2. Legacy Authentication (Fallback) When `mcp_headers.enabled` is `false`, the system falls back to using the `API_KEY` environment variable for all MCP server authentication. @@ -181,10 +181,8 @@ api: enabled: true # Enable MCP headers functionality servers: # MCP server configurations filesystem-tools: - auth_type: bearer # Authentication type: only bearer is supported env_var: API_KEY # Environment variable containing the token/key another-mcp-server: - auth_type: bearer env_var: ANOTHER_API_KEY # Use a different environment variable ``` diff --git a/src/lightspeed_evaluation/core/api/client.py b/src/lightspeed_evaluation/core/api/client.py index cad3714b..f4cd4bf3 100644 --- a/src/lightspeed_evaluation/core/api/client.py +++ b/src/lightspeed_evaluation/core/api/client.py @@ -97,6 +97,11 @@ def _setup_client(self) -> None: self.client.headers.update( {"MCP-HEADERS": json.dumps(mcp_headers_dict)} ) + else: + # Use API_KEY environment variable for authentication (backward compatibility) + api_key = os.getenv("API_KEY") + if api_key and self.client: + self.client.headers.update({"Authorization": f"Bearer {api_key}"}) except Exception as e: raise APIError(f"Failed to setup API client: {e}") from e @@ -140,11 +145,7 @@ def _build_mcp_headers(self) -> dict[str, dict[str, str]]: ) return mcp_headers - # Fallback to legacy API_KEY behavior for backward compatibility - api_key = os.getenv("API_KEY") - if api_key: - return {"filesystem-tools": {"Authorization": f"Bearer {api_key}"}} - + # No MCP headers configured return {} def query( diff --git a/src/lightspeed_evaluation/core/models/system.py b/src/lightspeed_evaluation/core/models/system.py index 3d0dfac5..3946466d 100644 --- a/src/lightspeed_evaluation/core/models/system.py +++ b/src/lightspeed_evaluation/core/models/system.py @@ -165,10 +165,6 @@ class MCPServerConfig(BaseModel): model_config = ConfigDict(extra="forbid") - auth_type: str = Field( - default="bearer", - description="Authentication type: only bearer is supported", - ) env_var: str = Field( ..., min_length=1, @@ -176,17 +172,9 @@ class MCPServerConfig(BaseModel): ) header_name: Optional[str] = Field( default=None, - description="Custom header name (optional, defaults based on auth_type)", + description="Custom header name (optional, defaults to 'Authorization')", ) - @field_validator("auth_type") - @classmethod - def validate_auth_type(cls, v: str) -> str: - """Validate auth_type is supported.""" - if v != "bearer": - raise ValueError("auth_type must be 'bearer'") - return v - class MCPHeadersConfig(BaseModel): """Configuration for MCP headers functionality."""