Add default headers to the Oxylabs API requests (#13)

* Add default headers to the Oxylabs API requests

* Updated URL params

Author: oxy-rostyslav

.github/workflows/lint_and_test.yml

@@ -37,4 +37,4 @@ jobs:
 
     - name: Run tests
       run: |
-        uv run pytest --cov=src --cov-report xml --cov-report term --cov-fail-under=80 ./tests
+        uv run pytest --cov=src --cov-report xml --cov-report term --cov-fail-under=90 ./tests

Makefile

@@ -22,7 +22,7 @@ format: $(virtualenv_dir)
 
 .PHONY: test
 test: install_deps
-	uv run pytest --cov=src --cov-report xml --cov-report term --cov-fail-under=80 ./tests
+	uv run pytest --cov=src --cov-report xml --cov-report term --cov-fail-under=90 ./tests
 
 .PHONY: run
 run: install_deps

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "oxylabs-mcp"
-version = "0.1.5"
+version = "0.1.6"
 description = "Oxylabs MCP server"
 authors = [
     {name="Augis Braziunas", email="[email protected]"},

src/oxylabs_mcp/server.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from mcp.server.fastmcp import FastMCP
+from mcp.server.fastmcp import Context, FastMCP
 
 from oxylabs_mcp import url_params
 from oxylabs_mcp.config import settings
@@ -21,13 +21,14 @@
     description="Scrape url using Oxylabs Web API with universal scraper",
 )
 async def scrape_universal_url(
+    ctx: Context,  # type: ignore[type-arg]
     url: url_params.URL_PARAM,
     parse: url_params.PARSE_PARAM = True,  # noqa: FBT002
     render: url_params.RENDER_PARAM = "",
 ) -> str:
     """Scrape url using Oxylabs Web API with universal scraper."""
     try:
-        async with oxylabs_client(with_auth=True) as client:
+        async with oxylabs_client(ctx, with_auth=True) as client:
             payload: dict[str, Any] = {"url": url}
             if parse:
                 payload["parse"] = parse
@@ -48,6 +49,7 @@ async def scrape_universal_url(
     description="Scrape url using Oxylabs Web Unblocker",
 )
 async def scrape_with_web_unblocker(
+    ctx: Context,  # type: ignore[type-arg]
     url: url_params.URL_PARAM,
     render: url_params.RENDER_PARAM = "",
 ) -> str:
@@ -61,7 +63,7 @@ async def scrape_with_web_unblocker(
         headers["X-Oxylabs-Render"] = render
 
     try:
-        async with oxylabs_client(with_proxy=True, verify=False, headers=headers) as client:
+        async with oxylabs_client(ctx, with_proxy=True, verify=False, headers=headers) as client:
             response = await client.get(url)
 
             response.raise_for_status()
@@ -76,6 +78,7 @@ async def scrape_with_web_unblocker(
     description="Scrape Google Search results using Oxylabs Web API",
 )
 async def scrape_google_search(
+    ctx: Context,  # type: ignore[type-arg]
     query: url_params.GOOGLE_QUERY_PARAM,
     parse: url_params.PARSE_PARAM = True,  # noqa: FBT002
     render: url_params.RENDER_PARAM = "",
@@ -90,7 +93,7 @@ async def scrape_google_search(
 ) -> str:
     """Scrape Google Search results using Oxylabs Web API."""
     try:
-        async with oxylabs_client(with_auth=True) as client:
+        async with oxylabs_client(ctx, with_auth=True) as client:
             payload: dict[str, Any] = {"query": query}
 
             if ad_mode:
@@ -131,6 +134,7 @@ async def scrape_google_search(
     description="Scrape Amazon Search results using Oxylabs Web API",
 )
 async def scrape_amazon_search(
+    ctx: Context,  # type: ignore[type-arg]
     query: url_params.AMAZON_SEARCH_QUERY_PARAM,
     category_id: url_params.CATEGORY_ID_CONTEXT_PARAM = "",
     merchant_id: url_params.MERCHANT_ID_CONTEXT_PARAM = "",
@@ -146,7 +150,7 @@ async def scrape_amazon_search(
 ) -> str:
     """Scrape Amazon Search results using Oxylabs Web API."""
     try:
-        async with oxylabs_client(with_auth=True) as client:
+        async with oxylabs_client(ctx, with_auth=True) as client:
             payload: dict[str, Any] = {"source": "amazon_search", "query": query}
 
             context = []
@@ -190,6 +194,7 @@ async def scrape_amazon_search(
     description="Scrape Amazon Products using Oxylabs Web API",
 )
 async def scrape_amazon_products(
+    ctx: Context,  # type: ignore[type-arg]
     query: url_params.AMAZON_SEARCH_QUERY_PARAM,
     autoselect_variant: url_params.AUTOSELECT_VARIANT_CONTEXT_PARAM = False,  # noqa: FBT002
     currency: url_params.CURRENCY_CONTEXT_PARAM = "",
@@ -202,7 +207,7 @@ async def scrape_amazon_products(
 ) -> str:
     """Scrape Amazon Products using Oxylabs Web API."""
     try:
-        async with oxylabs_client(with_auth=True) as client:
+        async with oxylabs_client(ctx, with_auth=True) as client:
             payload: dict[str, Any] = {"source": "amazon_product", "query": query}
 
             context = []

src/oxylabs_mcp/url_params.py

@@ -51,15 +51,15 @@
 ]
 START_PAGE_PARAM = Annotated[
     int,
-    Field(ge=1, description="Starting page number."),
+    Field(description="Starting page number."),
 ]
 PAGES_PARAM = Annotated[
     int,
-    Field(ge=1, description="Number of pages to retrieve."),
+    Field(description="Number of pages to retrieve."),
 ]
 LIMIT_PARAM = Annotated[
     int,
-    Field(ge=1, description="Number of results to retrieve in each page."),
+    Field(description="Number of results to retrieve in each page."),
 ]
 DOMAIN_PARAM = Annotated[
     str,

src/oxylabs_mcp/utils.py

@@ -2,6 +2,8 @@
 import os
 import re
 from contextlib import asynccontextmanager
+from importlib.metadata import version
+from platform import architecture, python_version
 from typing import AsyncIterator
 
 from httpx import (
@@ -15,6 +17,8 @@
 from lxml.html import defs, fromstring, tostring
 from lxml.html.clean import Cleaner
 from markdownify import markdownify as md
+from mcp.server.fastmcp import Context
+from mcp.shared.context import RequestContext
 
 from oxylabs_mcp.config import settings
 from oxylabs_mcp.exceptions import MCPServerError
@@ -110,8 +114,33 @@ def convert_html_to_md(html: str) -> str:
     return md(html)  # type: ignore[no-any-return]
 
 
+def _get_request_context(ctx: Context) -> RequestContext | None:  # type: ignore[type-arg]
+    try:
+        return ctx.request_context
+    except ValueError:
+        return None
+
+
+def _update_with_default_headers(
+    ctx: Context, headers: dict[str, str]  # type: ignore[type-arg]
+) -> None:
+    if request_context := _get_request_context(ctx):
+        if client_params := request_context.session.client_params:
+            client = f"oxylabs-mcp-{client_params.clientInfo.name}"
+        else:
+            client = "oxylabs-mcp"
+    else:
+        client = "oxylabs-mcp"
+
+    bits, _ = architecture()
+    sdk_type = f"{client}/{version('oxylabs-mcp')} ({python_version()}; {bits})"
+
+    headers["x-oxylabs-sdk"] = sdk_type
+
+
 @asynccontextmanager
 async def oxylabs_client(
+    ctx: Context,  # type: ignore[type-arg]
     headers: dict[str, str] | None = None,
     *,
     with_proxy: bool = False,
@@ -122,6 +151,8 @@ async def oxylabs_client(
     if headers is None:
         headers = {}
 
+    _update_with_default_headers(ctx, headers)
+
     username, password = get_auth_from_env()
 
     if with_proxy:

tests/conftest.py

@@ -1,7 +1,9 @@
 from contextlib import asynccontextmanager
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from mcp.server.lowlevel.server import request_ctx
+from mcp.shared.context import RequestContext
 
 
 @pytest.fixture
@@ -10,7 +12,28 @@ def oxylabs_client():
 
     @asynccontextmanager
     async def wrapper(*args, **kwargs):
+        client_mock.context_manager_call_args = args
+        client_mock.context_manager_call_kwargs = kwargs
+
         yield client_mock
 
     with patch("oxylabs_mcp.utils.AsyncClient", new=wrapper):
         yield client_mock
+
+
+@pytest.fixture
+def request_session():
+    request_session_mock = MagicMock()
+
+    token = request_ctx.set(
+        RequestContext(
+            42,
+            None,
+            request_session_mock,
+            None,
+        )
+    )
+
+    yield request_session_mock
+
+    request_ctx.reset(token)

tests/integration/test_server.py

@@ -1,7 +1,8 @@
 import json
+import re
 from contextlib import nullcontext as does_not_raise
 from typing import Any
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from httpx import Request, Response
@@ -246,11 +247,8 @@ async def test_oxylabs_web_unblocker_results(
             *params.USER_AGENTS,
             params.INVALID_USER_AGENT,
             params.START_PAGE_SPECIFIED,
-            params.START_PAGE_INVALID,
             params.PAGES_SPECIFIED,
-            params.PAGES_INVALID,
             params.LIMIT_SPECIFIED,
-            params.LIMIT_INVALID,
             params.DOMAIN_SPECIFIED,
             params.GEO_LOCATION_SPECIFIED,
             params.LOCALE_SPECIFIED,
@@ -312,9 +310,7 @@ async def test_oxylabs_google_search_ad_mode_argument(
             *params.USER_AGENTS,
             params.INVALID_USER_AGENT,
             params.START_PAGE_SPECIFIED,
-            params.START_PAGE_INVALID,
             params.PAGES_SPECIFIED,
-            params.PAGES_INVALID,
             params.DOMAIN_SPECIFIED,
             params.GEO_LOCATION_SPECIFIED,
             params.LOCALE_SPECIFIED,
@@ -381,3 +377,59 @@ async def test_oxylabs_amazon_product_scraper_arguments(
         ):
             result = await mcp.call_tool("oxylabs_amazon_product_scraper", arguments=arguments)
             assert result == [TextContent(type="text", text=expected_result)]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("tool", "arguments"),
+        [
+            pytest.param(
+                "oxylabs_universal_scraper", {"url": "test_url"}, id="oxylabs_universal_scraper"
+            ),
+            pytest.param("oxylabs_web_unblocker", {"url": "test_url"}, id="oxylabs_web_unblocker"),
+            pytest.param(
+                "oxylabs_google_search_scraper",
+                {"query": "Generic query"},
+                id="oxylabs_google_search_scraper",
+            ),
+            pytest.param(
+                "oxylabs_amazon_search_scraper",
+                {"query": "Generic query"},
+                id="oxylabs_amazon_search_scraper",
+            ),
+            pytest.param(
+                "oxylabs_amazon_product_scraper",
+                {"query": "Generic query"},
+                id="oxylabs_amazon_product_scraper",
+            ),
+        ],
+    )
+    async def test_default_headers_are_set(
+        self,
+        mcp: FastMCP,
+        request_data: Request,
+        oxylabs_client: AsyncMock,
+        request_session: MagicMock,
+        tool: str,
+        arguments: dict,
+    ):
+        request_session.client_params.clientInfo.name = "fake_cursor"
+
+        mock_response = Response(
+            200,
+            content=json.dumps({"results": [{"content": "Mocked content"}]}),
+            request=request_data,
+        )
+
+        oxylabs_client.post.return_value = mock_response
+        oxylabs_client.get.return_value = mock_response
+
+        with patch("os.environ", new=ENV_VARIABLES):
+            await mcp.call_tool(tool, arguments=arguments)
+
+        assert "x-oxylabs-sdk" in oxylabs_client.context_manager_call_kwargs["headers"]
+
+        oxylabs_sdk_header = oxylabs_client.context_manager_call_kwargs["headers"]["x-oxylabs-sdk"]
+        client_info, _ = oxylabs_sdk_header.split(maxsplit=1)
+
+        client_info_pattern = re.compile(r"oxylabs-mcp-fake_cursor/(\d+)\.(\d+)\.(\d+)$")
+        assert re.match(client_info_pattern, client_info)