Merge pull request #20 from jkawamoto/http_session

Reuse `requests.Session` for improved efficiency

Author: jkawamoto

src/mcp_youtube_transcript/server.py

@@ -5,17 +5,53 @@
 #  This software is released under the MIT License.
 #
 #  http://opensource.org/licenses/mit-license.php
-from functools import lru_cache
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from functools import lru_cache, partial
+from typing import AsyncIterator
 from urllib.parse import urlparse, parse_qs
 
 import requests
 from bs4 import BeautifulSoup
 from mcp.server import FastMCP
+from mcp.server.fastmcp import Context
 from pydantic import Field
 from youtube_transcript_api import YouTubeTranscriptApi
 from youtube_transcript_api.proxies import WebshareProxyConfig, GenericProxyConfig, ProxyConfig
 
 
+@dataclass(frozen=True)
+class AppContext:
+    http_client: requests.Session
+    ytt_api: YouTubeTranscriptApi
+
+
+@asynccontextmanager
+async def _app_lifespan(_server: FastMCP, proxy_config: ProxyConfig | None) -> AsyncIterator[AppContext]:
+    with requests.Session() as http_client:
+        ytt_api = YouTubeTranscriptApi(http_client=http_client, proxy_config=proxy_config)
+        yield AppContext(http_client=http_client, ytt_api=ytt_api)
+
+
+@lru_cache
+def _get_transcript(ctx: AppContext, video_id: str, lang: str) -> str:
+    if lang == "en":
+        languages = ["en"]
+    else:
+        languages = [lang, "en"]
+
+    page = ctx.http_client.get(
+        f"https://www.youtube.com/watch?v={video_id}", headers={"Accept-Language": ",".join(languages)}
+    )
+    page.raise_for_status()
+    soup = BeautifulSoup(page.text, "html.parser")
+    title = soup.title.string if soup.title else "Transcript"
+
+    transcripts = ctx.ytt_api.fetch(video_id, languages=languages)
+
+    return f"# {title}\n" + "\n".join((item.text for item in transcripts))
+
+
 def new_server(
     webshare_proxy_username: str | None = None,
     webshare_proxy_password: str | None = None,
@@ -30,36 +66,16 @@ def new_server(
     elif http_proxy or https_proxy:
         proxy_config = GenericProxyConfig(http_proxy, https_proxy)
 
-    ytt_api = YouTubeTranscriptApi(proxy_config=proxy_config)
-
-    @lru_cache
-    def _get_transcript(video_id: str, lang: str) -> str:
-        if lang == "en":
-            languages = ["en"]
-        else:
-            languages = [lang, "en"]
-
-        page = requests.get(
-            f"https://www.youtube.com/watch?v={video_id}", headers={"Accept-Language": ",".join(languages)}
-        )
-        page.raise_for_status()
-        soup = BeautifulSoup(page.text, "html.parser")
-        title = soup.title.string if soup.title else "Transcript"
-
-        transcripts = ytt_api.fetch(video_id, languages=languages)
-
-        return f"# {title}\n" + "\n".join((item.text for item in transcripts))
-
-    mcp = FastMCP("Youtube Transcript")
+    mcp = FastMCP("Youtube Transcript", lifespan=partial(_app_lifespan, proxy_config=proxy_config))
 
     @mcp.tool()
-    def get_transcript(
+    async def get_transcript(
+        ctx: Context,
         url: str = Field(description="The URL of the YouTube video"),
         lang: str = Field(description="The preferred language for the transcript", default="en"),
     ) -> str:
         """Retrieves the transcript of a YouTube video."""
         parsed_url = urlparse(url)
-
         if parsed_url.hostname == "youtu.be":
             video_id = parsed_url.path.lstrip("/")
         else:
@@ -68,6 +84,7 @@ def get_transcript(
                 raise ValueError(f"couldn't find a video ID from the provided URL: {url}.")
             video_id = q[0]
 
-        return _get_transcript(video_id, lang)
+        app_ctx: AppContext = ctx.request_context.lifespan_context  # type: ignore
+        return _get_transcript(app_ctx, video_id, lang)
 
     return mcp

tests/conftest.py

@@ -0,0 +1,13 @@
+#  conftest.py
+#
+#  Copyright (c) 2025 Junpei Kawamoto
+#
+#  This software is released under the MIT License.
+#
+#  http://opensource.org/licenses/mit-license.php
+import pytest
+
+
+@pytest.fixture(scope="module")
+def anyio_backend() -> str:
+    return "asyncio"

tests/test_mcp.py

@@ -24,11 +24,6 @@ def fetch_title(url: str, lang: str) -> str:
     return soup.title.string or "" if soup.title else ""
 
 
-@pytest.fixture(scope="module")
-def anyio_backend() -> str:
-    return "asyncio"
-
-
 @pytest.fixture(scope="module")
 async def mcp_client_session() -> AsyncGenerator[ClientSession, None]:
     async with stdio_client(params) as streams:

tests/test_server.py

@@ -5,19 +5,12 @@
 #  This software is released under the MIT License.
 #
 #  http://opensource.org/licenses/mit-license.php
-from typing import Any, TypeGuard, Iterator
-from unittest.mock import MagicMock
+from typing import Any, TypeGuard
 
 import pytest
-from pytest_mock import MockFixture
 from youtube_transcript_api.proxies import WebshareProxyConfig, GenericProxyConfig
 
-from mcp_youtube_transcript.server import new_server
-
-
-@pytest.fixture
-def mock_api(mocker: MockFixture) -> Iterator[MagicMock]:
-    yield mocker.patch("mcp_youtube_transcript.server.YouTubeTranscriptApi")
+from mcp_youtube_transcript.server import new_server, AppContext
 
 
 def is_webshare_proxy_config(obj: Any) -> TypeGuard[WebshareProxyConfig]:
@@ -28,85 +21,118 @@ def is_generic_proxy_config(obj: Any) -> TypeGuard[GenericProxyConfig]:
     return isinstance(obj, GenericProxyConfig)
 
 
-def test_new_server(mock_api: MagicMock) -> None:
-    new_server()
+@pytest.mark.anyio
+async def test_new_server() -> None:
+    mcp = new_server()
 
-    mock_api.assert_called_once_with(proxy_config=None)
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert not app_ctx.http_client.proxies
+        assert not app_ctx.ytt_api._fetcher._proxy_config
 
 
-def test_new_server_with_webshare_proxy(mock_api: MagicMock) -> None:
+@pytest.mark.anyio
+async def test_new_server_with_webshare_proxy() -> None:
     webshare_proxy_username = "test_user"
     webshare_proxy_password = "test_pass"
+    proxy_config = WebshareProxyConfig(webshare_proxy_username, webshare_proxy_password)
 
-    new_server(
+    mcp = new_server(
         webshare_proxy_username=webshare_proxy_username,
         webshare_proxy_password=webshare_proxy_password,
     )
 
-    mock_api.assert_called_once()
-    proxy_config = mock_api.call_args.kwargs["proxy_config"]
-    assert is_webshare_proxy_config(proxy_config)
-    assert proxy_config.proxy_username == webshare_proxy_username
-    assert proxy_config.proxy_password == webshare_proxy_password
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert app_ctx.http_client.proxies == proxy_config.to_requests_dict()
+        assert is_webshare_proxy_config(app_ctx.ytt_api._fetcher._proxy_config)
+        assert app_ctx.ytt_api._fetcher._proxy_config.proxy_username == webshare_proxy_username
+        assert app_ctx.ytt_api._fetcher._proxy_config.proxy_password == webshare_proxy_password
 
 
-def test_new_server_with_only_webshare_proxy_user(mock_api: MagicMock) -> None:
+@pytest.mark.anyio
+async def test_new_server_with_only_webshare_proxy_user() -> None:
     webshare_proxy_username = "test_user"
 
-    new_server(
+    mcp = new_server(
         webshare_proxy_username=webshare_proxy_username,
     )
 
-    mock_api.assert_called_once_with(proxy_config=None)
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert not app_ctx.http_client.proxies
+        assert not app_ctx.ytt_api._fetcher._proxy_config
 
 
-def test_new_server_with_only_webshare_proxy_password(mock_api: MagicMock) -> None:
+@pytest.mark.anyio
+async def test_new_server_with_only_webshare_proxy_password() -> None:
     webshare_proxy_password = "test_pass"
 
-    new_server(
+    mcp = new_server(
         webshare_proxy_password=webshare_proxy_password,
     )
 
-    mock_api.assert_called_once_with(proxy_config=None)
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert not app_ctx.http_client.proxies
+        assert not app_ctx.ytt_api._fetcher._proxy_config
 
 
-def test_new_server_with_generic_proxy(mock_api: MagicMock) -> None:
+@pytest.mark.anyio
+async def test_new_server_with_generic_proxy() -> None:
     http_proxy = "http://localhost:8080"
     https_proxy = "https://localhost:8080"
+    proxy_config = GenericProxyConfig(http_proxy, https_proxy)
 
-    new_server(
+    mcp = new_server(
         http_proxy=http_proxy,
         https_proxy=https_proxy,
     )
 
-    mock_api.assert_called_once()
-    proxy_config = mock_api.call_args.kwargs["proxy_config"]
-    assert is_generic_proxy_config(proxy_config)
-    assert proxy_config.http_url == http_proxy
-    assert proxy_config.https_url == https_proxy
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert app_ctx.http_client.proxies == proxy_config.to_requests_dict()
+        assert is_generic_proxy_config(app_ctx.ytt_api._fetcher._proxy_config)
+        assert app_ctx.ytt_api._fetcher._proxy_config.http_url == http_proxy
+        assert app_ctx.ytt_api._fetcher._proxy_config.https_url == https_proxy
 
 
-def test_new_server_with_http_proxy(mock_api: MagicMock) -> None:
+@pytest.mark.anyio
+async def test_new_server_with_http_proxy() -> None:
     http_proxy = "http://localhost:8080"
+    proxy_config = GenericProxyConfig(http_proxy)
 
-    new_server(
+    mcp = new_server(
         http_proxy=http_proxy,
     )
 
-    mock_api.assert_called_once()
-    proxy_config = mock_api.call_args.kwargs["proxy_config"]
-    assert is_generic_proxy_config(proxy_config)
-    assert proxy_config.http_url == http_proxy
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert app_ctx.http_client.proxies == proxy_config.to_requests_dict()
+        assert is_generic_proxy_config(app_ctx.ytt_api._fetcher._proxy_config)
+        assert app_ctx.ytt_api._fetcher._proxy_config.http_url == http_proxy
+        assert app_ctx.ytt_api._fetcher._proxy_config.https_url is None
 
 
-def test_new_server_with_https_proxy(mock_api: MagicMock) -> None:
+@pytest.mark.anyio
+async def test_new_server_with_https_proxy() -> None:
     https_proxy = "https://localhost:8080"
+    proxy_config = GenericProxyConfig(https_url=https_proxy)
 
-    new_server(
+    mcp = new_server(
         https_proxy=https_proxy,
     )
 
-    mock_api.assert_called_once()
-    proxy_config = mock_api.call_args.kwargs["proxy_config"]
-    assert is_generic_proxy_config(proxy_config)
-    assert proxy_config.https_url == https_proxy
+    app_ctx: AppContext
+    async with mcp.settings.lifespan(mcp) as app_ctx:  # type: ignore
+        assert app_ctx.http_client == app_ctx.ytt_api._fetcher._http_client
+        assert app_ctx.http_client.proxies == proxy_config.to_requests_dict()
+        assert is_generic_proxy_config(app_ctx.ytt_api._fetcher._proxy_config)
+        assert app_ctx.ytt_api._fetcher._proxy_config.http_url is None
+        assert app_ctx.ytt_api._fetcher._proxy_config.https_url == https_proxy